Add logging statements throughout model lifecycle

- Add logger initialization using __name__
- Log model loading process including device and config type
- Log CrossEncoder initialization stages
- Log prediction batching and processing
- Replace print statements with logger for OOM handling
- Log device changes in to_device function

modeling_zeranker.py

@@ -1,6 +1,7 @@
 from sentence_transformers import CrossEncoder as _CE
 
 import math
+import logging
 from typing import cast, Any
 import types
 
@@ -23,6 +24,8 @@ from transformers.tokenization_utils_fast import PreTrainedTokenizerFast
 # pyright: reportUnknownMemberType=false
 # pyright: reportUnknownVariableType=false
 
+logger = logging.getLogger(__name__)
+
 MODEL_PATH = "zeroentropy/zerank-2"
 PER_DEVICE_BATCH_SIZE_TOKENS = 10_000
 global_device = (
@@ -74,9 +77,12 @@ def load_model(
     if device is None:
         device = global_device
 
+    logger.info(f"Loading model from {MODEL_PATH} on device: {device}")
+
     config = AutoConfig.from_pretrained(MODEL_PATH)
     assert isinstance(config, PretrainedConfig)
 
+    logger.info(f"Loading model with config type: {config.model_type}")
     model = AutoModelForCausalLM.from_pretrained(
         MODEL_PATH,
         torch_dtype="auto",
@@ -93,6 +99,7 @@ def load_model(
         | Qwen3ForCausalLM,
     )
 
+    logger.info("Loading tokenizer")
     tokenizer = cast(
         AutoTokenizer,
         AutoTokenizer.from_pretrained(
@@ -105,6 +112,7 @@ def load_model(
     if tokenizer.pad_token is None:
         tokenizer.pad_token = tokenizer.eos_token
 
+    logger.info("Model and tokenizer loaded successfully")
     return tokenizer, model
 
 
@@ -113,16 +121,19 @@ _original_init = _CE.__init__
 
 
 def __init__(self, *args: Any, **kwargs: Any) -> None:
+    logger.info("Initializing CrossEncoder with eager model loading")
     # Call the original CrossEncoder __init__ first
     _original_init(self, *args, **kwargs)
 
     # Load the model immediately on instantiation
+    logger.info("Loading model on instantiation (no lazy loading)")
     self.inner_tokenizer, self.inner_model = load_model(global_device)
     self.inner_model.eval()
     self.inner_model.gradient_checkpointing_disable()
     self.inner_yes_token_id = self.inner_tokenizer.encode(
         "Yes", add_special_tokens=False
     )[0]
+    logger.info(f"CrossEncoder initialization complete. Yes token ID: {self.inner_yes_token_id}")
 
 
 def predict(
@@ -142,6 +153,8 @@ def predict(
             raise ValueError("query_documents or sentences must be provided")
         query_documents = [[sentence[0], sentence[1]] for sentence in sentences]
 
+    logger.info(f"Starting prediction for {len(query_documents)} query-document pairs")
+
     model = self.inner_model
     tokenizer = self.inner_tokenizer
 
@@ -170,9 +183,12 @@ def predict(
         batches[-1].append((query, document))
         max_length = max(max_length, 20 + len(query) + len(document))
 
+    logger.info(f"Created {len(batches)} batches for inference")
+
     # Inference all of the document batches
     all_logits: list[float] = []
-    for batch in batches:
+    for batch_idx, batch in enumerate(batches):
+        logger.debug(f"Processing batch {batch_idx + 1}/{len(batches)} with {len(batch)} pairs")
         batch_inputs = format_pointwise_datapoints(
             tokenizer,
             batch,
@@ -184,9 +200,9 @@ def predict(
             with torch.inference_mode():
                 outputs = model(**batch_inputs, use_cache=False)
         except torch.OutOfMemoryError:
-            print(f"GPU OOM! {torch.cuda.memory_reserved()}")
+            logger.warning(f"GPU OOM! Memory reserved: {torch.cuda.memory_reserved()}")
             torch.cuda.empty_cache()
-            print(f"GPU After OOM Cache Clear: {torch.cuda.memory_reserved()}")
+            logger.info(f"GPU cache cleared. Memory reserved: {torch.cuda.memory_reserved()}")
             outputs = model(**batch_inputs, use_cache=False)
 
         # Extract the logits
@@ -209,11 +225,13 @@ def predict(
     # Unsort by indices
     scores = [score for _, score in sorted(zip(permutation, scores, strict=True))]
 
+    logger.info(f"Prediction complete. Generated {len(scores)} scores")
     return scores
 
 
 def to_device(self: _CE, new_device: torch.device) -> None:
     global global_device
+    logger.info(f"Changing device from {global_device} to {new_device}")
     global_device = new_device