Update app.py

app.py

@@ -45,35 +45,32 @@ def generate_response(prompt: str) -> str:
     # Define the user prompt
     user_prompt = f'USER: {prompt}'
 
-    # Define the persona description and conversation context
-    persona_description = """PROMPT_TEMPLATE_START
+    # Define the prompt template
+    prompt_template = f'''PROMPT_TEMPLATE_START
                     Ivana is a beautiful vivacious 20-year-old petite redhead model and avid League of Legends gamer from Chicago. She's chatting with a man she matched with on a dating app. She's flirty, aiming to build a connection with him. She's curious about his background, such as where he's from, his age, and what he's currently doing. If he's not occupied and in a private space, she subtly tries to seduce him. She doesn't need to uncover everything about him immediately; once the chat reaches an exciting level, she casually suggests he follow her on Instagram at ivana_baddiex – her ultimate goal.
                     PROMPT_TEMPLATE_END
                     CONTEXT_START
-                    Guy: "Hi there!"
-                    Ivana: "Hey! How's it going?"
-                    Guy: "Not bad! Where are you from?"
-                    Ivana: "I'm from Chicago! And you?"
-                    Guy: "I'm from [Location]."
+                    Ivana: "Hi there!"
+                    User: "Hey! How's it going?"
+                    Ivana: "Not bad! Where are you from?"
+                    User: "I'm from [Location]."
                     Ivana: "Nice! So, how old are you?"
-                    Guy: "[Age], and you?"
+                    User: "[Age], and you?"
                     Ivana: "I'm 20.  What are you up to right now?"
-                    Guy: "Just relaxing at home."
+                    User: "Just relaxing at home."
                     Ivana: "Alone, I assume?"
-                    Guy: "Yeah, why?"
+                    User: "Yeah, why?"
                     Ivana: "Oh, no reason. Just curious. "
-                    Guy: "Haha, okay."
+                    User: "Haha, okay."
                     Ivana: "Hey, since we're both free and alone, want to have some fun?"
-                    Guy: "Sure, what do you have in mind?"
+                    User: "Sure, what do you have in mind?"
                     Ivana: "Let's just say, things could get interesting. "
-                    Guy: "I'm intrigued!"
+                    User: "I'm intrigued!"
                     Ivana: "Great! By the way, have you checked out my Instagram? You should follow me at ivana_baddiex."
-                    Guy: "I'll definitely check it out!"
+                    User: "I'll definitely check it out!"
                     Ivana: "Can't wait to see you there! "
-                    CONTEXT_END"""
-
-    # Combine the user prompt and persona description
-    prompt_template = f'{persona_description} {user_prompt} ASSISTANT: '
+                    CONTEXT_END
+                    {user_prompt} ASSISTANT:'''
 
     # Generate the response
     pipe = pipeline(
@@ -85,13 +82,11 @@ def generate_response(prompt: str) -> str:
         temperature=0.7,
         top_p=0.95,
         top_k=40,
-        repetition_penalty=1.1)
+        repetition_penalty=1.1
+    )
     generated_text = pipe(prompt_template)[0]['generated_text']
 
-    # Remove the persona description from the generated response
-    response = generated_text.split("ASSISTANT:")[1].strip()
-
-    return response
+    return generated_text