Update app.py

app.py

@@ -6,6 +6,7 @@ model_name = "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B"
 print("Loading model...")
 
 tokenizer = AutoTokenizer.from_pretrained(model_name)
+
 model = AutoModelForCausalLM.from_pretrained(
     model_name,
     device_map="auto"
@@ -14,33 +15,45 @@ model = AutoModelForCausalLM.from_pretrained(
 print("Model loaded!")
 
 def chat(message, history):
-    conversation = ""
+    try:
+        conversation = ""
+
+        # Limit history to avoid overload
+        for user, bot in history[-3:]:
+            conversation += f"User: {user}\nAssistant: {bot}\n"
+
+        conversation += f"User: {message}\nAssistant:"
 
-    for user, bot in history:
-        conversation += f"User: {user}\nAssistant: {bot}\n"
+        inputs = tokenizer(
+            conversation,
+            return_tensors="pt",
+            truncation=True,
+            max_length=1024
+        ).to(model.device)
 
-    conversation += f"User: {message}\nAssistant:"
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=150,
+            temperature=0.7,
+            top_p=0.9,
+            repetition_penalty=1.1,
+            do_sample=True
+        )
 
-    inputs = tokenizer(conversation, return_tensors="pt").to(model.device)
+        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-    outputs = model.generate(
-        **inputs,
-        max_new_tokens=300,
-        temperature=0.8,
-        top_p=0.95,
-        repetition_penalty=1.1,
-        do_sample=True
-    )
+        # Extract only latest response
+        response = response.split("Assistant:")[-1].strip()
 
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    response = response.split("Assistant:")[-1].strip()
+        return response
 
-    return response
+    except Exception as e:
+        return "⚠️ Error: Model overloaded. Try again with a shorter question."
 
 iface = gr.ChatInterface(
     fn=chat,
     title="Emalawi19 AI 🤖",
-    description="Your own AI assistant powered by DeepSeek"
+    description="Chat with your own AI assistant powered by DeepSeek"
 )
 
 iface.launch()