Fix: Add pad_token and attention_mask for proper generation"

handler.py

@@ -15,6 +15,10 @@ class EndpointHandler:
         
         print(f"Loading tokenizer from {path}...")
         self.tokenizer = AutoTokenizer.from_pretrained(path)
+
+        # ✅ ADD THIS: Set pad token to prevent corruption
+        if self.tokenizer.pad_token is None:
+            self.tokenizer.pad_token = self.tokenizer.eos_token
         
         print(f"Loading model from {path} on device: {self.device}...")
         self.model = AutoModelForCausalLM.from_pretrained(
@@ -62,11 +66,15 @@ class EndpointHandler:
         print(f"Generating with parameters: {gen_kwargs}")
         
         # Tokenize input
-        enc = self.tokenizer(inputs, return_tensors="pt").to(self.device)
+        enc = self.tokenizer(inputs, return_tensors="pt",padding=True,           # Enable padding
+        truncation=True,        #  Truncate if needed
+        max_length=2048).to(self.device)
         
         # Generate with no_grad for efficiency
         with torch.no_grad():
-            out = self.model.generate(**enc, **gen_kwargs)
+            out = self.model.generate(**enc, 
+                                      **gen_kwargs,
+                                      pad_token_id=self.tokenizer.pad_token_id) # Tell model which token is padding
         
         # Decode output
         generated_text = self.tokenizer.decode(out[0], skip_special_tokens=True)