Update handler.py

handler.py

@@ -1,4 +1,6 @@
+import gc
 import os
+import threading
 from pathlib import Path
 from typing import Any, Dict, List
 
@@ -8,7 +10,7 @@ from transformers import AutoTokenizer, AutoModelForCausalLM
 
 class EndpointHandler:
     def __init__(self, path: str = ""):
-        model_dir = Path(path or os.getenv("HF_MODEL_DIR", ".")).resolve()
+        model_dir = Path(path or os.getenv("HF_MODEL_DIR", "")).resolve()
 
         if not model_dir.exists():
             raise FileNotFoundError(f"Model directory does not exist: {model_dir}")
@@ -49,7 +51,7 @@ class EndpointHandler:
 
         self.model = AutoModelForCausalLM.from_pretrained(
             str(model_dir),
-            torch_dtype=self.torch_dtype,
+            dtype=self.torch_dtype,
             trust_remote_code=True,
         ).to(self.device).eval()
 
@@ -65,7 +67,10 @@ class EndpointHandler:
         self.token_true_id = yes_ids[0]
         self.token_false_id = no_ids[0]
 
-        self.max_length = 8192
+        self.max_length = int(os.getenv("HANDLER_MAX_LENGTH", "8192"))
+        self.batch_size = int(os.getenv("HANDLER_BATCH_SIZE", "8"))
+        self.max_documents = int(os.getenv("HANDLER_MAX_DOCUMENTS", "64"))
+        self._semaphore = threading.Semaphore(int(os.getenv("HANDLER_MAX_CONCURRENT", "5")))
 
         self.prefix = (
             "<|im_start|>system\n"
@@ -114,6 +119,17 @@ class EndpointHandler:
         probs = torch.nn.functional.softmax(pair_scores, dim=1)[:, 1]
         return probs.tolist()
 
+    def _score_in_batches(self, pairs: List[str]) -> List[float]:
+        all_scores = []
+        for i in range(0, len(pairs), self.batch_size):
+            batch = pairs[i : i + self.batch_size]
+            model_inputs = self._process_inputs(batch)
+            scores = self._score(model_inputs)
+            all_scores.extend(scores)
+            del model_inputs
+            gc.collect()
+        return all_scores
+
     def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
         payload = data.get("inputs", data)
 
@@ -128,9 +144,27 @@ class EndpointHandler:
         if not isinstance(documents, list) or len(documents) == 0:
             raise ValueError("`documents` must be a non-empty list of strings.")
 
+        if len(documents) > self.max_documents:
+            raise ValueError(
+                f"`documents` exceeds max allowed ({self.max_documents}). "
+                f"Got {len(documents)}."
+            )
+
         pairs = [self._format_one(instruction, query, doc) for doc in documents]
-        model_inputs = self._process_inputs(pairs)
-        scores = self._score(model_inputs)
+        acquired = self._semaphore.acquire(timeout=int(os.getenv("HANDLER_QUEUE_TIMEOUT", "60")))
+        if not acquired:
+            raise RuntimeError(
+                "Server is busy. Another request is being processed. Please retry."
+            )
+        try:
+            scores = self._score_in_batches(pairs)
+        except MemoryError:
+            gc.collect()
+            raise RuntimeError(
+                "Out of memory while scoring. Try sending fewer or shorter documents."
+            )
+        finally:
+            self._semaphore.release()
 
         results = []
         for i, (doc, score) in enumerate(zip(documents, scores)):