| """ |
| 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 |
| 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), |
| ).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") |
|
|
| |
| 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} |
|
|
