Update app.py

app.py

@@ -43,24 +43,19 @@ def generate_response(prompt, max_length=1024):
     
     outputs = pipe(messages, max_new_tokens=max_length)
     
-    # Find assistant's response in the output
+    # Extract just the assistant's response
     try:
-        # The output contains the full conversation history
-        generated_text = outputs[0]["generated_text"]
-        # Look for the last assistant message
-        assistant_prefix = "{'role': 'assistant', 'content': '"
-        assistant_start = generated_text.rfind(assistant_prefix)
-        if assistant_start != -1:
-            # Move past the prefix
-            content_start = assistant_start + len(assistant_prefix)
-            # Find the end of the content (before the closing quote and brace)
-            content_end = generated_text.rfind("'}")
-            if content_end != -1:
-                return generated_text[content_start:content_end]
+        # Get the message list from the output
+        message_list = eval(outputs[0]["generated_text"])  # Safely convert string to list
+        # Get the last message (assistant's response)
+        assistant_message = message_list[-1]
+        if assistant_message["role"] == "assistant":
+            return assistant_message["content"]
     except Exception as e:
         print(f"Error extracting response: {e}")
-    
-    # Fallback: return the raw generated text if extraction fails
+        # If extraction fails, return the raw output
+        return str(outputs[0]["generated_text"])
+        
     return outputs[0]["generated_text"]
 
 # Example with proper line breaks