Update app.py

app.py

@@ -63,9 +63,12 @@ def generate_response(prompt: str) -> str:
         top_k=40,
         repetition_penalty=1.1
     )
-    generated_text = pipe(prompt_template)[0]['generated_text']
+        generated_text = pipe(persona_prompt + user_prompt)[0]['generated_text']
+        
+        # Remove the "ASSISTANT:" prefix from the generated text
+        assistant_response = generated_text.replace("ASSISTANT:", "").strip()
             
-    return {"user": prompt, "assistant": generated_text}
+    return {"user": prompt, "assistant": assistant_response}
 
   
 def generate_prompt_response(persona_prompt: str, prompt: str) -> dict:
@@ -91,7 +94,7 @@ def generate_prompt_response(persona_prompt: str, prompt: str) -> dict:
         )
         generated_text = pipe(persona_prompt + user_prompt)[0]['generated_text']
         
-            # Remove the "ASSISTANT:" prefix from the generated text
+        # 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