Update handler.py

handler.py

@@ -1,59 +1,115 @@
+import os
 from typing import Any, Dict, List
+
 import torch
-from transformers import AutoTokenizer, AutoModelForSequenceClassification
+from transformers import AutoTokenizer, AutoModelForCausalLM
+
 
 class EndpointHandler:
     def __init__(self, path: str = ""):
+        model_dir = path if path else os.getenv("HF_MODEL_DIR", ".")
         self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.dtype = torch.float16 if self.device == "cuda" else torch.float32
+
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            model_dir,
+            padding_side="left",
+            trust_remote_code=True,
+        )
+        self.model = AutoModelForCausalLM.from_pretrained(
+            model_dir,
+            torch_dtype=self.dtype,
+            trust_remote_code=True,
+        ).to(self.device).eval()
+
+        self.token_false_id = self.tokenizer.convert_tokens_to_ids("no")
+        self.token_true_id = self.tokenizer.convert_tokens_to_ids("yes")
+
+        self.max_length = 8192
+
+        self.prefix = (
+            '<|im_start|>system\n'
+            'Judge whether the Document meets the requirements based on the Query '
+            'and the Instruct provided. Note that the answer can only be "yes" or "no".'
+            '<|im_end|>\n'
+            '<|im_start|>user\n'
+        )
+        self.suffix = '<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n'
+
+        self.prefix_tokens = self.tokenizer.encode(self.prefix, add_special_tokens=False)
+        self.suffix_tokens = self.tokenizer.encode(self.suffix, add_special_tokens=False)
+
+    def _format_one(self, instruction: str, query: str, document: str) -> str:
+        return f"<Instruct>: {instruction}\n<Query>: {query}\n<Document>: {document}"
+
+    def _process_inputs(self, pairs: List[str]) -> Dict[str, torch.Tensor]:
+        inputs = self.tokenizer(
+            pairs,
+            padding=False,
+            truncation="longest_first",
+            return_attention_mask=False,
+            max_length=self.max_length - len(self.prefix_tokens) - len(self.suffix_tokens),
+        )
+
+        for i, ids in enumerate(inputs["input_ids"]):
+            inputs["input_ids"][i] = self.prefix_tokens + ids + self.suffix_tokens
+
+        padded = self.tokenizer.pad(
+            inputs,
+            padding=True,
+            return_tensors="pt",
+            max_length=self.max_length,
+        )
 
-        self.tokenizer = AutoTokenizer.from_pretrained(path)
-        self.model = AutoModelForSequenceClassification.from_pretrained(path).to(self.device)
-        self.model.eval()
+        for k in padded:
+            padded[k] = padded[k].to(self.device)
 
-        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))
+        return padded
 
-        self.model.config.pad_token_id = self.tokenizer.pad_token_id
+    @torch.no_grad()
+    def _score(self, model_inputs: Dict[str, torch.Tensor]) -> List[float]:
+        logits = self.model(**model_inputs).logits[:, -1, :]
+        false_scores = logits[:, self.token_false_id]
+        true_scores = logits[:, self.token_true_id]
+        pair_scores = torch.stack([false_scores, true_scores], dim=1)
+        probs = torch.nn.functional.softmax(pair_scores, dim=1)[:, 1]
+        return probs.tolist()
 
-    @torch.inference_mode()
     def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
-        payload = data["inputs"]
-        params = data.get("parameters", {})
+        payload = data.get("inputs", data)
 
+        instruction = payload.get(
+            "instruction",
+            "Given a web search query, retrieve relevant passages that answer the query",
+        )
         query = payload["query"]
         documents = payload["documents"]
-        instruction = params.get("instruction")
+        return_documents = payload.get("return_documents", True)
 
-        if instruction:
-            query_text = f"Instruct: {instruction}\nQuery: {query}"
-        else:
-            query_text = query
+        if not isinstance(documents, list) or len(documents) == 0:
+            raise ValueError("`documents` must be a non-empty list of strings.")
 
-        pairs = [[query_text, doc] for doc in documents]
+        pairs = [self._format_one(instruction, query, doc) for doc in documents]
+        model_inputs = self._process_inputs(pairs)
+        scores = self._score(model_inputs)
 
-        enc = self.tokenizer(
-            pairs,
-            padding=True,
-            truncation=True,
-            return_tensors="pt"
-        ).to(self.device)
-
-        outputs = self.model(**enc)
-        logits = outputs.logits
-
-        if logits.shape[-1] == 1:
-            scores = logits.squeeze(-1).float().cpu().tolist()
-        else:
-            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,
-        )
+        results = []
+        for i, (doc, score) in enumerate(zip(documents, scores)):
+            item = {
+                "index": i,
+                "relevance_score": float(score),
+            }
+            if return_documents:
+                item["document"] = doc
+            results.append(item)
+
+        results.sort(key=lambda x: x["relevance_score"], reverse=True)
 
-        return {"scores": scores, "ranked": ranked}
+        return {
+            "results": results,
+            "meta": {
+                "instruction": instruction,
+                "query": query,
+                "count": len(documents),
+            },
+        }