handler.py

"""HuggingFace Inference Endpoints custom handler for the Tally adherence package.

The checkpoint is NOT a standard model — the real entry point is `AdherenceModel` (a wrapper in
modeling_adherence.py: baked weights + deterministic guard + scope gate + attack cutoff), and it is NOT a
PreTrainedModel and has no `auto_map`. So the default TGI / transformers handler cannot serve the guarded
stack — it would load a plain Qwen3ForCausalLM (weights only, no guards) or fail. This handler loads the
real class and calls `.chat()`, so the endpoint serves the FULL product.

To use it, the Inference Endpoint must be created with task = "Custom" (so it picks up handler.py); a
Text-Generation / TGI task ignores this file. GPU required; device_map="auto" so the 8B shards across
multiple small GPUs (e.g. 4x T4 = 64GB) instead of OOMing on a single 16GB card.
"""
from __future__ import annotations

import importlib.util
import os
import sys
from typing import Any, Dict, List


class EndpointHandler:
    def __init__(self, path: str = "") -> None:
        spec = importlib.util.spec_from_file_location("modeling_adherence",
                                                      os.path.join(path, "modeling_adherence.py"))
        ma = importlib.util.module_from_spec(spec)
        # register BEFORE exec — modeling_adherence uses `from __future__ import annotations`, so @dataclass
        # resolves its field types via sys.modules[cls.__module__]; unregistered => NoneType.__dict__ crash.
        sys.modules["modeling_adherence"] = ma
        spec.loader.exec_module(ma)
        self.model = ma.AdherenceModel.from_pretrained(path, torch_dtype="auto", device_map="auto")

    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, str]]:
        inputs = data.get("inputs", data)
        if isinstance(inputs, str):                       # plain prompt
            messages = [{"role": "user", "content": inputs}]
        elif isinstance(inputs, list):                    # OpenAI-style chat messages
            messages = [{"role": m.get("role", "user"), "content": m.get("content", "")}
                        if isinstance(m, dict) else {"role": "user", "content": str(m)} for m in inputs]
        else:
            messages = [{"role": "user", "content": str(inputs)}]
        params = data.get("parameters") or {}
        out = self.model.chat(messages, max_new_tokens=int(params.get("max_new_tokens", 256)),
                              temperature=params.get("temperature"))
        return [{"generated_text": out}]