Update handler.py

handler.py

@@ -1,6 +1,5 @@
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 import torch
-import json
 
 class EndpointHandler:
     def __init__(self, model_dir):
@@ -9,21 +8,29 @@ class EndpointHandler:
         self.model.eval()
 
     def preprocess(self, data):
-        # Validar entrada
-        if not data or not isinstance(data, dict) or "inputs" not in data or data["inputs"] is None:
-            raise ValueError("La entrada debe ser un diccionario con la clave 'inputs' y un valor válido")
-
+        if not isinstance(data, dict) or "inputs" not in data or data["inputs"] is None:
+            raise ValueError("La entrada debe ser un diccionario con la clave 'inputs' y un valor válido.")
+        
         input_text = "Generate a valid JSON capturing data from this text: " + data["inputs"]
-        if not input_text.strip():
-            raise ValueError("El texto de entrada no puede estar vacío")
-
-        # Tokenización de la entrada
         tokens = self.tokenizer(input_text, return_tensors="pt", truncation=True, padding=True)
+        if not tokens or not tokens["input_ids"]:
+            raise ValueError("Error al tokenizar el texto de entrada. Verifica el texto.")
         return tokens
 
     def inference(self, tokens):
+        generate_kwargs = {
+            "max_length": 1000,
+            "num_beams": 5,
+            "do_sample": True,
+            "temperature": 0.3,
+            "top_k": 50,
+            "top_p": 0.9,
+            "repetition_penalty": 2.5
+        }
         with torch.no_grad():
-            outputs = self.model.generate(**tokens)
+            outputs = self.model.generate(**tokens, **generate_kwargs)
+        if outputs is None or len(outputs) == 0:
+            raise ValueError("El modelo no generó ninguna salida.")
         return outputs
 
     def postprocess(self, outputs):