Upload handler.py with huggingface_hub

handler.py

@@ -52,12 +52,12 @@ class EndpointHandler:
         inputs = data.pop("inputs", data)
         parameters = data.pop("parameters", {})
 
-        # Default generation parameters with safe values
+        # Default generation parameters - use greedy decoding to avoid sampling issues
         max_new_tokens = parameters.get("max_new_tokens", 512)
-        temperature = max(0.1, min(parameters.get("temperature", 0.8), 2.0))  # Clamp between 0.1 and 2.0
-        top_p = max(0.1, min(parameters.get("top_p", 0.95), 1.0))  # Clamp between 0.1 and 1.0
-        do_sample = parameters.get("do_sample", True)
-        repetition_penalty = max(1.0, min(parameters.get("repetition_penalty", 1.05), 2.0))  # Clamp between 1.0 and 2.0
+
+        # Use greedy decoding (do_sample=False) to avoid probability tensor issues
+        # This is more stable for models with potential embedding issues
+        do_sample = False  # Force greedy decoding
 
         # Tokenize input
         input_ids = self.tokenizer(
@@ -67,20 +67,15 @@ class EndpointHandler:
             max_length=self.model.config.max_position_embeddings - max_new_tokens
         ).input_ids.to(self.model.device)
 
-        # Generate response with safe parameters
+        # Generate response with greedy decoding (no sampling)
         with torch.no_grad():
             outputs = self.model.generate(
                 input_ids,
                 max_new_tokens=max_new_tokens,
-                temperature=temperature,
-                top_p=top_p,
-                top_k=50,  # Add top_k for stability
-                do_sample=do_sample,
-                repetition_penalty=repetition_penalty,
+                do_sample=False,  # Greedy decoding - most stable
+                num_beams=1,  # No beam search for speed
                 pad_token_id=self.tokenizer.pad_token_id if self.tokenizer.pad_token_id is not None else self.tokenizer.eos_token_id,
                 eos_token_id=self.tokenizer.eos_token_id,
-                bad_words_ids=None,  # Ensure no bad words restriction causing issues
-                min_length=1,  # Ensure at least 1 token is generated
             )
 
         # Decode output