| from transformers import AutoModelForCausalLM, AutoTokenizer |
| from jsonschema import validate, ValidationError |
| from typing import Dict, Any, List |
| import jsonschema |
| import torch |
| import json |
|
|
| class EndpointHandler(): |
| def __init__(self, path=""): |
| self.model = AutoModelForCausalLM.from_pretrained( |
| path, |
| trust_remote_code=True, |
| torch_dtype=torch.bfloat16, |
| device_map="auto" |
| ) |
| self.tokenizer = AutoTokenizer.from_pretrained(path, trust_remote_code=True) |
| self.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": ["calentamiento", "drills", "enfriamiento"] |
| } |
|
|
| def parse_and_validate_json(self, 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=self.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(self, text, **parameters): |
| default_parameters = { |
| "max_new_tokens": 512, |
| "temperature": 0.5, |
| "top_p": 0.99, |
| "top_k": 20, |
| "repetition_penalty": 1.2 |
| } |
| default_parameters.update(parameters) |
| inputs = self.tokenizer(text, return_tensors="pt").to("cuda:0") |
| inputs_length = len(inputs["input_ids"][0]) |
| with torch.inference_mode(): |
| outputs = self.model.generate(**inputs, **default_parameters) |
| return self.tokenizer.decode(outputs[0][inputs_length:], skip_special_tokens=True) |
|
|
| def model_inference_with_validation(self, filled_template, **parameters): |
| while True: |
| result = self.model_inference(filled_template, **parameters) |
| is_valid, response_or_error = self.parse_and_validate_json(result) |
| if is_valid: |
| return {'generated_text': [response_or_error]} |
| else: |
| continue |
|
|
| def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]: |
| input_data = data.pop("inputs", data) |
| parameters = data.pop("parameters", {}) |
|
|
| template = """ |
| [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) |
|
|
| return self.model_inference_with_validation(filled_template, **parameters) |
