update handler.py for debugging

tensor size mismatch error typically occurs when there's an inconsistency between model configuration and input processing. Fixing by:
Uses fixed padding with max_length
Adds debug printing
Simplifies input handling
Uses consistent tensor dimensions

handler.py

@@ -30,54 +30,46 @@ class EndpointHandler:
             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 as array
-                inputs = []
-                for msg in messages:
-                    role = msg.get("role", "")
-                    content = msg.get("content", "")
-                    inputs.append(f"{role}: {content}")
-                input_text = "\n".join(inputs)
+                input_text = data.get("inputs")[0] if isinstance(data.get("inputs"), list) else str(data.get("inputs"))
 
-            # Get generation parameters
-            params = {**self.default_params}
-            if "parameters" in data:
-                params.update(data["parameters"])
-                
-            # Remove pad_token_id from params if it's going to be set explicitly
-            params.pop('pad_token_id', None)
+            # Print debug information
+            print(f"Input text: {input_text}")
 
-            # Tokenize input
+            # Tokenize with fixed dimensions
             tokenizer_output = self.tokenizer(
                 input_text,
-                return_tensors="pt",
-                padding=True,
+                padding='max_length',  # Changed to max_length
                 truncation=True,
-                max_length=512,
+                max_length=512,  # Fixed length
+                return_tensors="pt",
                 return_attention_mask=True
             )
 
+            # Print tensor shapes for debugging
+            print(f"Input ids shape: {tokenizer_output['input_ids'].shape}")
+            print(f"Attention mask shape: {tokenizer_output['attention_mask'].shape}")
+
             # Generate response
             with torch.no_grad():
                 outputs = self.model.generate(
                     tokenizer_output["input_ids"],
                     attention_mask=tokenizer_output["attention_mask"],
-                    pad_token_id=self.tokenizer.pad_token_id,  # Set it only here
-                    **params
+                    max_length=512,
+                    pad_token_id=self.tokenizer.pad_token_id,
+                    do_sample=True,
+                    temperature=0.7,
+                    top_p=0.7,
+                    top_k=50
                 )
 
-            # Decode response and ensure array output
+            # Decode response
             generated_text = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
             
-            # Always return an array as required by the endpoint
             return [{"generated_text": generated_text}]
             
         except Exception as e:
             print(f"Error in generation: {str(e)}")
+            print(f"Model config: {self.model.config}")
             return {"error": str(e)}
 
     def preprocess(self, request):