Update app.py

app.py

@@ -26,22 +26,42 @@ def load_model():
 def respond(message, history):
     load_model()
 
-    prompt = ""
+    prompt = "<bos>"
     for user, bot in history:
-        prompt += f"User: {user}\nAssistant: {bot}\n"
-    prompt += f"User: {message}\nAssistant:"
+        prompt += (
+            "<start_of_turn>user\n"
+            f"{user}\n"
+            "<end_of_turn>\n"
+            "<start_of_turn>model\n"
+            f"{bot}\n"
+            "<end_of_turn>\n"
+        )
+
+    prompt += (
+        "<start_of_turn>user\n"
+        f"{message}\n"
+        "<end_of_turn>\n"
+        "<start_of_turn>model\n"
+    )
+
+    inputs = tokenizer(
+        prompt,
+        return_tensors="pt",
+        truncation=True,
+        max_length=2048
+    )
 
-    inputs = tokenizer(prompt, return_tensors="pt", truncation=True)
     outputs = model.generate(
         **inputs,
-        max_new_tokens=150,
-        do_sample=True,
+        max_new_tokens=200,
         temperature=0.7,
-        top_p=0.9
+        top_p=0.9,
+        do_sample=True,
+        eos_token_id=tokenizer.eos_token_id
     )
 
-    text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    reply = text.split("Assistant:")[-1].strip()
+    decoded = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    reply = decoded.split("<start_of_turn>model")[-1].strip()
     return reply
 
 gr.ChatInterface(