Update handler.py

handler.py

@@ -1,69 +1,68 @@
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 import torch
-import json
-
-
-model_name = "jla25/squareV4"
-
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
-
 
 class EndpointHandler:
     def __init__(self, model_dir):
+        """
+        Inicializa el handler con el modelo y tokenizador.
+        """
+        # Cargar el tokenizador y el modelo desde el directorio proporcionado
         self.tokenizer = AutoTokenizer.from_pretrained(model_dir)
         self.model = AutoModelForSeq2SeqLM.from_pretrained(model_dir)
-        self.model.eval()
+        self.model.eval()  # Poner el modelo en modo evaluación
 
     def preprocess(self, data):
-        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.")
+        """
+        Preprocesa los datos de entrada para el modelo.
+        """
+        # Validar entrada
+        if not isinstance(data, dict) or "inputs" not in data:
+            raise ValueError("Entrada inválida. Debe ser un diccionario con la clave 'inputs'.")
 
-        # Prompt personalizado para guiar al modelo
-        input_text = f"Generate a valid JSON capturing data from this text:{data['inputs']}"
-        print(f"Prompt generado para el modelo: {input_text}")
-        input_text = input_text.encode("utf-8").decode("utf-8")
-        tokens = self.tokenizer(input_text, return_tensors="pt", truncation=True, padding="max_length", max_length=1024)
+        input_text = f"Generate a valid JSON capturing data from this text: {data['inputs']}"
+        # Tokenizar entrada
+        tokens = self.tokenizer(
+            input_text,
+            return_tensors="pt",
+            truncation=True,
+            padding="max_length",
+            max_length=512
+        )
         return tokens
 
-    def inference(self, tokens):
+    def inference(self, inputs):
+        """
+        Realiza la inferencia con el modelo.
+        """
         generate_kwargs = {
             "max_length": 512,
             "num_beams": 5,
             "do_sample": False,
-            "temperature": 0.3,
+            "temperature": 0.7,
             "top_k": 50,
-            "top_p": 0.8,
-            "repetition_penalty": 2.5
+            "top_p": 0.9,
+            "repetition_penalty": 2.0,
+            "early_stopping": True  # Asegurar que no sea None
         }
         with torch.no_grad():
-            outputs = self.model.generate(**tokens, **generate_kwargs)
+            outputs = self.model.generate(**inputs, **generate_kwargs)
         return outputs
 
-    def clean_output(self, output):
-        try:
-            start_index = output.index("{")
-            end_index = output.rindex("}") + 1
-            return output[start_index:end_index]
-        except ValueError:
-            return output
-
-    def postprocess(self, outputs):
-        decoded_output = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
-        cleaned_output = self.clean_output(decoded_output)
-
-        # Imprimir siempre el texto generado para depuración
-        print(f"Texto generado por el modelo: {decoded_output}")
-        print(f"JSON limpiado: {cleaned_output}")
-
-        return {"response": cleaned_output}
+    def postprocess(self, model_outputs):
+        """
+        Procesa las salidas del modelo para devolver resultados.
+        """
+        # Decodificar la salida generada por el modelo
+        decoded_output = self.tokenizer.decode(model_outputs[0], skip_special_tokens=True)
+        return {"response": decoded_output}
 
     def __call__(self, data):
+        """
+        Ejecuta el pipeline de preprocesamiento, inferencia y postprocesamiento.
+        """
+        # Preprocesar entrada
         tokens = self.preprocess(data)
-        outputs = self.inference(tokens)
-        result = self.postprocess(outputs)
-        return result
-
-
-# Crear una instancia del handler
-handler = EndpointHandler(model_name)
+        # Realizar inferencia
+        model_outputs = self.inference(tokens)
+        # Postprocesar y devolver resultados
+        return self.postprocess(model_outputs)