File size: 1,812 Bytes
c57c722 4d8a834 c57c722 133bed4 4d8a834 a7f3b7d 133bed4 c57c722 24a1035 c57c722 24a1035 c57c722 24a1035 c57c722 |
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 |
from typing import Dict, List, Any
from transformers import pipeline
import logging
# Configuración básica del logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
class EndpointHandler():
def __init__(self, path=""):
logger.info(f"Modelo cargado desde el path: {path}")
# Cargar el pipeline con el modelo especificado
try:
self.pipeline = pipeline(model=path, truncation=True)
except Exception as e:
logger.exception(f"Error cargando el modelo desde el path {path}: {e}")
path = "AndresR2909/finetuning-bert-text-classification"
self.pipeline = pipeline(model=path, truncation=True)
def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
"""
data args:
inputs (:obj: `str`)
date (:obj: `str`)
Return:
A :obj:`list` | `dict`: will be serialized and returned
"""
# get inputs
input = data.get("inputs",data)
date = data.get("date", None)
# Textos para análisis
texts = [input]
# Realizar la inferencia
outputs = self.pipeline(texts)
# Ajustar las etiquetas a 0 y 1
label_mapping = {"LABEL_0": 0, "LABEL_1": 1}
label_names = {0: "sin_intencion", 1: "intencion_suicida"}
# Modificar las salidas con el formato deseado
adjusted_results = [
{
"input": text,
"clasiffication": str(label_mapping[result['label']]), # "clasiffication" como string
"label": label_names[label_mapping[result['label']]] # Mapeo de etiquetas a nombres
}
for result, text in zip(outputs, texts)
]
return adjusted_results |