| """ |
| 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 = ""): |
| |
| 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: |
| |
| |
| 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: |
| 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()) |
| ] |
|
|
