File size: 4,134 Bytes
e1154a4 | 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 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 | """
HuggingFace Inference Endpoint handler — capture test (T2).
Commit this as `handler.py` with `capture.py` and `requirements.txt` beside it in a model
repo. HuggingFace's Default container detects the handler and runs this instead of the stock
pipeline.
The signals are returned INSIDE the HTTP response under "_capture" rather than posted
anywhere. That is deliberate:
* it proves capture works on this surface even if outbound network is blocked, so the
capture question and the egress question stay independent;
* it needs no API key, no database, and no platform to be up;
* and it happens to be the fallback transport design we would need anyway if T1 fails.
Env vars (all optional):
CAPTURE_ATTENTION=1 load with eager attention and capture attention metrics
(disables SDPA/FlashAttention — measure the cost, do not
leave this on)
CAPTURE_TASK=causal for decoder LMs; default "classification"
CAPTURE_MAX_LAYERS=4 how many transformer blocks to sample
CAPTURE_EXPORT_DIMS=256 how much of the embedding to return
"""
from __future__ import annotations
import os
import time
import traceback
from typing import Any
import torch
from capture import CaptureProbe, load_model
ATTENTION = os.environ.get("CAPTURE_ATTENTION", "0") == "1"
TASK = os.environ.get("CAPTURE_TASK", "classification")
MAX_LAYERS = int(os.environ.get("CAPTURE_MAX_LAYERS", "4"))
EXPORT_DIMS = int(os.environ.get("CAPTURE_EXPORT_DIMS", "256"))
class EndpointHandler:
def __init__(self, path: str = ""):
# HF mounts the selected model repo at `path` (/repository in the container).
self.model, self.tokenizer = load_model(
path or ".", task=TASK, capture_attention=ATTENTION
)
self.probe = CaptureProbe(
self.model, capture_attention=ATTENTION, max_layers=MAX_LAYERS
)
self.env = {
"capture_attention": ATTENTION,
"task": TASK,
"attn_implementation": getattr(self.model.config, "_attn_implementation", "unknown"),
"torch": torch.__version__,
}
def __call__(self, data: dict[str, Any]) -> dict[str, Any]:
text = data.pop("inputs", data)
started = time.perf_counter()
enc = self.tokenizer(text, return_tensors="pt", truncation=True, max_length=256)
try:
with self.probe.capture() as cap:
with torch.no_grad():
out = self.model(
**enc,
output_hidden_states=True,
output_attentions=ATTENTION,
)
signals = self.probe.summarise(
cap, inputs=enc["input_ids"], outputs=out, export_dims=EXPORT_DIMS
)
capture_error = None
except Exception: # noqa: BLE001
# Capture must never take the endpoint down — that is the whole point of
# testing it before anyone depends on it.
signals, capture_error = {}, traceback.format_exc(limit=4)
with torch.no_grad():
out = self.model(**enc)
total_ms = round((time.perf_counter() - started) * 1000, 2)
prediction = self._prediction(out)
return {
"prediction": prediction,
"_capture": signals,
"_meta": {
"total_ms": total_ms,
"env": self.env,
"capture_error": capture_error,
},
}
def _prediction(self, out: Any) -> Any:
logits = getattr(out, "logits", None)
if logits is None:
return None
if logits.dim() == 3: # causal LM
return {"next_token_id": int(logits[0, -1].argmax())}
probs = torch.softmax(logits[0], dim=-1)
id2label = getattr(self.model.config, "id2label", {})
return [
{"label": id2label.get(i, str(i)), "score": round(float(p), 6)}
for i, p in enumerate(probs.tolist())
]
|