updated handler.py to resolve tokenization errors.

handler.py

@@ -31,48 +31,59 @@ class EndpointHandler:
         Returns:
             Generated text
         """
-        # Handle input
-        if isinstance(data.get("inputs"), str):
-            input_text = data["inputs"]
-        else:
-            # Extract messages from input
-            messages = data.get("inputs", {}).get("messages", [])
-            if not messages:
-                return {"error": "No messages provided"}
-                
-            # Format input text
-            input_text = ""
-            for msg in messages:
-                role = msg.get("role", "")
-                content = msg.get("content", "")
-                input_text += f"{role}: {content}\n"
+        try:
+            # Handle input
+            if isinstance(data.get("inputs"), str):
+                input_text = data["inputs"]
+            else:
+                # Extract messages from input
+                messages = data.get("inputs", {}).get("messages", [])
+                if not messages:
+                    return {"error": "No messages provided"}
+                    
+                # Format input text
+                input_text = ""
+                for msg in messages:
+                    role = msg.get("role", "")
+                    content = msg.get("content", "")
+                    input_text += f"{role}: {content}\n"
 
-        # Get generation parameters
-        params = {**self.default_params}
-        if "parameters" in data:
-            params.update(data["parameters"])
+            # Get generation parameters
+            params = {**self.default_params}
+            if "parameters" in data:
+                params.update(data["parameters"])
 
-        # Tokenize input
-        inputs = self.tokenizer(
-            input_text,
-            return_tensors="pt",
-            padding=True,
-            truncation=True,
-            max_length=512
-        )
-
-        # Generate response
-        with torch.no_grad():
-            outputs = self.model.generate(
-                inputs["input_ids"],
-                attention_mask=inputs["attention_mask"],
-                **params
+            # Ensure proper tokenization with padding and attention mask
+            tokenizer_output = self.tokenizer(
+                input_text,
+                return_tensors="pt",
+                padding=True,
+                truncation=True,
+                max_length=512,
+                return_attention_mask=True
             )
 
-        # Decode response
-        generated_text = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+            # Move tensors to the same device as the model
+            input_ids = tokenizer_output["input_ids"]
+            attention_mask = tokenizer_output["attention_mask"]
+
+            # Generate response
+            with torch.no_grad():
+                outputs = self.model.generate(
+                    input_ids,
+                    attention_mask=attention_mask,
+                    pad_token_id=self.tokenizer.pad_token_id,
+                    **params
+                )
 
-        return [{"generated_text": generated_text}]
+            # Decode response
+            generated_text = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+            
+            return [{"generated_text": generated_text}]
+            
+        except Exception as e:
+            print(f"Error in generation: {str(e)}")
+            return {"error": str(e)}
 
     def preprocess(self, request):
         """