Update app.py

app.py

@@ -91,10 +91,11 @@ def generate_prompt_response(persona_prompt: str, prompt: str) -> dict:
         )
         generated_text = pipe(persona_prompt + user_prompt)[0]['generated_text']
         
-        # Extract only the assistant's response from the generated text
-        assistant_response = generated_text.split(user_prompt)[-1].strip()
-            
-        return {"user": prompt,"assistant": assistant_response}
+            # Remove the "ASSISTANT:" prefix from the generated text
+        assistant_response = generated_text.replace("ASSISTANT:", "").strip()
+    
+        # Return the user prompt and assistant's response as a dictionary
+        return {"user": prompt, "assistant": assistant_response}
 
     except Exception as e:
         # Handle any exceptions and return an error message