Update app.py

app.py

@@ -1,69 +1,76 @@
-import os
 import asyncio
-import pandas as pd
-from datasets import load_dataset
-from transformers import AutoTokenizer, AutoModelForCausalLM
-import torch
+import random
 import gradio as gr
+from datasets import load_dataset
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 
-# Paths
-PROMPTS_CSV = "prompts.csv"
-MODEL_NAME = "LiquidAI/LFM2-2.6B"
+print("===== Application Startup =====")
 
-# Check for dataset, download if missing
-if not os.path.exists(PROMPTS_CSV):
-    print("prompts.csv not found. Downloading dataset from Hugging Face...")
-    dataset = load_dataset("fka/awesome-chatgpt-prompts", split="train")
-    df = pd.DataFrame(dataset)
-    df.to_csv(PROMPTS_CSV, index=False)
-    print("Dataset saved to prompts.csv")
-else:
-    df = pd.read_csv(PROMPTS_CSV)
+# -----------------------
+# Load model
+# -----------------------
+print("Loading model...")
+model_name = "gpt2"  # you can swap this for a larger model if you have GPU
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name)
+generator = pipeline("text-generation", model=model, tokenizer=tokenizer)
+print("Model loaded successfully.")
 
-all_prompts = df['prompt'].tolist()
-print(f"Total prompts available: {len(all_prompts)}")
+# -----------------------
+# Load dataset
+# -----------------------
+print("Fetching dataset...")
+dataset = load_dataset("lvwerra/stack-exchange-paired", split="train[:200]")  
+# limit to 200 for speed – you can increase if you want
+print(f"Total prompts available: {len(dataset)}")
 
-# Load first 20 prompts for fast startup
-fast_prompts = all_prompts[:20]
-remaining_prompts = all_prompts[20:]
+# Split dataset
+initial_prompts = dataset[:20]   # first 20 for fast startup
+remaining_prompts = dataset[20:] # remaining ~180
 
-# Load tokenizer and model
-print("Loading model...")
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
-device = "cuda" if torch.cuda.is_available() else "cpu"
-model.to(device)
-print(f"Model loaded on {device}")
+# Storage for loaded prompts
+prompts = []
+for item in initial_prompts:
+    prompts.append(item["question"])
 
-# Async function to load remaining prompts
-async def load_remaining_prompts():
-    global fast_prompts
-    print("Loading remaining prompts asynchronously...")
-    await asyncio.sleep(1)  # simulate async loading
-    fast_prompts.extend(remaining_prompts)
-    print("All prompts loaded.")
+print(f"Loaded {len(prompts)} initial prompts for fast startup.")
 
-# Function to generate response
-def generate_response(prompt, max_tokens=100):
-    inputs = tokenizer(prompt, return_tensors="pt").to(device)
-    outputs = model.generate(**inputs, max_new_tokens=max_tokens)
-    response = tokenizer.decode(outputs[0][inputs["input_ids"].shape[-1]:], skip_special_tokens=True)
-    return response
+# -----------------------
+# Async loading of remaining prompts
+# -----------------------
+async def load_remaining_prompts():
+    print("Background: Loading remaining prompts...")
+    await asyncio.sleep(2)  # simulate delay
+    for item in remaining_prompts:
+        prompts.append(item["question"])
+    print(f"Background: Finished loading. Total prompts now = {len(prompts)}")
 
+# -----------------------
 # Gradio interface
-def chat_with_prompt(prompt_idx):
-    prompt = fast_prompts[prompt_idx]
-    response = generate_response(prompt)
-    return f"Prompt:\n{prompt}\n\nResponse:\n{response}"
-
-with gr.Blocks() as demo:
-    gr.Markdown("## ChatGPT Prompt Tester")
-    prompt_dropdown = gr.Dropdown(choices=[str(i) for i in range(len(fast_prompts))], label="Select Prompt Index")
-    output_text = gr.Textbox(label="Model Response", lines=15)
-    prompt_dropdown.change(chat_with_prompt, inputs=prompt_dropdown, outputs=output_text)
+# -----------------------
+def chat_with_model(user_input):
+    """Respond to user with a random dataset prompt + model output."""
+    if not prompts:
+        return "Prompts not ready yet. Please wait..."
+    prompt = random.choice(prompts)
+    response = generator(f"{prompt}\n\nUser: {user_input}\nAI:", 
+                         max_length=100, 
+                         num_return_sequences=1,
+                         do_sample=True)[0]["generated_text"]
+    return response
 
-# Run async loading in the background
-asyncio.create_task(load_remaining_prompts())
+demo = gr.Interface(
+    fn=chat_with_model,
+    inputs=gr.Textbox(lines=2, placeholder="Ask me something..."),
+    outputs="text",
+    title="Fast Prompt Loader Chatbot",
+    description="Loads 20 prompts fast, then background loads 200+ prompts"
+)
 
-# Launch Gradio
-demo.launch(server_name="0.0.0.0", server_port=7860)
+# -----------------------
+# App runner
+# -----------------------
+if __name__ == "__main__":
+    loop = asyncio.get_event_loop()
+    loop.create_task(load_remaining_prompts())  # schedule async loading
+    demo.launch(server_name="0.0.0.0", server_port=7860)