File size: 1,884 Bytes
70adf2b | 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 52 | """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"
# What to say for each label the classifier can predict.
# Every label used in data.csv needs an entry here.
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): # several texts -> list of results
return self._predict([str(t) for t in inputs])
return self._predict([str(inputs)])[0] # one text -> one result
|