chore: use gradio instead of streamlit

README.md

@@ -3,8 +3,8 @@ title: Customer Service Chatbot
 emoji: 🔮
 colorFrom: indigo
 colorTo: indigo
-sdk: streamlit
-sdk_version: 1.34.0
+sdk: gradio
+sdk_version: 5.23.0
 app_file: app.py
 pinned: false
 ---

app.py

@@ -77,9 +77,20 @@ def predict(message, history):
     inputs = tokenizer(prompt, return_tensors="pt").to(device)
     # 3. Generate response
     outputs = model.generate(**inputs, max_new_tokens=100)
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    decoded = tokenizer.decode(outputs[0], skip_special_tokens=True)
 
     # print(f"Response: {response}, outputs: {outputs}")
+    
+    # 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:
+        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)