Update handler.py

handler.py

@@ -2,17 +2,14 @@ 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")
+    def __init__(self, path: str = ""):
+        self.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.tokenizer = AutoTokenizer.from_pretrained(path)
+        self.model = AutoModelForSequenceClassification.from_pretrained(path).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
@@ -24,21 +21,12 @@ class EndpointHandler:
 
     @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"
-        }
-        """
+        payload = data["inputs"]
+        params = data.get("parameters", {})
 
-        query = data["query"]
-        documents = data["documents"]
-        instruction = data.get("instruction")
+        query = payload["query"]
+        documents = payload["documents"]
+        instruction = params.get("instruction")
 
         if instruction:
             query_text = f"Instruct: {instruction}\nQuery: {query}"
@@ -47,22 +35,19 @@ class EndpointHandler:
 
         pairs = [[query_text, doc] for doc in documents]
 
-        inputs = self.tokenizer(
+        enc = 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]
+        outputs = self.model(**enc)
         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(
@@ -71,7 +56,4 @@ class EndpointHandler:
             reverse=True,
         )
 
-        return {
-            "scores": scores,
-            "ranked": ranked,
-        }
+        return {"scores": scores, "ranked": ranked}