File size: 1,074 Bytes
1772d91 | 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 | """Hugging Face Inference Endpoints entry point — deploy this repo as a CPU/GPU API.
Request: {"inputs": "Subject: <subject>\n\n<body>"} -> {"spam": p, "ham": 1 - p}
{"inputs": ["Subject: ...", "Subject: ..."]} -> [{"spam": p, "ham": 1 - p}, ...]
"""
import sys
from pathlib import Path
HERE = Path(__file__).resolve().parent
sys.path.insert(0, str(HERE))
import model as M # noqa: E402
class EndpointHandler:
def __init__(self, path: str = ""):
self.predictor = M.load(path or HERE, "cuda" if M.cuda_available() else "cpu")
def __call__(self, data: dict):
inputs = data.pop("inputs", data)
if isinstance(inputs, dict): # {"subject": ..., "body": ...}
inputs = M.compose_email(inputs.get("subject", ""), inputs.get("body", inputs.get("text", "")))
if isinstance(inputs, (list, tuple)):
probs = self.predictor.predict_proba([str(t) for t in inputs])
return [{"spam": p, "ham": 1.0 - p} for p in probs]
return self.predictor.predict(str(inputs))
|