|
|
"""
|
|
|
Handler para el Inference Endpoint del clasificador de emails
|
|
|
"""
|
|
|
|
|
|
import torch
|
|
|
import numpy as np
|
|
|
import pickle
|
|
|
from transformers import AutoTokenizer, AutoModelForSequenceClassification
|
|
|
from huggingface_hub import hf_hub_download
|
|
|
|
|
|
class EndpointHandler:
|
|
|
def __init__(self):
|
|
|
self.model = None
|
|
|
self.tokenizer = None
|
|
|
self.encoder = None
|
|
|
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
|
self.load_model()
|
|
|
|
|
|
def load_model(self):
|
|
|
"""Cargar el modelo"""
|
|
|
try:
|
|
|
|
|
|
self.model = AutoModelForSequenceClassification.from_pretrained("vertigoq3/email-classifier-bert")
|
|
|
self.tokenizer = AutoTokenizer.from_pretrained("vertigoq3/email-classifier-bert")
|
|
|
|
|
|
|
|
|
self.model.to(self.device)
|
|
|
self.model.eval()
|
|
|
|
|
|
|
|
|
encoder_path = hf_hub_download(
|
|
|
repo_id="vertigoq3/email-classifier-bert",
|
|
|
filename="label_encoder.pkl"
|
|
|
)
|
|
|
|
|
|
with open(encoder_path, "rb") as f:
|
|
|
self.encoder = pickle.load(f)
|
|
|
|
|
|
except Exception as e:
|
|
|
print(f"Error al cargar modelo: {e}")
|
|
|
raise
|
|
|
|
|
|
def __call__(self, inputs):
|
|
|
"""Procesar una solicitud de inferencia"""
|
|
|
try:
|
|
|
if isinstance(inputs, str):
|
|
|
text = inputs
|
|
|
elif isinstance(inputs, dict) and "inputs" in inputs:
|
|
|
text = inputs["inputs"]
|
|
|
else:
|
|
|
text = str(inputs)
|
|
|
|
|
|
|
|
|
tokenized = self.tokenizer(
|
|
|
text,
|
|
|
return_tensors="pt",
|
|
|
truncation=True,
|
|
|
padding=True,
|
|
|
max_length=512
|
|
|
)
|
|
|
tokenized = {k: v.to(self.device) for k, v in tokenized.items()}
|
|
|
|
|
|
|
|
|
with torch.no_grad():
|
|
|
outputs = self.model(**tokenized)
|
|
|
logits = outputs.logits
|
|
|
probabilities = torch.softmax(logits, dim=-1)
|
|
|
predicted_class_id = torch.argmax(probabilities, dim=-1).item()
|
|
|
predicted_class = self.encoder.inverse_transform([predicted_class_id])[0]
|
|
|
confidence = float(probabilities[0][predicted_class_id])
|
|
|
|
|
|
return {
|
|
|
"predicted_class": predicted_class,
|
|
|
"confidence": confidence,
|
|
|
"all_probabilities": {
|
|
|
self.encoder.classes_[i]: float(probabilities[0][i])
|
|
|
for i in range(len(self.encoder.classes_))
|
|
|
}
|
|
|
}
|
|
|
|
|
|
except Exception as e:
|
|
|
return {"error": str(e)}
|
|
|
|
|
|
|
|
|
handler = EndpointHandler()
|
|
|
|
