| """Custom handler for Hugging Face Inference Endpoints. |
| |
| HF creates EndpointHandler(path) once when the endpoint starts, then calls it |
| for every request with the JSON body, e.g. {"inputs": "hello how are you"}. |
| """ |
| import os |
| from typing import Any, Dict, List, Union |
|
|
| import joblib |
| from sentence_transformers import SentenceTransformer |
|
|
| ENCODER_NAME = "sentence-transformers/all-MiniLM-L6-v2" |
|
|
| |
| |
| REPLIES = { |
| "greeting": "hello", |
| "computer": "output", |
| "other": "I don't understand", |
| } |
|
|
|
|
| class EndpointHandler: |
| def __init__(self, path: str = ""): |
| base = path or os.path.dirname(os.path.abspath(__file__)) |
| self.encoder = SentenceTransformer(ENCODER_NAME) |
| self.clf = joblib.load(os.path.join(base, "classifier.joblib")) |
|
|
| def _predict(self, texts: List[str]) -> List[Dict[str, Any]]: |
| if not texts: |
| return [] |
| embeddings = self.encoder.encode(texts, normalize_embeddings=True) |
| probs = self.clf.predict_proba(embeddings) |
| results = [] |
| for row in probs: |
| best = int(row.argmax()) |
| label = str(self.clf.classes_[best]) |
| results.append({ |
| "reply": REPLIES[label], |
| "label": label, |
| "score": round(float(row[best]), 4), |
| }) |
| return results |
|
|
| def __call__(self, data: Dict[str, Any]) -> Union[Dict[str, Any], List[Dict[str, Any]]]: |
| if not isinstance(data, dict) or "inputs" not in data: |
| raise ValueError('Send JSON like {"inputs": "hi"}') |
| inputs = data["inputs"] |
| if isinstance(inputs, list): |
| return self._predict([str(t) for t in inputs]) |
| return self._predict([str(inputs)])[0] |
|
|
