Update app.py

app.py

@@ -10,7 +10,8 @@ torch.set_num_threads(1)
 print("🤖 Starting Model Loading...")
 
 try:
-    MODEL_NAME = "google/mt5-base"
+    # Use small version for faster inference and better generalization
+    MODEL_NAME = "google/mt5-small"
     TOKENIZER = AutoTokenizer.from_pretrained(MODEL_NAME)
     CHAT_MODEL = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME)
     print("✅ Model loaded successfully.")
@@ -31,7 +32,7 @@ def detect_script(text):
 
 
 # ---------------------------------------------------------
-# 3️⃣ Text Chat Function
+# 3️⃣ Chat Function
 # ---------------------------------------------------------
 def generate_reply(user_text):
     """Generate Sindhi or Roman Sindhi reply based on input language."""
@@ -40,20 +41,27 @@ def generate_reply(user_text):
 
     script_type = detect_script(user_text)
 
+    # Improved instruction-style prompts
     if script_type == "sindhi":
-        prompt = f"هيٺين سوال جو جواب سنڌي ۾ ڏيو:\n{user_text}"
+        prompt = f"توھان جو ڪم ھي آھي ته ھيٺين جملي جو جواب سنڌي ۾ ڏيو:\nسوال: {user_text}\nجواب:"
     else:
-        prompt = f"هيٺين سوال جو جواب رومن سنڌي ۾ ڏيو:\n{user_text}"
+        prompt = f"Tuhanjo kaam aahe ta neeche likhe sawal jo jawab Roman Sindhi mein likho:\nSawaal: {user_text}\nJawab:"
 
-    inputs = TOKENIZER(prompt, return_tensors="pt")
+    inputs = TOKENIZER(prompt, return_tensors="pt", truncation=True)
     outputs = CHAT_MODEL.generate(
         **inputs,
         max_new_tokens=100,
         do_sample=True,
         top_p=0.95,
-        top_k=50
+        top_k=40,
+        temperature=0.8
     )
-    return TOKENIZER.decode(outputs[0], skip_special_tokens=True)
+
+    reply = TOKENIZER.decode(outputs[0], skip_special_tokens=True)
+
+    # Remove unwanted tokens (like <extra_id_0>, etc.)
+    reply = re.sub(r"<.*?>", "", reply).strip()
+    return reply
 
 
 # ---------------------------------------------------------