Update handler.py

handler.py

@@ -1,23 +1,65 @@
 from transformers import AutoModelForCausalLM, AutoTokenizer
+from jsonschema import validate, ValidationError
 from typing import Dict, Any, List
+import jsonschema
 import torch
 import json
 
-def model_inference(model, tokenizer, text, **kwargs):
-    default_kwargs = {
-        "max_new_tokens": 1024,
+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"]
+}
+
+def parse_and_validate_json(output):
+    output = output.replace("\n\n", "\n").strip()
+    outputs = output.split("\n")
+    
+    for single_output in outputs:
+        try:
+            corrected_json = single_output.replace('}\n"', '}, "')
+            data = json.loads(corrected_json)
+            validate(instance=data, schema=schema)
+            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):
+    default_parameters = {
+        "max_new_tokens": 300,
         "temperature": 0.5,
-        "top_p": 0.95,
-        "top_k": 150,
-        "repetition_penalty": 1.4
+        "top_p": 0.99,
+        "top_k": 20,
+        "repetition_penalty": 1.2
     }
-    default_kwargs.update(kwargs)
+    default_parameters.update(parameters)
     inputs = tokenizer(text, return_tensors="pt").to("cuda:0")
     inputs_length = len(inputs["input_ids"][0])
     with torch.inference_mode():
-        outputs = model.generate(**inputs, **default_kwargs)
+        outputs = model.generate(**inputs, **default_parameters)
     return tokenizer.decode(outputs[0][inputs_length:], skip_special_tokens=True)
 
+
+def model_inference_with_validation(filled_template, **parameters):
+    while True:
+        result = model_inference(self.model, self.my_tokenizer, filled_template, **parameters)
+        is_valid, response_or_error = parse_and_validate_json(result)
+        if is_valid:
+            return {'generated_text': [result]}
+
+
+
 class EndpointHandler():
     def __init__(self, path=""):
         self.model = AutoModelForCausalLM.from_pretrained(
@@ -31,14 +73,17 @@ class EndpointHandler():
     def __call__(self, data: Dict[str, Any]) -> Dict[str, str]:
         input_data = data.pop("inputs", data)
         parameters = data.pop("parameters", {})
+
         template = """
-        <s>[INST] <<SYS>>
+        [INST] <<SYS>>
         A continuación se presentan parámetros fisiológicos y ambientales capturados por los sensores de un smartwatch utilizado por un atleta. Basándote en estos datos y en tu conocimiento de las rutinas de entrenamiento para atletas, crea una rutina diaria de entrenamiento personalizada para el atleta.
         <</SYS>>
         {data} [/INST]
         """.strip()
+        
         data_text = json.dumps(input_data, indent=4)
         filled_template = template.replace("{data}", data_text)
-        result = model_inference(self.model, self.my_tokenizer, filled_template, **parameters)
-        
-        return {'generated_text': [result]}
+
+        return model_inference_with_validation(filled_template, **parameters)
+
+