Update app.py

app.py

@@ -19,14 +19,14 @@ else:
 @spaces.GPU
 def generate_response(passage: str, question: str) -> str:
     # Prepare the input text by combining the passage and question
-    chat = [{"role": "user", "content": f"Passage: {passage}\nQuestion: {question}"}]
+    message = [f"Passage: {passage}\nQuestion: {question}"]
+    inputs = tokenizer(message, return_tensors='pt', return_token_type_ids=False).to('cuda')
 
-    prompt = tokenizer.apply_chat_template(chat, tokenize=False, add_generation_prompt=True)
-    inputs = tokenizer.encode(prompt, add_special_tokens=False, return_tensors="pt")
-    response = model.generate(input_ids=inputs.to(model.device), max_new_tokens=100)
-    
-    response = tokenizer.batch_decode(response, skip_special_tokens=True)[0].split("<|assistant|>")[-1].strip()
+    response = model.generate(**inputs, max_new_tokens=100)
     
+    response = tokenizer.batch_decode(response, skip_special_tokens=True)[0]
+
+    response = response[len(message[0]):].strip().split('\n')[0]
     
     return response