File size: 2,368 Bytes
d161eae dac55f6 d161eae dac55f6 d161eae dac55f6 d161eae dac55f6 d161eae dac55f6 d161eae dac55f6 d161eae dac55f6 d161eae dac55f6 d161eae dac55f6 d161eae dac55f6 d161eae dac55f6 |
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 |
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
import torch
class EndpointHandler:
def __init__(self, model_dir):
"""
Inicializa el handler con el modelo y tokenizador.
"""
# Cargar el tokenizador y el modelo desde el directorio proporcionado
self.tokenizer = AutoTokenizer.from_pretrained(model_dir)
self.model = AutoModelForSeq2SeqLM.from_pretrained(model_dir)
self.model.eval() # Poner el modelo en modo evaluación
def preprocess(self, data):
"""
Preprocesa los datos de entrada para el modelo.
"""
# Validar entrada
if not isinstance(data, dict) or "inputs" not in data:
raise ValueError("Entrada inválida. Debe ser un diccionario con la clave 'inputs'.")
input_text = f"Generate a valid JSON capturing data from this text: {data['inputs']}"
# Tokenizar entrada
tokens = self.tokenizer(
input_text,
return_tensors="pt",
truncation=True,
padding="max_length",
max_length=512
)
return tokens
def inference(self, inputs):
"""
Realiza la inferencia con el modelo.
"""
generate_kwargs = {
"max_length": 512,
"num_beams": 5,
"do_sample": False,
"temperature": 0.7,
"top_k": 50,
"top_p": 0.9,
"repetition_penalty": 2.0,
"early_stopping": True # Asegurar que no sea None
}
with torch.no_grad():
outputs = self.model.generate(**inputs, **generate_kwargs)
return outputs
def postprocess(self, model_outputs):
"""
Procesa las salidas del modelo para devolver resultados.
"""
# Decodificar la salida generada por el modelo
decoded_output = self.tokenizer.decode(model_outputs[0], skip_special_tokens=True)
return {"response": decoded_output}
def __call__(self, data):
"""
Ejecuta el pipeline de preprocesamiento, inferencia y postprocesamiento.
"""
# Preprocesar entrada
tokens = self.preprocess(data)
# Realizar inferencia
model_outputs = self.inference(tokens)
# Postprocesar y devolver resultados
return self.postprocess(model_outputs)
|