Update modeling_zeranker.py

modeling_zeranker.py

@@ -1,7 +1,7 @@
 from sentence_transformers import CrossEncoder as _CE
 
 import math
-from typing import cast
+from typing import cast, Any
 import types
 
 import torch
@@ -21,8 +21,11 @@ from transformers.tokenization_utils_fast import PreTrainedTokenizerFast
 # pyright: reportUnknownMemberType=false
 # pyright: reportUnknownVariableType=false
 
-MODEL_PATH = "zeroentropy/ze-rerank-small-v0.3.0"
+MODEL_PATH = "zeroentropy/zerank-1-small"
 PER_DEVICE_BATCH_SIZE_TOKENS = 15_000
+global_device = (
+    torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+)
 
 
 def format_pointwise_datapoints(
@@ -67,7 +70,7 @@ def load_model(
     | Qwen3ForCausalLM,
 ]:
     if device is None:
-        device = torch.device("cpu")
+        device = global_device
 
     config = AutoConfig.from_pretrained(MODEL_PATH)
     assert isinstance(config, PretrainedConfig)
@@ -80,7 +83,6 @@ def load_model(
     )
     if config.model_type == "llama":
         model.config.attn_implementation = "flash_attention_2"
-    print(f"Model Type: {config.model_type}")
     assert isinstance(
         model,
         LlamaForCausalLM
@@ -104,13 +106,30 @@ def load_model(
     return tokenizer, model
 
 
-def predict(self, query_documents: list[tuple[str, str]]) -> list[float]:
+def predict(
+    self,
+    query_documents: list[tuple[str, str]] | None = None,
+    *,
+    sentences: Any = None,
+    batch_size: Any = None,
+    show_progress_bar: Any = None,
+    activation_fn: Any = None,
+    apply_softmax: Any = None,
+    convert_to_numpy: Any = None,
+    convert_to_tensor: Any = None,
+) -> list[float]:
+    if query_documents is None:
+        if sentences is None:
+            raise ValueError("query_documents or sentences must be provided")
+        query_documents = [[sentence[0], sentence[1]] for sentence in sentences]
+
     if not hasattr(self, "inner_model"):
-        self.inner_tokenizer, self.inner_model = load_model(torch.device("cuda"))
+        self.inner_tokenizer, self.inner_model = load_model(global_device)
         self.inner_model.gradient_checkpointing_enable()
         self.inner_model.eval()
-        self.inner_yes_token_id = self.inner_tokenizer.encode("Yes", add_special_tokens=False)[0]
-        print("patched")
+        self.inner_yes_token_id = self.inner_tokenizer.encode(
+            "Yes", add_special_tokens=False
+        )[0]
 
     model = self.inner_model
     tokenizer = self.inner_tokenizer
@@ -120,11 +139,11 @@ def predict(self, query_documents: list[tuple[str, str]]) -> list[float]:
     ]
     # Sort
     permutation = list(range(len(query_documents)))
-    permutation.sort(key=lambda i: -len(query_documents[i][0]) - len(query_documents[i][1]))
+    permutation.sort(
+        key=lambda i: -len(query_documents[i][0]) - len(query_documents[i][1])
+    )
     query_documents = [query_documents[i] for i in permutation]
 
-    device = torch.device("cuda")
-
     # Extract document batches from this line of datapoints
     max_length = 0
     batches: list[list[tuple[str, str]]] = []
@@ -148,7 +167,7 @@ def predict(self, query_documents: list[tuple[str, str]]) -> list[float]:
             batch,
         )
 
-        batch_inputs = batch_inputs.to(device)
+        batch_inputs = batch_inputs.to(global_device)
 
         try:
             outputs = model(**batch_inputs, use_cache=False)
@@ -164,7 +183,7 @@ def predict(self, query_documents: list[tuple[str, str]]) -> list[float]:
         last_positions = attention_mask.sum(dim=1) - 1
 
         batch_size = logits.shape[0]
-        batch_indices = torch.arange(batch_size, device=device)
+        batch_indices = torch.arange(batch_size, device=global_device)
         last_logits = logits[batch_indices, last_positions]
 
         yes_logits = last_logits[:, self.inner_yes_token_id]
@@ -181,8 +200,15 @@ def predict(self, query_documents: list[tuple[str, str]]) -> list[float]:
     return scores
 
 
+def to_device(self: _CE, new_device: torch.device) -> None:
+    global global_device
+    global_device = new_device
+
+
 _CE.predict = predict
 
 from transformers import Qwen3Config
 
 ZEConfig = Qwen3Config
+
+_CE.to = to_device