added mask

app.py

@@ -129,11 +129,20 @@ def generate_response(prompt):
     
     full_prompt += "Assistant: "
     
-    # Generate response
-    inputs = llm_tokenizer(full_prompt, return_tensors="pt").to(device)
+    # Generate response with explicit attention mask
+    inputs = llm_tokenizer(
+        full_prompt, 
+        return_tensors="pt",
+        padding=True,
+        truncation=True,
+        max_length=512,
+        return_attention_mask=True  # Explicitly request attention mask
+    ).to(device)
+    
     with torch.no_grad():
         output = llm_model.generate(
-            **inputs, 
+            input_ids=inputs["input_ids"],
+            attention_mask=inputs["attention_mask"],  # Pass the attention mask
             max_new_tokens=128,
             do_sample=True,
             temperature=0.7,