45 lines
2.5 KiB
Python
45 lines
2.5 KiB
Python
|
"""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}]
|