Update handler.py

handler.py

@@ -5,18 +5,6 @@ import jsonschema
 import torch
 import json
 
-schema = {
-    "type": "object",
-    "properties": {
-        "objetivoDistancia": {"type": "string"},
-        "calentamiento": {"type": "string"},
-        "drills": {"type": "array", "items": {"type": "object"}},
-        "sprints": {"type": "array", "items": {"type": "object"}},
-        "enfriamiento": {"type": "string"}
-    },
-    "required": ["objetivoDistancia", "calentamiento", "drills", "sprints", "enfriamiento"]
-}
-
 class EndpointHandler():
     def __init__(self, path=""):
         self.model = AutoModelForCausalLM.from_pretrained(
@@ -25,9 +13,9 @@ class EndpointHandler():
             torch_dtype=torch.bfloat16,
             device_map="auto"
         )
-        self.my_tokenizer = AutoTokenizer.from_pretrained(path, trust_remote_code=True)
+        self.tokenizer = AutoTokenizer.from_pretrained(path, trust_remote_code=True)
 
-    def parse_and_validate_json(output):
+    def parse_and_validate_json(self, output):
         output = output.replace("\n\n", "\n").strip()
         outputs = output.split("\n")
         
@@ -39,10 +27,9 @@ class EndpointHandler():
                 return True, corrected_json
             except (json.JSONDecodeError, ValidationError):
                 continue
-      
         return False, "No se encontró ningún JSON válido según el esquema requerido."
 
-    def model_inference(model, tokenizer, text, **parameters):
+    def model_inference(self, text, **parameters):
         default_parameters = {
             "max_new_tokens": 300,
             "temperature": 0.5,
@@ -51,20 +38,22 @@ class EndpointHandler():
             "repetition_penalty": 1.2
         }
         default_parameters.update(parameters)
-        inputs = tokenizer(text, return_tensors="pt").to("cuda:0")
+        inputs = self.tokenizer(text, return_tensors="pt").to("cuda:0")
         inputs_length = len(inputs["input_ids"][0])
         with torch.inference_mode():
-            outputs = model.generate(**inputs, **default_parameters)
-        return tokenizer.decode(outputs[0][inputs_length:], skip_special_tokens=True)
+            outputs = self.model.generate(**inputs, **default_parameters)
+        return self.tokenizer.decode(outputs[0][inputs_length:], skip_special_tokens=True)
 
-    def model_inference_with_validation(filled_template, **parameters):
+    def model_inference_with_validation(self, filled_template, **parameters):
         while True:
-            result = self.model_inference(self.model, self.my_tokenizer, filled_template, **parameters)
+            result = self.model_inference(filled_template, **parameters)
             is_valid, response_or_error = self.parse_and_validate_json(result)
             if is_valid:
-                return {'generated_text': [result]}
+                return {'generated_text': [response_or_error]}
+            else:
+                continue
 
-    def __call__(self, data: Dict[str, Any]) -> Dict[str, str]:
+    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
         input_data = data.pop("inputs", data)
         parameters = data.pop("parameters", {})
 
@@ -78,6 +67,4 @@ class EndpointHandler():
         data_text = json.dumps(input_data, indent=4)
         filled_template = template.replace("{data}", data_text)
 
-        return self.model_inference_with_validation(filled_template, **parameters)
-
-        
+        return self.model_inference_with_validation(filled_template, **parameters)