Create handler.py

handler.py

@@ -0,0 +1,77 @@
+from typing import Any, Dict, List
+import torch
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
+
+
+class EndpointHandler:
+    def __init__(self, model_dir: str, **kwargs: Any) -> None:
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+        self.tokenizer = AutoTokenizer.from_pretrained(model_dir)
+        self.model = AutoModelForSequenceClassification.from_pretrained(model_dir)
+        self.model.to(self.device)
+        self.model.eval()
+
+        # Fix batching when the tokenizer has no pad token
+        if self.tokenizer.pad_token is None:
+            if self.tokenizer.eos_token is not None:
+                self.tokenizer.pad_token = self.tokenizer.eos_token
+            else:
+                self.tokenizer.add_special_tokens({"pad_token": "[PAD]"})
+                self.model.resize_token_embeddings(len(self.tokenizer))
+
+        self.model.config.pad_token_id = self.tokenizer.pad_token_id
+
+    @torch.inference_mode()
+    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
+        """
+        Expected request body:
+        {
+          "query": "What is the capital of China?",
+          "documents": [
+            "The capital of China is Beijing.",
+            "Gravity is a force..."
+          ],
+          "instruction": "Given a web search query, retrieve relevant passages that answer the query"
+        }
+        """
+
+        query = data["query"]
+        documents = data["documents"]
+        instruction = data.get("instruction")
+
+        if instruction:
+            query_text = f"Instruct: {instruction}\nQuery: {query}"
+        else:
+            query_text = query
+
+        pairs = [[query_text, doc] for doc in documents]
+
+        inputs = self.tokenizer(
+            pairs,
+            padding=True,
+            truncation=True,
+            return_tensors="pt"
+        ).to(self.device)
+
+        outputs = self.model(**inputs)
+
+        # Common reranker head shape: [batch, 1] or [batch, 2]
+        logits = outputs.logits
+
+        if logits.shape[-1] == 1:
+            scores = logits.squeeze(-1).float().cpu().tolist()
+        else:
+            # If binary classification style, use the positive class
+            scores = logits[:, -1].float().cpu().tolist()
+
+        ranked = sorted(
+            [{"index": i, "score": s, "document": documents[i]} for i, s in enumerate(scores)],
+            key=lambda x: x["score"],
+            reverse=True,
+        )
+
+        return {
+            "scores": scores,
+            "ranked": ranked,
+        }