Update handler.py

handler.py

@@ -5,14 +5,14 @@ import json
 
 model_name = "jla25/squareV3"
 
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=True)
+model = AutoModelForSeq2SeqLM.from_pretrained(model_name, use_auth_token=True)
 
 
 class EndpointHandler:
     def __init__(self, model_dir):
-        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
-        self.model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+        self.tokenizer = AutoTokenizer.from_pretrained(model_dir, use_auth_token=True)
+        self.model = AutoModelForSeq2SeqLM.from_pretrained(model_dir, use_auth_token=True)
         self.model.eval()
 
     def preprocess(self, data):
@@ -20,9 +20,9 @@ class EndpointHandler:
             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 = ({data['inputs']})
-        # Imprimir el texto generado para el prompt
+        input_text = f"{data['inputs']}"
         print(f"Prompt generado para el modelo: {input_text}")
+        
         tokens = self.tokenizer(input_text, return_tensors="pt", truncation=True, padding="max_length", max_length=1024)
         return tokens
 
@@ -41,45 +41,29 @@ class EndpointHandler:
         return outputs
 
     def clean_output(self, output):
-        # Extraer el JSON dentro del texto generado
         try:
             start_index = output.index("{")
             end_index = output.rindex("}") + 1
             return output[start_index:end_index]
         except ValueError:
-            # Si no hay un JSON válido en el texto
             return output
 
-    def validate_json(self, json_text):
-        # Validar el JSON generado
-        try:
-            json_data = json.loads(json_text)
-            if "values" in json_data and isinstance(json_data["values"], list):
-                return {"is_valid": True, "json_data": json_data}
-            else:
-                return {"is_valid": False, "error": "El JSON no contiene el formato esperado."}
-        except json.JSONDecodeError as e:
-            return {"is_valid": False, "error": f"Error decodificando JSON: {str(e)}"}
-
     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: {decoded_output}")
+        print(f"Texto generado por el modelo: {decoded_output}")
         print(f"JSON limpiado: {cleaned_output}")
 
-        # Validar el JSON generado
-        validation_result = self.validate_json(cleaned_output)
-
-        if not validation_result["is_valid"]:
-            print(f"Error en la validación: {validation_result['error']}")
-            raise ValueError(f"JSON inválido: {validation_result['error']}")
-
-        return {"response": validation_result["json_data"]}
+        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)