Updates.

app.py

@@ -71,9 +71,12 @@ You are a medical assistant trained on clinical reasoning data. Given the follow
 def generate_local_answer(prompt, max_new_tokens=512):
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     print(f"Using device: {device}")
-    input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(device)
+    inputs = tokenizer(prompt, return_tensors="pt", padding=True).to(device)
+    input_ids = inputs["input_ids"]
+    attention_mask = inputs["attention_mask"]
     output = generation_model.generate(
         input_ids=input_ids,
+        attention_mask=attention_mask,
         max_new_tokens=max_new_tokens,
         temperature=0.5,
         do_sample=True,