handler.py

"""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))