初始化项目,由ModelHub XC社区提供模型
Model: AttentioResearch/tally-8b-flagship Source: Original Platform
This commit is contained in:
44
handler.py
Normal file
44
handler.py
Normal file
@@ -0,0 +1,44 @@
|
||||
"""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}]
|
||||
Reference in New Issue
Block a user