Update app.py

app.py

@@ -88,9 +88,10 @@ def generate_response(prompt: str) -> str:
 
     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)
+    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])
+    
 
     return generated_text