Create handler.py

handler.py

@@ -0,0 +1,40 @@
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+import torch
+
+class ModelHandler:
+    def __init__(self, model_dir):
+        # Cargar el modelo y el tokenizador desde el directorio del modelo
+        self.tokenizer = AutoTokenizer.from_pretrained(model_dir)
+        self.model = AutoModelForSeq2SeqLM.from_pretrained(model_dir)
+        self.model.eval()  # Configurar el modelo en modo de evaluación
+
+    def preprocess(self, data):
+        # Preprocesamiento de la entrada
+        if isinstance(data, dict) and "input_text" in data:
+            input_text = data["input_text"]
+        else:
+            raise ValueError("Esperando un diccionario con la clave 'inputs'")
+
+        # Tokenización de la entrada
+        tokens = self.tokenizer(input_text, return_tensors="pt", truncation=True, padding=True)
+        return tokens
+
+    def inference(self, tokens):
+        # Realizar la inferencia
+        with torch.no_grad():
+            outputs = self.model(**tokens)
+            # Obtener las predicciones y aplicar softmax para probabilidades
+            probabilities = torch.softmax(outputs.logits, dim=-1)
+        return probabilities
+
+    def postprocess(self, probabilities):
+        # Postprocesamiento para devolver la salida en formato JSON
+        predictions = torch.argmax(probabilities, dim=-1)
+        return {"predictions": predictions.tolist(), "probabilities": probabilities.tolist()}
+
+    def __call__(self, data):
+        # Llamada principal del handler para procesamiento completo
+        tokens = self.preprocess(data)
+        probabilities = self.inference(tokens)
+        result = self.postprocess(probabilities)
+        return result