Upload inference.py

inference.py

@@ -0,0 +1,114 @@
+# inference.py
+
+import os
+import torch
+import joblib
+import numpy as np
+from transformers import BertTokenizer, BertModel
+
+
+class EndpointHandler:
+    """
+    Custom handler for Hugging Face Inference Endpoints.
+
+    Expected input JSON:
+        {"inputs": "some text"}
+      or {"inputs": ["text 1", "text 2", ...]}
+
+    Output:
+        For single input:
+            {"label": "...", "confidence": 0.95}
+        For multiple:
+            [
+                {"label": "...", "confidence": 0.95},
+                {"label": "...", "confidence": 0.80},
+                ...
+            ]
+    """
+
+    def __init__(self, path: str = "."):
+        # 1. Device setup
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        print(f"[handler] Using device: {self.device}")
+
+        # 2. Load BERT
+        print("[handler] Loading BERT tokenizer and model...")
+        self.tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
+        self.bert_model = BertModel.from_pretrained("bert-base-uncased")
+        self.bert_model.to(self.device)
+        self.bert_model.eval()
+
+        # 3. Load MLP, scaler, label encoder
+        print("[handler] Loading classification components...")
+        mlp_path = os.path.join(path, "mlp_query_classifier.joblib")
+        scaler_path = os.path.join(path, "scaler_query_classifier.joblib")
+        le_path = os.path.join(path, "label_encoder_query_classifier.joblib")
+
+        self.mlp = joblib.load(mlp_path)
+        self.scaler = joblib.load(scaler_path)
+        self.le = joblib.load(le_path)
+
+        print("[handler] Loaded MLP, scaler, and label encoder.")
+
+    # ------------ Helper: BERT embeddings ------------
+    def get_bert_embeddings(self, text_list):
+        inputs = self.tokenizer(
+            text_list,
+            padding=True,
+            truncation=True,
+            max_length=128,
+            return_tensors="pt"
+        ).to(self.device)
+
+        with torch.no_grad():
+            outputs = self.bert_model(**inputs)
+
+        # CLS token embedding
+        cls_embeddings = outputs.last_hidden_state[:, 0, :]
+        return cls_embeddings.cpu().numpy()
+
+    # ------------ Main entry point ------------
+    def __call__(self, data):
+        """
+        data: dict with key "inputs"
+        """
+        if "inputs" not in data:
+            raise ValueError("Input JSON must have an 'inputs' field.")
+
+        texts = data["inputs"]
+
+        # Normalize to list
+        is_single = False
+        if isinstance(texts, str):
+            texts = [texts]
+            is_single = True
+
+        # 1) BERT embedding
+        embeddings = self.get_bert_embeddings(texts)
+
+        # 2) Scale with same scaler as training
+        embeddings_scaled = self.scaler.transform(embeddings)
+
+        # 3) Predict class indices
+        pred_indices = self.mlp.predict(embeddings_scaled)
+
+        # 4) Map indices to labels
+        labels = self.le.inverse_transform(pred_indices)
+
+        # 5) Optionally, get probabilities
+        results = []
+        for i, idx in enumerate(pred_indices):
+            label = labels[i]
+            try:
+                probs = self.mlp.predict_proba(embeddings_scaled[i : i + 1])[0]
+                confidence = float(np.max(probs))
+            except Exception:
+                confidence = None
+
+            result = {"label": label}
+            results.append(result)
+
+        # If the user sent a single string, return a single dict
+        if is_single:
+            return results[0]
+        return results