Update app.py

app.py

@@ -62,11 +62,12 @@ def generate_response(prompt: str) -> str:
         repetition_penalty=1.1
     )
     generated_text = pipe(prompt_template)[0]['generated_text']
+            # Extract the user prompt and assistant's reply
+    assistant_response = generated_text.split(user_prompt)[-1].strip()
+            
+    return {"user": prompt, "assistant": assistant_response}
 
-    return generated_text
-
-
-
+  
 def generate_prompt_response(persona_prompt: str, prompt: str) -> dict:
     try:
         # Validate inputs