Integrate with Sentence Transformers (#1)

- Integrate with Sentence Transformers (95c9c83d21761bf279784ed6f64f404422aa4940)

1_LogitScore/config.json

@@ -0,0 +1,5 @@
+{
+    "true_token_id": 9175,
+    "false_token_id": 2665,
+    "module_input_name": "causal_logits"
+}

README.md

@@ -15,6 +15,7 @@ tags:
 - document-reranking
 - vidore
 - beir
+- sentence-transformers
 datasets:
 - vidore/colpali_train_set
 - lightonai/embeddings-fine-tuning
@@ -126,12 +127,60 @@ This is the best text result in the LightOn-rerank family (2B pointwise: 49.13,
 - **Fine-tuning:** joint text+vision LoRA (r=32, α=32, rsLoRA, merged into the released weights), mixed-modality batches (2 text + 2 vision groups per micro-batch), vision loss weight 1.3, lr 1e-4, 1 epoch (465 steps)
 - **Data:** same 213k groups as the listwise models — 107k text groups (NQ, TriviaQA, MS MARCO; each a `[pos, neg_0, neg_1, neg_2]` 4-list with hard negatives mined via the NV-Retriever approach with GTE-ModernBERT) + 106k vision groups (ColPali train set with negatives mined by Nomic). For pointwise training each group is flattened into (query, document, Yes/No) triples.
 - **Languages:** English (training), French (zero-shot transfer)
-- **Requirements:** `transformers >= 5.4.0` (`qwen3_5` architecture)
+- **Requirements:** `transformers >= 5.4.0` (`qwen3_5` architecture); `sentence-transformers >= 5.4.0` for the CrossEncoder usage
 
 ## Usage: pointwise reranking
 
 Each candidate is scored independently as `logit("Yes") − logit("No")` at the first generated position; sort candidates by descending score. The model was trained with a fixed system prompt and user template: use them verbatim for best results.
 
+### Using Sentence Transformers
+
+Install Sentence Transformers (`>= 5.4.0`):
+
+```bash
+pip install "sentence-transformers[image]"
+```
+
+The trained system prompt and user templates are baked into the bundled `reranker` chat template — including the non-thinking generation prompt, so no `enable_thinking` handling is needed — and query-document pairs are formatted correctly out of the box:
+
+```python
+from sentence_transformers import CrossEncoder
+
+model = CrossEncoder("lightonai/LightOn-rerank-PW-4B")
+
+query = "What is late interaction in neural information retrieval?"
+documents = [
+    "ColBERT computes token-level query-document interactions at search time...",
+    "The Eiffel Tower is located on the Champ de Mars in Paris.",
+]
+
+pairs = [(query, doc) for doc in documents]
+scores = model.predict(pairs)
+print(scores)
+# [-3.     -8.0625]
+
+rankings = model.rank(query, documents)
+print(rankings)
+# [{'corpus_id': 0, 'score': -3.0}, {'corpus_id': 1, 'score': -8.0625}]
+```
+
+To rerank **page images**, pass a `PIL.Image` (or an image URL or file path string) as the document. Text and image candidates can be mixed in the same call:
+
+```python
+from PIL import Image
+
+pairs = [
+    (query, Image.open("page_1.png")),
+    (query, "https://example.com/page_2.png"),
+    (query, "A text passage candidate for the same query."),
+]
+scores = model.predict(pairs)
+```
+
+Scores are raw `logit("Yes") − logit("No")` differences. You can map them to 0...1 probabilities with `model.predict(pairs, activation_fn=torch.nn.Sigmoid())`.
+
+### Using Transformers
+
 ```python
 import torch
 from transformers import AutoModelForImageTextToText, AutoProcessor
@@ -179,6 +228,7 @@ inputs = processor(
 with torch.inference_mode():
     logits = model(**inputs).logits[:, -1]
 scores = (logits[:, YES_TOKEN_ID] - logits[:, NO_TOKEN_ID]).tolist()
+# [-3.0, -8.0625]
 
 ranked = sorted(zip(scores, documents), reverse=True)
 ```
@@ -222,7 +272,7 @@ Full-page document images can exceed 8k tokens, so keep `--max-model-len` at 163
 
 ## Notes & limitations
 
-- **Thinking must be disabled for scoring.** Qwen3.5-4B's chat template enables `<think>` by default; with it on, the first generated token is a thinking token and the Yes/No logprobs are distorted. Pass `enable_thinking=False` to `apply_chat_template` (or `chat_template_kwargs={"enable_thinking": False}` via the vLLM OpenAI client) as shown above.
+- **Thinking must be disabled for scoring.** Qwen3.5-4B's chat template enables `<think>` by default; with it on, the first generated token is a thinking token and the Yes/No logprobs are distorted. Pass `enable_thinking=False` to `apply_chat_template` (or `chat_template_kwargs={"enable_thinking": False}` via the vLLM OpenAI client) as shown above. The Sentence Transformers path handles this automatically: the bundled `reranker` chat template hardcodes the non-thinking generation prompt.
 - Pointwise scoring does not benefit from the 2B→4B scale-up on vision (59.80 vs 59.87 for the 2B), while generative listwise gains +2.0 points over the same step. For vision workloads prefer [LightOn-rerank-LW-4B](https://huggingface.co/lightonai/LightOn-rerank-LW-4B) at this size, or [LightOn-rerank-PW-2B](https://huggingface.co/lightonai/LightOn-rerank-PW-2B) for the same visual quality at lower cost.
 - Training data is English-only. French works zero-shot (the backbone is multilingual) but is slightly behind English on average.
 - BEIR contamination flag: NQ and MSMARCO are part of the text training data; headline text figures use decontaminated means that exclude them.

additional_chat_templates/reranker.jinja

@@ -0,0 +1,99 @@
+{%- set default_system = "Judge whether the document is relevant to the query. Answer Yes or No." -%}
+{%- set ns = namespace(system_text="", doc_visual=false, doc_text=false) -%}
+{%- for message in messages if message.role == "system" -%}
+    {%- if message.content is string -%}
+        {%- set ns.system_text = ns.system_text + message.content -%}
+    {%- else -%}
+        {%- for content in message.content -%}
+            {%- if 'text' in content -%}
+                {%- set ns.system_text = ns.system_text + content.text -%}
+            {%- endif -%}
+        {%- endfor -%}
+    {%- endif -%}
+{%- endfor -%}
+{%- if not ns.system_text -%}
+    {%- set ns.system_text = default_system -%}
+{%- endif -%}
+{%- for message in messages if message.role == "document" -%}
+    {%- if message.content is string -%}
+        {%- if message.content -%}
+            {%- set ns.doc_text = true -%}
+        {%- endif -%}
+    {%- else -%}
+        {%- for content in message.content -%}
+            {%- if content.type == 'image' or 'image' in content or 'image_url' in content or content.type == 'video' or 'video' in content -%}
+                {%- set ns.doc_visual = true -%}
+            {%- elif 'text' in content and content.text -%}
+                {%- set ns.doc_text = true -%}
+            {%- endif -%}
+        {%- endfor -%}
+    {%- endif -%}
+{%- endfor -%}
+{%- macro render_content(message) -%}
+    {%- if message.content is string -%}
+        {{- message.content -}}
+    {%- else -%}
+        {%- for content in message.content -%}
+            {%- if content.type == 'image' or 'image' in content or 'image_url' in content -%}
+                {{- '<|vision_start|><|image_pad|><|vision_end|>' -}}
+            {%- elif content.type == 'video' or 'video' in content -%}
+                {{- '<|vision_start|><|video_pad|><|vision_end|>' -}}
+            {%- elif 'text' in content -%}
+                {{- content.text -}}
+            {%- endif -%}
+        {%- endfor -%}
+    {%- endif -%}
+{%- endmacro -%}
+{%- macro render_visuals(message) -%}
+    {%- if message.content is not string -%}
+        {%- for content in message.content -%}
+            {%- if content.type == 'image' or 'image' in content or 'image_url' in content -%}
+                {{- '<|vision_start|><|image_pad|><|vision_end|>' -}}
+            {%- elif content.type == 'video' or 'video' in content -%}
+                {{- '<|vision_start|><|video_pad|><|vision_end|>' -}}
+            {%- endif -%}
+        {%- endfor -%}
+    {%- endif -%}
+{%- endmacro -%}
+{%- macro render_texts(message) -%}
+    {%- if message.content is string -%}
+        {{- message.content -}}
+    {%- else -%}
+        {%- for content in message.content -%}
+            {%- if content.type != 'image' and 'image' not in content and 'image_url' not in content and content.type != 'video' and 'video' not in content and 'text' in content -%}
+                {{- content.text -}}
+            {%- endif -%}
+        {%- endfor -%}
+    {%- endif -%}
+{%- endmacro -%}
+{%- set user_body -%}
+    {%- if ns.doc_visual -%}
+        {%- for message in messages if message.role == "document" -%}
+            {{- render_visuals(message) -}}
+        {%- endfor -%}
+        {{- 'Given a query, determine if the document image is relevant. The query is: ' -}}
+        {%- for message in messages if message.role == "query" -%}
+            {{- render_content(message) -}}
+        {%- endfor -%}
+        {%- if ns.doc_text -%}
+            {{- '\n\nDocument: ' -}}
+            {%- for message in messages if message.role == "document" -%}
+                {{- render_texts(message) -}}
+            {%- endfor -%}
+        {%- endif -%}
+    {%- else -%}
+        {{- 'Given a query, determine if the document is relevant. The query is: ' -}}
+        {%- for message in messages if message.role == "query" -%}
+            {{- render_content(message) -}}
+        {%- endfor -%}
+        {{- '\n\nDocument: ' -}}
+        {%- for message in messages if message.role == "document" -%}
+            {{- render_content(message) -}}
+        {%- endfor -%}
+    {%- endif -%}
+{%- endset -%}
+{{- '<|im_start|>system\n' + ns.system_text|trim + '<|im_end|>\n' -}}
+{{- '<|im_start|>user\n' + user_body|trim + '<|im_end|>\n' -}}
+{%- if add_generation_prompt -%}
+    {{- '<|im_start|>assistant\n<think>\n\n</think>\n\n' -}}
+{%- endif -%}

config_sentence_transformers.json

@@ -0,0 +1,11 @@
+{
+  "__version__": {
+    "pytorch": "2.10.0+cu128",
+    "sentence_transformers": "5.6.0",
+    "transformers": "5.13.1"
+  },
+  "activation_fn": "torch.nn.modules.linear.Identity",
+  "default_prompt_name": null,
+  "model_type": "CrossEncoder",
+  "prompts": {}
+}

modules.json

@@ -0,0 +1,14 @@
+[
+  {
+    "idx": 0,
+    "name": "0",
+    "path": "",
+    "type": "sentence_transformers.base.modules.transformer.Transformer"
+  },
+  {
+    "idx": 1,
+    "name": "1",
+    "path": "1_LogitScore",
+    "type": "sentence_transformers.cross_encoder.modules.logit_score.LogitScore"
+  }
+]

sentence_bert_config.json

@@ -0,0 +1,29 @@
+{
+    "transformer_task": "any-to-any",
+    "modality_config": {
+        "text": {
+            "method": "forward",
+            "method_output_name": "logits"
+        },
+        "image": {
+            "method": "forward",
+            "method_output_name": "logits"
+        },
+        "video": {
+            "method": "forward",
+            "method_output_name": "logits"
+        },
+        "message": {
+            "method": "forward",
+            "method_output_name": "logits",
+            "format": "structured"
+        }
+    },
+    "module_output_name": "causal_logits",
+    "processing_kwargs": {
+        "chat_template": {
+            "chat_template": "reranker",
+            "add_generation_prompt": true
+        }
+    }
+}