Create handler.py

handler.py

@@ -0,0 +1,70 @@
+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):
+        self.tokenizer = AutoTokenizer.from_pretrained(model_dir)
+        self.model = AutoModelForSeq2SeqLM.from_pretrained(model_dir)
+        self.model.eval()
+
+    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.")
+
+        # 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)
+        return tokens
+
+    def inference(self, tokens):
+        generate_kwargs = {
+            "max_length": 512,
+            "num_beams": 5,
+            "do_sample": False,
+            "temperature": 0.3,
+            "top_k": 50,
+            "top_p": 0.8,
+            "early_stopping": True,  # Añadir explicitamente esta configuración
+            "repetition_penalty": 2.5
+        }
+        with torch.no_grad():
+            outputs = self.model.generate(**tokens, **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 __call__(self, data):
+        tokens = self.preprocess(data)
+        outputs = self.inference(tokens)
+        result = self.postprocess(outputs)
+        return result
+
+
+# Crear una instancia del handler
+handler = EndpointHandler(model_name)