Update app.py

app.py

@@ -1,29 +1,31 @@
 import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 
-# Load model and tokenizer (it will automatically pull from Hugging Face hub)
 model_id = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
 
 tokenizer = AutoTokenizer.from_pretrained(model_id)
 model = AutoModelForCausalLM.from_pretrained(model_id)
 
-# Create text-generation pipeline
 generator = pipeline(
-    "text-generation", 
-    model=model, 
-    tokenizer=tokenizer, 
+    "text-generation",
+    model=model,
+    tokenizer=tokenizer,
     max_new_tokens=200,
     temperature=0.7,
     top_p=0.9
 )
 
-# Define chatbot function
-def chat(prompt):
-    response = generator(prompt)
-    return response[0]['generated_text']
+def chat(user_input):
+    prompt = f"You are a helpful assistant.\nUser: {user_input}\nAssistant:"
+    response = generator(prompt, return_full_text=False)
+    return response[0]['generated_text'].strip()
 
-# Build Gradio interface
-iface = gr.Interface(fn=chat, inputs="text", outputs="text", title="Self-Hosted AI Chatbot (CPU)")
+iface = gr.Interface(
+    fn=chat,
+    inputs=gr.Textbox(lines=2, placeholder="Ask a full question, like 'Where is Delhi?'"),
+    outputs="text",
+    title="TinyLlama Chatbot 🤖",
+    description="Lightweight chatbot powered by TinyLlama. Works better with complete questions!"
+)
 
-# Launch app
 iface.launch()