File size: 3,402 Bytes
43d5672 2a15982 43d5672 2a15982 655f992 d462649 1a10f46 43d5672 9684946 db0321e 1a10f46 e170593 db0321e 0782331 0282e9c 3567112 0282e9c db0321e 0782331 fb2a5f9 0282e9c fb2a5f9 0782331 c703b48 e1dd33d c703b48 0782331 c703b48 0782331 c703b48 0782331 c703b48 0782331 fb2a5f9 c703b48 0782331 c703b48 0782331 758892d 2927fde 2a15982 43d5672 2a15982 43d5672 2a15982 43d5672 2a15982 0782331 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 | 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) |