Update app.py

app.py

@@ -1,85 +1,50 @@
 import asyncio
-from datasets import load_dataset
 from transformers import AutoTokenizer, AutoModelForCausalLM
-import torch
+from datasets import load_dataset
+
+# -------------------------------
+# Model loading
+# -------------------------------
+MODEL_NAME = "LiquidAI/LFM2-2.6B"
 
-# ---------- CONFIG ----------
-MODEL_NAME = "LiquidAI/LFM2-2.6B"  # Replace with your model
-MAX_PROMPTS_INITIAL = 20           # Initial prompts to load fast
-MAX_PROMPTS_TOTAL = 200            # Load all prompts later
-MAX_NEW_TOKENS = 50
-DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+print("Loading model...")
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
+print("Model loaded.")
 
-# ---------- FUNCTION TO LOAD PROMPTS ----------
+# -------------------------------
+# Fetch prompts from Hugging Face
+# -------------------------------
 async def fetch_prompts():
-    """
-    Load prompts using Hugging Face datasets library.
-    Returns a list of prompts.
-    """
+    print("Fetching prompts from Hugging Face dataset...")
     dataset = load_dataset("fka/awesome-chatgpt-prompts", split="train")
-    # Extract prompts as a simple list
-    prompts = [item["prompt"] for item in dataset]
-    return prompts
-
-# ---------- FUNCTION TO LOAD MODEL ----------
-def load_model():
-    """
-    Load tokenizer and model to device.
-    """
-    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-    model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
-    model.to(DEVICE)
-    return tokenizer, model
-
-# ---------- FUNCTION TO GENERATE RESPONSE ----------
-def generate_response(tokenizer, model, prompt, max_new_tokens=MAX_NEW_TOKENS):
-    """
-    Generate a response from the model for a given prompt.
-    """
-    messages = [{"role": "user", "content": prompt}]
-    inputs = tokenizer.apply_chat_template(
-        messages,
-        add_generation_prompt=True,
-        tokenize=True,
-        return_dict=True,
-        return_tensors="pt"
-    ).to(DEVICE)
-
-    outputs = model.generate(**inputs, max_new_tokens=max_new_tokens)
-    response = tokenizer.decode(outputs[0][inputs["input_ids"].shape[-1]:])
-    return response
+    all_prompts = dataset['prompt']  # List of 200+ prompts
+    print(f"Total prompts available: {len(all_prompts)}")
+    return all_prompts
 
-# ---------- MAIN FUNCTION ----------
+# -------------------------------
+# Main async function
+# -------------------------------
 async def main():
-    print("Loading model...")
-    tokenizer, model = load_model()
-    print("Model loaded on", DEVICE)
-
-    print("Fetching initial prompts...")
     all_prompts = await fetch_prompts()
-    print(f"Total prompts available: {len(all_prompts)}")
 
-    # Load initial batch
-    initial_prompts = all_prompts[:MAX_PROMPTS_INITIAL]
-    print(f"Using first {MAX_PROMPTS_INITIAL} prompts for fast startup...")
+    # Use first 20 for fast startup
+    fast_prompts = all_prompts[:20]
+    print("Using first 20 prompts for fast startup...")
+    for i, p in enumerate(fast_prompts, 1):
+        print(f"[Prompt {i}] {p}")
 
-    # Generate responses for initial prompts
-    for i, prompt in enumerate(initial_prompts, 1):
-        response = generate_response(tokenizer, model, prompt)
-        print(f"[Prompt {i}] {prompt}")
-        print(f"[Response {i}] {response}\n")
+    # Load remaining prompts asynchronously
+    remaining_prompts = all_prompts[20:]
+    print(f"Loading remaining {len(remaining_prompts)} prompts asynchronously...")
 
-    # Optionally, load all remaining prompts asynchronously after initial batch
-    remaining_prompts = all_prompts[MAX_PROMPTS_INITIAL:MAX_PROMPTS_TOTAL]
-    if remaining_prompts:
-        print(f"Loading remaining {len(remaining_prompts)} prompts asynchronously...")
-        # Here you could generate responses asynchronously if desired
-        # Example: loop over remaining prompts and generate later
-        await asyncio.sleep(0.1)  # placeholder for async tasks
+    # Example: You could process remaining prompts in background
+    # For demonstration, just print count after "async load"
+    await asyncio.sleep(1)
+    print("Remaining prompts loaded.")
 
-# ---------- ENTRY POINT ----------
+# -------------------------------
+# Entry point
+# -------------------------------
 if __name__ == "__main__":
-    try:
-        asyncio.run(main())
-    except Exception as e:
-        print("Fatal error:", str(e))
+    asyncio.run(main())