Update app.py

app.py

@@ -2,7 +2,8 @@ import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
 
-MODEL_NAME = "deepseek-ai/deepseek-coder-1.3b-instruct"
+# 🔥 Better lightweight model
+MODEL_NAME = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
 
 print("🔄 Loading model from Hugging Face...")
 
@@ -11,45 +12,43 @@ model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
 
 print("✅ Model loaded!")
 
-chat_history_ids = None
-
+# 🔹 Chat function
 def chat(user_input):
-    global chat_history_ids
+    # 🔥 Strong system prompt (fixes bad answers)
+    prompt = f"""You are a helpful AI assistant.
+You always give clear, correct, and short answers.
+If asked for code, return clean and working code only.
 
-    new_input_ids = tokenizer.encode(
-        user_input + tokenizer.eos_token,
-        return_tensors='pt'
-    )
+User: {user_input}
+Assistant:"""
 
-    if chat_history_ids is not None:
-        bot_input_ids = torch.cat([chat_history_ids, new_input_ids], dim=-1)
-    else:
-        bot_input_ids = new_input_ids
+    inputs = tokenizer(prompt, return_tensors="pt")
 
-    chat_history_ids = model.generate(
-        bot_input_ids,
-        max_length=1000,
-        pad_token_id=tokenizer.eos_token_id,
+    output_ids = model.generate(
+        inputs["input_ids"],
+        max_new_tokens=150,        # ⚡ faster
         do_sample=True,
-        top_k=50,
-        top_p=0.95,
-        temperature=0.75
+        temperature=0.7,
+        top_p=0.9,
+        repetition_penalty=1.2
     )
 
-    response = tokenizer.decode(
-        chat_history_ids[:, bot_input_ids.shape[-1]:][0],
-        skip_special_tokens=True
-    )
+    response = tokenizer.decode(output_ids[0], skip_special_tokens=True)
+
+    # 🔥 Clean output
+    if "Assistant:" in response:
+        response = response.split("Assistant:")[-1].strip()
 
     return response
 
 
+# 🔹 Simple UI
 iface = gr.Interface(
     fn=chat,
-    inputs="text",
+    inputs=gr.Textbox(placeholder="Type your message here..."),
     outputs="text",
     title="Emalawi19 AI 🤖",
-    description="Chat with your AI powered by Hugging Face"
+    description="Fast & smarter AI powered by Hugging Face"
 )
 
 iface.launch()