Update app.py

app.py

@@ -51,7 +51,7 @@ model, tokenizer = load_model_and_tokenizer()
 def generate_response(model, tokenizer, instruction, max_new_tokens=2048):
     """Generate a response from the model based on an instruction."""
     try:
-        # Encode input with truncation
+        # Encode input with truncation and create an attention mask
         inputs = tokenizer.encode(
             instruction,
             return_tensors="pt",
@@ -59,15 +59,24 @@ def generate_response(model, tokenizer, instruction, max_new_tokens=2048):
             max_length=tokenizer.model_max_length
         ).to(model.device)
 
+        # Create attention mask (1 for real tokens, 0 for padding tokens)
+        attention_mask = torch.ones(inputs.shape, device=model.device)
+
+        print(f"Model input tokens: {inputs}")  # Debugging line
+        print(f"Attention mask: {attention_mask}")  # Debugging line
+
         # Generate response
         outputs = model.generate(
             inputs,
+            attention_mask=attention_mask,  # Pass the attention mask here
             max_new_tokens=max_new_tokens,
             temperature=0.7,
             top_p=0.9,
             do_sample=True,
         )
 
+        print(f"Model output tokens: {outputs}")  # Debugging line
+
         # Decode and strip input prompt from response
         response = tokenizer.decode(outputs[0], skip_special_tokens=True)
         generated_text = response[len(instruction):].strip()