File size: 2,725 Bytes
1de92bf | 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 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 | """
Hugging Face Inference Endpoint custom handler.
Loads a base model with multiple LoRA adapters (one per mode).
Adapters stay resident in memory; set_adapter() switches cheaply per request.
The model is trained to emit <ui>{...json...}</ui> then prose.
"""
import json
import os
import re
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel
BASE_MODEL = os.getenv("BASE_MODEL", "Qwen/Qwen2.5-7B-Instruct")
# Override via env vars to point at your own adapter repos on the Hub
ADAPTERS: dict[str, str] = {
"support": os.getenv("ADAPTER_SUPPORT", "your-org/adapter-support"),
"analytics": os.getenv("ADAPTER_ANALYTICS", "your-org/adapter-analytics"),
"form": os.getenv("ADAPTER_FORM", "your-org/adapter-form"),
}
_UI_RE = re.compile(r"<ui>(.*?)</ui>", re.DOTALL)
class EndpointHandler:
def __init__(self, path: str = ""):
self.tok = AutoTokenizer.from_pretrained(BASE_MODEL)
base = AutoModelForCausalLM.from_pretrained(
BASE_MODEL,
torch_dtype=torch.bfloat16,
device_map="auto",
)
names = list(ADAPTERS.keys())
# First adapter via from_pretrained so PEFT wraps the model
self.model = PeftModel.from_pretrained(
base, ADAPTERS[names[0]], adapter_name=names[0]
)
# Remaining adapters attached in-place
for name in names[1:]:
self.model.load_adapter(ADAPTERS[name], adapter_name=name)
self.model.eval()
def __call__(self, data: dict) -> dict:
inp = data.get("inputs", {})
messages = inp.get("messages", [])
mode = inp.get("mode", "support")
if mode not in ADAPTERS:
mode = "support"
self.model.set_adapter(mode)
prompt = self.tok.apply_chat_template(
messages, add_generation_prompt=True, tokenize=False
)
ids = self.tok(prompt, return_tensors="pt").to(self.model.device)
with torch.no_grad():
out = self.model.generate(
**ids,
max_new_tokens=512,
do_sample=False,
)
new_tokens = out[0][ids.input_ids.shape[1]:]
text = self.tok.decode(new_tokens, skip_special_tokens=True)
ui_spec, clean_text = _extract_ui_spec(text)
return {"text": clean_text, "ui_spec": ui_spec, "adapter": mode}
def _extract_ui_spec(text: str) -> tuple[dict | None, str]:
m = _UI_RE.search(text)
if not m:
return None, text
try:
spec = json.loads(m.group(1).strip())
except json.JSONDecodeError:
return None, text
return spec, _UI_RE.sub("", text).strip()
|