add: max_query/doc_length parametrization

modeling.py

@@ -1,10 +1,11 @@
-import numpy as np
 from dataclasses import dataclass
+from typing import Dict, List, Optional, Tuple
+
+import numpy as np
 import torch
 from torch import nn
-from typing import Optional, List, Dict, Tuple
-from transformers.models.qwen3 import modeling_qwen3
 from transformers.modeling_outputs import CausalLMOutputWithPast
+from transformers.models.qwen3 import modeling_qwen3
 
 
 @dataclass
@@ -51,9 +52,12 @@ def format_docs_prompts_func(
     )
 
     if instruction:
-        prompt += f'<instruct>\n{instruction}\n</instruct>\n'
+        prompt += f"<instruct>\n{instruction}\n</instruct>\n"
 
-    doc_prompts = [f'<passage id="{i}">\n{doc}{doc_emb_token}\n</passage>' for i, doc in enumerate(docs)]
+    doc_prompts = [
+        f'<passage id="{i}">\n{doc}{doc_emb_token}\n</passage>'
+        for i, doc in enumerate(docs)
+    ]
     prompt += "\n".join(doc_prompts) + "\n"
     prompt += f"<query>\n{query}{query_emb_token}\n</query>"
 
@@ -76,14 +80,19 @@ class JinaForRanking(modeling_qwen3.Qwen3ForCausalLM):
 
         self.post_init()
 
-        self.special_tokens = {"query_embed_token": "<|rerank_token|>", "doc_embed_token": "<|embed_token|>"}
+        self.special_tokens = {
+            "query_embed_token": "<|rerank_token|>",
+            "doc_embed_token": "<|embed_token|>",
+        }
         self.doc_embed_token_id = 151670
         self.query_embed_token_id = 151671
 
     def forward(self, *args, **kwargs) -> CausalLMOutputWithScores:
         kwargs.pop("output_hidden_states", None)
         kwargs.pop("use_cache", None)
-        assert kwargs.pop("labels", None) is None, "labels should not be passed to forward()"
+        assert kwargs.pop("labels", None) is None, (
+            "labels should not be passed to forward()"
+        )
         input_ids = kwargs.pop("input_ids", None)
 
         outputs = super().forward(
@@ -107,7 +116,9 @@ class JinaForRanking(modeling_qwen3.Qwen3ForCausalLM):
         query_embeds = self.projector(query_embeds)
 
         query_embeds_expanded = query_embeds.expand_as(doc_embeds)
-        scores = torch.nn.functional.cosine_similarity(doc_embeds, query_embeds_expanded, dim=-1).squeeze(-1)
+        scores = torch.nn.functional.cosine_similarity(
+            doc_embeds, query_embeds_expanded, dim=-1
+        ).squeeze(-1)
 
         return CausalLMOutputWithScores(
             loss=None,
@@ -124,13 +135,17 @@ class JinaForRanking(modeling_qwen3.Qwen3ForCausalLM):
         if not hasattr(self, "_tokenizer"):
             from transformers import AutoTokenizer
 
-            self._tokenizer = AutoTokenizer.from_pretrained(self.name_or_path, trust_remote_code=True)
+            self._tokenizer = AutoTokenizer.from_pretrained(
+                self.name_or_path, trust_remote_code=True
+            )
 
             if self._tokenizer.pad_token is None:
                 self._tokenizer.pad_token = self._tokenizer.unk_token
-                self._tokenizer.pad_token_id = self._tokenizer.convert_tokens_to_ids(self._tokenizer.pad_token)
+                self._tokenizer.pad_token_id = self._tokenizer.convert_tokens_to_ids(
+                    self._tokenizer.pad_token
+                )
 
-            self._tokenizer.padding_side = 'left'
+            self._tokenizer.padding_side = "left"
 
     def _truncate_texts(
         self,
@@ -144,17 +159,21 @@ class JinaForRanking(modeling_qwen3.Qwen3ForCausalLM):
         docs = []
         doc_lengths = []
         for doc in documents:
-            doc_tokens = self._tokenizer(doc, truncation=True, max_length=max_doc_length)
-            if len(doc_tokens['input_ids']) >= max_doc_length:
-                doc = self._tokenizer.decode(doc_tokens['input_ids'])
-            doc_lengths.append(len(doc_tokens['input_ids']))
+            doc_tokens = self._tokenizer(
+                doc, truncation=True, max_length=max_doc_length
+            )
+            if len(doc_tokens["input_ids"]) >= max_doc_length:
+                doc = self._tokenizer.decode(doc_tokens["input_ids"])
+            doc_lengths.append(len(doc_tokens["input_ids"]))
             docs.append(doc)
 
-        query_tokens = self._tokenizer(query, truncation=True, max_length=max_query_length)
-        if len(query_tokens['input_ids']) >= max_query_length:
-            query = self._tokenizer.decode(query_tokens['input_ids'])
+        query_tokens = self._tokenizer(
+            query, truncation=True, max_length=max_query_length
+        )
+        if len(query_tokens["input_ids"]) >= max_query_length:
+            query = self._tokenizer.decode(query_tokens["input_ids"])
 
-        query_length = len(query_tokens['input_ids'])
+        query_length = len(query_tokens["input_ids"])
 
         return query, docs, doc_lengths, query_length
 
@@ -200,6 +219,8 @@ class JinaForRanking(modeling_qwen3.Qwen3ForCausalLM):
         documents: List[str],
         top_n: Optional[int] = None,
         return_embeddings: bool = False,
+        max_doc_length: int = 2048,
+        max_query_length: int = 512,
     ) -> List[dict]:
         """
         Rerank documents by relevance to a query.
@@ -221,14 +242,14 @@ class JinaForRanking(modeling_qwen3.Qwen3ForCausalLM):
 
         # Derived from model configuration
         max_length = self._tokenizer.model_max_length
-        max_query_length = 512
-        max_doc_length = 2048
 
         # Derive block_size from max_length to fit documents efficiently
         # Heuristic: allow ~125 docs per batch for typical doc sizes
         block_size = 125
 
-        query, docs, doc_lengths, query_length = self._truncate_texts(query, documents, max_query_length, max_doc_length)
+        query, docs, doc_lengths, query_length = self._truncate_texts(
+            query, documents, max_query_length, max_doc_length
+        )
 
         length_capacity = max_length - 2 * query_length
 
@@ -242,7 +263,9 @@ class JinaForRanking(modeling_qwen3.Qwen3ForCausalLM):
             length_capacity -= length
 
             if len(block_docs) >= block_size or length_capacity <= max_doc_length:
-                outputs = self._compute_single_batch(query, block_docs, instruction=None)
+                outputs = self._compute_single_batch(
+                    query, block_docs, instruction=None
+                )
 
                 doc_embeddings.extend(outputs.doc_embeds[0].cpu().float().numpy())
                 query_embeddings.append(outputs.query_embeds[0].cpu().float().numpy())
@@ -277,10 +300,12 @@ class JinaForRanking(modeling_qwen3.Qwen3ForCausalLM):
 
         return [
             {
-                'document': documents[scores_argsort[i]],
-                'relevance_score': scores[0][scores_argsort[i]],
-                'index': scores_argsort[i],
-                'embedding': doc_embeddings[scores_argsort[i]] if return_embeddings else None,
+                "document": documents[scores_argsort[i]],
+                "relevance_score": scores[0][scores_argsort[i]],
+                "index": scores_argsort[i],
+                "embedding": doc_embeddings[scores_argsort[i]]
+                if return_embeddings
+                else None,
             }
             for i in range(top_n)
         ]