Update Gradio app with multiple files

app.py

@@ -80,11 +80,22 @@ def chat_with_model(message, history):
         full_response = tokenizer.batch_decode(output)[0]
         
         # Extract only the assistant's response
-        response_start = full_response.find('<|assistant|>') + len('<|assistant|>')
-        assistant_response = full_response[response_start:].strip()
-        
-        # Clean up the response
-        assistant_response = assistant_response.replace('<|endoftext|>', '').strip()
+        # Find the start of assistant role
+        assistant_start = full_response.find('<|start_of_role|>assistant<|end_of_role|>')
+        if assistant_start != -1:
+            assistant_start += len('<|start_of_role|>assistant<|end_of_role|>')
+            assistant_response = full_response[assistant_start:].strip()
+        else:
+            # Fallback to original method if pattern not found
+            response_start = full_response.find('<|assistant|>')
+            if response_start != -1:
+                response_start += len('<|assistant|>')
+                assistant_response = full_response[response_start:].strip()
+            else:
+                assistant_response = full_response.strip()
+        
+        # Clean up the response - remove end markers
+        assistant_response = assistant_response.replace('<|endoftext|>', '').replace('<|end_of_text|>', '').strip()
         
         return assistant_response