File size: 3,162 Bytes
dbb1d3a | 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 | """
handler.py — HuggingFace Inference Endpoint handler for SriRamanaAtmic/AtmicReRanker (BERT-large
cross-encoder, BertForSequenceClassification, 2 labels). Upload to the root of the model repo.
Scores match backend/reranker.py exactly, so the backend can fall back to local CPU scoring without
changing results: same num_labels=2 load, same tokenisation ([query, passage] pairs, passage-only
truncation at 512 tokens), fp32, softmax(logits)[:, 1], passages scored shortest first. TF32 is off
(L4 / A10G would otherwise round matmuls and move scores by ~1e-3).
Input: {"inputs": {"query": "Who wrote Guru Ramana?", "passages": ["...", "..."]}}
Output: {"scores": [0.93, 0.12, ...]} P(relevant) in [0, 1], in the order of passages
A warm-up is an ordinary call with one short pair.
Env (endpoint settings, optional): RERANK_BATCH (pairs per forward pass, default 32),
RERANK_MAX_PASSAGES (largest request accepted, default 128).
"""
import os
import torch
from transformers import BertConfig, BertForSequenceClassification, BertTokenizerFast
torch.backends.cuda.matmul.allow_tf32 = False # exact fp32, as on the backend's CPU
torch.backends.cudnn.allow_tf32 = False
BATCH = int(os.environ.get("RERANK_BATCH", "32"))
MAX_PASSAGES = int(os.environ.get("RERANK_MAX_PASSAGES", "128"))
class EndpointHandler:
def __init__(self, path=""):
self.device = "cuda" if torch.cuda.is_available() else "cpu"
self.tokenizer = BertTokenizerFast.from_pretrained(path)
self.model = BertForSequenceClassification.from_pretrained(
path, config=BertConfig.from_pretrained(path, num_labels=2), # config.json leaves num_labels unset
).to(self.device).eval()
def __call__(self, data: dict) -> dict:
inputs = data.get("inputs", data)
if not isinstance(inputs, dict):
raise ValueError('expected {"inputs": {"query": str, "passages": [str, ...]}}')
query, passages = inputs.get("query"), inputs.get("passages")
if not isinstance(query, str) or not query.strip():
raise ValueError("query must be a non-empty string")
if not isinstance(passages, list) or not all(isinstance(p, str) for p in passages):
raise ValueError("passages must be a list of strings")
if len(passages) > MAX_PASSAGES:
raise ValueError(f"at most {MAX_PASSAGES} passages per request")
# Shortest first: a batch is padded to its longest pair, so similar lengths waste least work.
order = sorted(range(len(passages)), key=lambda i: len(passages[i]))
scores = [0.0] * len(passages)
with torch.inference_mode():
for s in range(0, len(order), BATCH):
idx = order[s:s + BATCH]
enc = self.tokenizer(
[[query, passages[i]] for i in idx],
padding=True, truncation="only_second", max_length=512, return_tensors="pt",
).to(self.device)
probs = self.model(**enc).logits.float().softmax(-1)[:, 1].tolist()
for i, p in zip(idx, probs):
scores[i] = p
return {"scores": scores}
|