Add inference endpoint handler

handler.py

@@ -0,0 +1,85 @@
+"""
+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:
+            # Cargar modelo y tokenizer
+            self.model = AutoModelForSequenceClassification.from_pretrained("vertigoq3/email-classifier-bert")
+            self.tokenizer = AutoTokenizer.from_pretrained("vertigoq3/email-classifier-bert")
+            
+            # Mover al dispositivo
+            self.model.to(self.device)
+            self.model.eval()
+            
+            # Cargar encoder
+            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)
+            
+            # Tokenizar
+            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()}
+            
+            # Clasificar
+            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)}
+
+# Crear instancia global
+handler = EndpointHandler()