chore: parse assistant message

app.py

@@ -77,20 +77,24 @@ def predict(message, history):
     inputs = tokenizer(prompt, return_tensors="pt").to(device)
     # 3. Generate response
     outputs = model.generate(**inputs, max_new_tokens=100)
-    decoded = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    # skip_special_tokens=False: we want to keep special tokens for easier parsing
+    decoded = tokenizer.decode(outputs[0], skip_special_tokens=False)
 
-    # print(f"Response: {response}, outputs: {outputs}")
+    # print(f"decoded: {decoded}\n")
+    # print(f"outputs: {outputs}\n")
     
     # Extract only the assistant's message (after the last user message)
-    # This works for most chat templates that append the assistant's reply at the end
-    if "<|im_start|>assistant" in decoded:
+    if "<|start_header_id|>assistant<|end_header_id|>" in decoded:
+        response = decoded.split("<|start_header_id|>assistant<|end_header_id|>")[-1]
+        response = response.replace("<|eot_id|>", "").strip()
+    elif "<|im_start|>assistant" in decoded:
+        # This works for most chat templates that append the assistant's reply at the end
         response = decoded.split("<|im_start|>assistant")[-1]
-        # Remove possible end tokens or markers
         response = response.replace("<|im_end|>", "").strip()
     else:
         # Fallback: just return the decoded output
         response = decoded.strip()
-
+        
     return response
 
 demo = gr.ChatInterface(predict, type="messages", examples=examples)