Update app.py

app.py

@@ -16,15 +16,14 @@ else:
 @spaces.GPU
 def generate_response(passage: str, question: str) -> str:
     # Prepare the input text by combining the passage and question
-    messages = [{"role": "user", "content": f"Passage: {passage}\nQuestion: {question}"}]
-    inputs = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, return_tensors="pt")
+    chat = [{"role": "user", "content": f"Passage: {passage}\nQuestion: {question}"}]
+
+    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(olmo.device), max_new_tokens=100)
     
-    # Generate text, focusing only on the new tokens added by the model
-    outputs = model.generate(**inputs, max_new_tokens=150)
+    response = tokenizer.batch_decode(response, skip_special_tokens=True)[0]
     
-    # Decode only the generated part, skipping the prompt input
-    # generated_tokens = outputs[0][inputs.input_ids.shape[-1]:]  # Ignore input tokens in the output
-    response = tokenizer.batch_decode(outputs, skip_special_tokens=True)[0]
     
     return response