Update app.py

app.py

@@ -94,7 +94,6 @@ def generate_response(prompt: str) -> str:
     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.
                     PROMPT_TEMPLATE_END
-
                     CONTEXT_START
                     Guy: "Hi there!"
                     Ivana: "Hey! How's it going?"
@@ -118,14 +117,15 @@ def generate_response(prompt: str) -> str:
                     Ivana: "Can't wait to see you there! "
                     CONTEXT_END"""
 
-    # prompt_template = f'''{PERSONA_DESC} ASSISTANT: '''.encode()
-    prompt_template = f'{PERSONA_DESC}\n\nASSISTANT: {prompt}\n'.encode()
-    input_ids = tokenizer(prompt_template, return_tensors='pt').input_ids.cuda()
-    output = model.generate(inputs=input_ids, temperature=0.7, do_sample=True, top_p=0.95, top_k=40, max_new_tokens=512)
-    generated_text = tokenizer.decode(output[0])
+    prompt_template = f'{PERSONA_DESC}\n\nASSISTANT: {prompt}\n'
+
+    input_ids = tokenizer(prompt_template, return_tensors='pt').input_ids.to(model.device)
+    output = model.generate(input_ids, temperature=0.7, do_sample=True, top_p=0.95, top_k=40, max_length=512)
+    generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
 
     return generated_text
 
 
 
 
+