Update app.py

app.py

@@ -1,28 +1,31 @@
-import torch 
 import gradio as gr
+import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 
-model_name = "tiiuae/falcon-rw-1b"
+model_name = "sshleifer/tiny-gpt2"  # Light model for free CPU Space
 
-# Load tokenizer and model
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype="auto")
+model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype="auto")
 
-# Create text generation pipeline
 generator = pipeline("text-generation", model=model, tokenizer=tokenizer)
 
-# Chat function (ChatGPT-style with faster, shorter replies)
 def chat(user_input, history=[]):
-    # Simple system prompt
-    system_prompt = (
-        "You are a helpful, smart, and friendly assistant. "
-        "Reply politely, clearly, and concisely like ChatGPT."
-    )
-
-    # Build conversation input
-    full_input = system_prompt + "\n"
-    for turn in history:
-        full_input += f"User: {turn[0]}\nAI: {turn[1]}\n"
-    full_input += f"User: {user_input}\nAI:"
-
-    # Generate reply (
+    prompt = "You are a helpful and friendly assistant.\n"
+    for msg in history:
+        prompt += f"User: {msg[0]}\nAI: {msg[1]}\n"
+    prompt += f"User: {user_input}\nAI:"
+
+    response = generator(prompt, max_new_tokens=60, do_sample=True, temperature=0.7)[0]["generated_text"]
+    reply = response.split("AI:")[-1].split("User:")[0].strip()
+    history.append((user_input, reply))
+    return history, history
+
+with gr.Blocks() as demo:
+    gr.Markdown("# 🤖 Tiny GPT-2 Chatbot")
+    chatbot = gr.Chatbot()
+    msg = gr.Textbox(label="Ask something...")
+    state = gr.State([])
+
+    msg.submit(chat, [msg, state], [chatbot, state])
+
+demo.launch()