Update app.py

app.py

@@ -42,9 +42,12 @@ app = FastAPI(root_path="/api/v1")
 # Function to generate a response using the model
     
 def generate_response(prompt: str) -> str:
-    PERSONA_NAME = "Ivana"
-    PERSONA_DESC = f"""PROMPT_TEMPLATE_START
-                    {PERSONA_NAME} 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.
+    # Define the user prompt
+    user_prompt = f'USER: {prompt}'
+
+    # Define the persona description and conversation context
+    persona_description = """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!"
@@ -68,27 +71,27 @@ def generate_response(prompt: str) -> str:
                     Guy: "I'll definitely check it out!"
                     Ivana: "Can't wait to see you there! "
                     CONTEXT_END"""
-    prompt_template=f''' {PERSONA_NAME} 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. USER: {prompt} ASSISTANT:
 
-'''
-    #prompt_template = f'{PERSONA_DESC}\n\nASSISTANT: {prompt}\n'
+    # Combine the user prompt and persona description
+    prompt_template = f'{persona_description} {user_prompt} ASSISTANT: '
 
+    # Generate the response
     pipe = pipeline(
-    "text-generation",
-    model=model,
-    tokenizer=tokenizer,
-    max_new_tokens=512,
-    do_sample=True,
-    temperature=0.7,
-    top_p=0.95,
-    top_k=40,
-    repetition_penalty=1.1
-)
-    generated_text = (pipe(prompt_template)[0]['generated_text'])
-    
-
-    return generated_text
-
+        "text-generation",
+        model=model,
+        tokenizer=tokenizer,
+        max_new_tokens=512,
+        do_sample=True,
+        temperature=0.7,
+        top_p=0.95,
+        top_k=40,
+        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