Upload handler.py with huggingface_hub

handler.py

@@ -5,7 +5,8 @@ import torch
 
 class EndpointHandler:
     """
-    Custom handler for DoloresAI model on HuggingFace Inference Endpoints.
+    Custom handler for DoloresAI model - GREEDY DECODING ONLY
+    This avoids sampling issues with resized embeddings.
     """
 
     def __init__(self, path=""):
@@ -30,18 +31,14 @@ class EndpointHandler:
 
     def __call__(self, data: Dict[str, Any]) -> List[Dict[str, str]]:
         """
-        Process inference requests.
+        Process inference requests using GREEDY DECODING ONLY.
 
         Args:
             data (Dict): Input data with format:
                 {
                     "inputs": str,  # The prompt text
                     "parameters": {  # Optional generation parameters
-                        "max_new_tokens": int,
-                        "temperature": float,
-                        "top_p": float,
-                        "do_sample": bool,
-                        "repetition_penalty": float
+                        "max_new_tokens": int
                     }
                 }
 
@@ -52,13 +49,9 @@ class EndpointHandler:
         inputs = data.pop("inputs", data)
         parameters = data.pop("parameters", {})
 
-        # Default generation parameters - use greedy decoding to avoid sampling issues
+        # Get max tokens (only parameter we use)
         max_new_tokens = parameters.get("max_new_tokens", 512)
 
-        # 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(
             inputs,
@@ -67,13 +60,14 @@ class EndpointHandler:
             max_length=self.model.config.max_position_embeddings - max_new_tokens
         ).input_ids.to(self.model.device)
 
-        # Generate response with greedy decoding (no sampling)
+        # Generate response with GREEDY DECODING ONLY
+        # This is stable and avoids NaN/inf errors from sampling
         with torch.no_grad():
             outputs = self.model.generate(
                 input_ids,
                 max_new_tokens=max_new_tokens,
-                do_sample=False,  # Greedy decoding - most stable
-                num_beams=1,  # No beam search for speed
+                do_sample=False,  # GREEDY - no sampling
+                num_beams=1,  # No beam search
                 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,
             )