ameliechatelain's picture
Model card: add exact reranking prompts, mark flash-attn optional
c25e8b6 verified
|
Raw
History Blame
12.8 kB
---
license: apache-2.0
language:
- en
- fr
base_model:
- Qwen/Qwen3.5-4B
pipeline_tag: image-text-to-text
library_name: transformers
tags:
- reranker
- cross-encoder
- multimodal
- text-ranking
- document-reranking
- vidore
- beir
---
<div align="center">
<img src="https://cdn-uploads.huggingface.co/production/uploads/609bbe2f4932693ca2009d6a/kbQOAarw0eaApow3M9HIl.png" alt="LightOn" width="512">
[![Website](https://img.shields.io/badge/LightOn-Website-blue?logo=google-chrome)](https://lighton.ai)
[![LinkedIn](https://img.shields.io/badge/LightOn-LinkedIn-0A66C2?logo=linkedin)](https://www.linkedin.com/company/lighton/)
[![X](https://img.shields.io/badge/@LightOnIO-X-black?logo=x)](https://x.com/LightOnIO)
📝 Blog post: *coming soon* <!-- TODO: add blog post link -->
</div>
<h1 align="center">LightOn-rerank-LW-4B</h1>
<h3 align="center">Unified Text + Visual Document Reranker by LightOn</h3>
<p align="center">
<a href="https://huggingface.co/lightonai/LightOn-rerank-PW-0.8B">PW-0.8B</a> |
<a href="https://huggingface.co/lightonai/LightOn-rerank-LW-0.8B">LW-0.8B</a> |
<a href="https://huggingface.co/lightonai/LightOn-rerank-PW-2B">PW-2B</a> |
<a href="https://huggingface.co/lightonai/LightOn-rerank-LW-2B">LW-2B</a> |
<a href="https://huggingface.co/lightonai/LightOn-rerank-PW-4B">PW-4B</a> |
<a href="https://huggingface.co/lightonai/LightOn-rerank-LW-4B">LW-4B</a>
</p>
---
## About the LightOn-rerank family
Production retrieval pipelines usually need two rerankers: one for text passages and one for visual documents (PDF pages, slides, scans). **LightOn-rerank** models are *unified* cross-encoder rerankers: a single model scores both text passages and document page images against a query, on top of any first-stage retriever (BM25, dense embeddings, or ColPali-family late-interaction models).
The models are built on Qwen3.5 vision-language backbones (hybrid linear + full attention) and jointly fine-tuned on text and visual reranking data with mixed-modality batches (LoRA, merged into the released weights). Training data is English-only; French performance transfers zero-shot from the multilingual backbone.
The family comes in two scoring flavours × three sizes (0.8B / 2B / 4B):
- **PW (pointwise)** — each candidate is scored independently: the model judges whether the document answers the query, and the score is `logit("Yes") − logit("No")`. One forward pass per candidate, no generation — simple to serve (vLLM-compatible) and embarrassingly parallel.
- **LW (listwise)** — generative listwise ranking: 4 candidates are placed in a single prompt and the model generates a permutation (`[2] > [4] > [1] > [3]`). Larger candidate pools are ranked with a sliding window (window 4, stride 2, bottom-to-top). Cross-document attention makes LW markedly stronger on hard visual reranking, and unlike pointwise scoring it keeps improving with backbone size.
**LightOn-rerank-LW-4B** is the strongest model of the family: on ViDoRe V3 it outperforms the official Qwen3-VL-Reranker-8B (0.6469 vs 0.6423 NDCG@10) at half the parameter count, winning 14 of 16 domain×language splits against our 2B listwise model, with French gaining even more than English.
## Results
**ViDoRe V3** (visual document reranking, 8 domains × EN/FR queries), NDCG@10, ColQwen2.5-v0.2 first stage, retrieve 100 / rerank 100. All models — including baselines — were re-evaluated under this same two-stage protocol, so numbers are mutually comparable but not comparable to vendor-reported end-to-end results.
| Model | Params | Scoring | ViDoRe V3 overall @10 |
|---|---|---|---|
| **LightOn-rerank-LW-4B (this model)** | 4.5B | listwise | **0.6469** |
| *Qwen3-VL-Reranker-8B* | *8B* | *pointwise (pooling)* | *0.6423* |
| [LightOn-rerank-LW-2B](https://huggingface.co/lightonai/LightOn-rerank-LW-2B) | 2.2B | listwise | 0.6266 |
| [LightOn-rerank-PW-2B](https://huggingface.co/lightonai/LightOn-rerank-PW-2B) | 2.2B | pointwise | 0.5987 |
| [LightOn-rerank-PW-4B](https://huggingface.co/lightonai/LightOn-rerank-PW-4B) | 4.5B | pointwise | 0.5980 |
| *jina-reranker-m0* | *2.4B* | *pointwise* | *0.5939* |
| *Qwen3-VL-Reranker-2B* | *2B* | *pointwise (pooling)* | *0.5918* |
| [LightOn-rerank-LW-0.8B](https://huggingface.co/lightonai/LightOn-rerank-LW-0.8B) | 0.85B | listwise | 0.5825 |
| [LightOn-rerank-PW-0.8B](https://huggingface.co/lightonai/LightOn-rerank-PW-0.8B) | 0.85B | pointwise | 0.4820 |
### ViDoRe V3 detail (NDCG, ColQwen2.5 first stage, rerank-100, sliding window 4/2)
| Domain | EN @5 | EN @10 | FR @5 | FR @10 |
|---|---|---|---|---|
| finance_en | 0.7329 | 0.7412 | 0.6526 | 0.6635 |
| finance_fr | 0.4702 | 0.4920 | 0.5023 | 0.5255 |
| computer_science | 0.8159 | 0.8284 | 0.7922 | 0.8087 |
| hr | 0.7033 | 0.7100 | 0.6610 | 0.6618 |
| energy | 0.6915 | 0.7115 | 0.7160 | 0.7326 |
| industrial | 0.5933 | 0.5995 | 0.5509 | 0.5495 |
| pharmaceuticals | 0.6726 | 0.6810 | 0.6502 | 0.6615 |
| physics | 0.4690 | 0.4918 | 0.4554 | 0.4920 |
| **mean** | **0.6436** | **0.6569** | **0.6226** | **0.6369** |
Overall @10: **0.6469** (EN 0.6569 / FR 0.6369) — best result under this protocol, above the official Qwen3-VL-Reranker-8B (0.6423) at half the parameters.
### Text reranking — BEIR (13 datasets, NDCG@10, BM25 first stage, rerank-100, sliding window 4/2)
Mean @10 **0.4808**; clean mean (11 datasets, excluding contaminated NQ/MSMARCO) **0.4833**. Strongest datasets: fever 0.7928, scifact 0.7660, hotpotqa 0.7304, trec-covid 0.7152. Text gains over the 2B listwise model are modest (+0.003 clean mean on the 8-dataset overlap) — the scale-up pays off mainly on visual reranking.
## Model Details
- **Model type:** multimodal cross-encoder reranker — **generative listwise**: 4 candidates per prompt, ranked by generating a permutation; sliding window (4, stride 2) for larger pools
- **Base model:** [Qwen/Qwen3.5-4B](https://huggingface.co/Qwen/Qwen3.5-4B) (Qwen3.5 hybrid linear + full attention VLM)
- **Parameters:** ≈4.5B (bfloat16, 9.1 GB)
- **Inputs:** query (text) + candidate document(s) — text passage **or** page image
- **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 5e-5, warmup 30%, 1 epoch (419 steps)
- **Data:** gold ranking permutations built from labelled pos/hard-neg structure (no teacher): 107k text groups (NQ, TriviaQA, MSMARCO) + 106k visual query–page groups (ColPali-family hard negatives at ranks [0, 5, 10])
- **Languages:** English (training), French (zero-shot transfer)
- **Requirements:** `transformers >= 5.4.0` (`qwen3_5` architecture)
- **Internal experiment ID:** `exp39`
## Usage — generative listwise reranking
The model ranks **4 candidates per prompt** by generating a permutation string such as `[2] > [1] > [4] > [3]`. Candidate pools larger than 4 are ranked with a sliding window (window 4, stride 2) moving from the bottom of the list to the top, so the best candidates bubble up to the front. If a generation cannot be parsed, fall back to the input order (observed fallback rate in our evals: ≈0.01%).
```python
import re
import torch
from transformers import AutoModelForImageTextToText, AutoProcessor
model_id = "lightonai/LightOn-rerank-LW-4B"
model = AutoModelForImageTextToText.from_pretrained(
model_id,
dtype=torch.bfloat16,
attn_implementation="flash_attention_2", # optional — remove if flash-attn is not installed
device_map="cuda",
).eval()
processor = AutoProcessor.from_pretrained(model_id)
PROMPT = "<|im_start|>user\n{user}<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n"
PERM_RE = re.compile(r"\[(\d)\]\s*>\s*\[(\d)\]\s*>\s*\[(\d)\]\s*>\s*\[(\d)\]")
def rank_window(query: str, docs: list[str]) -> list[int]:
"""Rank exactly 4 text passages; returns window indices, most relevant first."""
body = "\n".join(f"[{i + 1}]: {d}" for i, d in enumerate(docs))
user = f"Query: {query}\n\nRank these passages from most to least relevant:\n{body}\n\nRanking:"
inputs = processor(text=[PROMPT.format(user=user)], return_tensors="pt").to(model.device)
out = model.generate(
**inputs, max_new_tokens=30, do_sample=False,
pad_token_id=processor.tokenizer.eos_token_id,
)
completion = processor.decode(out[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)
m = PERM_RE.search(completion)
return [int(g) - 1 for g in m.groups()] if m else list(range(4))
def rerank(query: str, docs: list, window: int = 4, stride: int = 2) -> list[int]:
"""Sliding-window listwise rerank; returns document indices, most relevant first."""
order = list(range(len(docs)))
positions = list(range(max(0, len(docs) - window), -1, -stride))
if positions and positions[-1] != 0:
positions.append(0)
for pos in positions:
end = min(pos + window, len(docs))
p = max(0, end - window)
if end - p < 2:
continue
perm = rank_window(query, [docs[i] for i in order[p:end]])
order[p:end] = [order[p + j] for j in perm]
return order
query = "What is late interaction in neural information retrieval?"
documents = ["passage 1 ...", "passage 2 ...", "passage 3 ...", "passage 4 ...", "passage 5 ..."]
print(rerank(query, documents))
```
For **page images**, build the window prompt with interleaved image placeholders instead:
```python
def rank_window_images(query: str, images: list) -> list[int]: # 4 PIL images
content = [{"type": "text", "text": f"Query: {query}\n\nRank these documents from most to least relevant:\n"}]
for i, img in enumerate(images):
content += [
{"type": "text", "text": f"[{i + 1}]: "},
{"type": "image", "image": img},
{"type": "text", "text": "\n"},
]
content.append({"type": "text", "text": "\nRanking:"})
text = processor.apply_chat_template(
[{"role": "user", "content": content}], tokenize=False, add_generation_prompt=True
)
text += "<think>\n\n</think>\n\n" # REQUIRED on the 4B backbone
inputs = processor(text=[text], images=list(images), return_tensors="pt").to(model.device)
out = model.generate(
**inputs, max_new_tokens=30, do_sample=False,
pad_token_id=processor.tokenizer.eos_token_id,
)
completion = processor.decode(out[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)
m = PERM_RE.search(completion)
return [int(g) - 1 for g in m.groups()] if m else list(range(4))
```
If a window has fewer than 4 candidates, pad it by repeating the last candidate and drop the duplicates from the returned order. Our evaluations ran this model with HF `generate()`; at the time of our evals vLLM could not serve the Qwen3.5-4B hybrid architecture for generation.
## Notes & limitations
- **Prepend an empty thinking block to the assistant turn.** Qwen3.5-4B is a thinking model; without the `<think>\n\n</think>\n\n` prefix (already included in the snippets above), reasoning tokens consume the generation budget and the ranking permutation never appears.
- Training data is English-only. French works zero-shot (the backbone is multilingual) but is slightly behind English on average.
- Vision hard negatives were mined with ColPali-family retrievers; the model pairs best with a ColPali-family first stage (e.g. ColQwen2.5) for visual reranking.
- BEIR contamination flag: NQ and MSMARCO are part of the text training data; headline text figures use clean means that exclude them.
## The LightOn-rerank family
| Model | Backbone | Scoring | ViDoRe V3 @10 |
|---|---|---|---|
| [LightOn-rerank-PW-0.8B](https://huggingface.co/lightonai/LightOn-rerank-PW-0.8B) | Qwen3.5-0.8B | pointwise | 0.4820 |
| [LightOn-rerank-LW-0.8B](https://huggingface.co/lightonai/LightOn-rerank-LW-0.8B) | Qwen3.5-0.8B | listwise | 0.5825 |
| [LightOn-rerank-PW-2B](https://huggingface.co/lightonai/LightOn-rerank-PW-2B) | Qwen3.5-2B | pointwise | 0.5987 |
| [LightOn-rerank-LW-2B](https://huggingface.co/lightonai/LightOn-rerank-LW-2B) | Qwen3.5-2B | listwise | 0.6266 |
| [LightOn-rerank-PW-4B](https://huggingface.co/lightonai/LightOn-rerank-PW-4B) | Qwen3.5-4B | pointwise | 0.5980 |
| [LightOn-rerank-LW-4B](https://huggingface.co/lightonai/LightOn-rerank-LW-4B) | Qwen3.5-4B | listwise | **0.6469** |
Rule of thumb: **LW** models are stronger at every size (and the gap grows with size); **PW** models are cheaper to serve and score candidates independently. For the best quality pick [LW-4B](https://huggingface.co/lightonai/LightOn-rerank-LW-4B); for the best quality/cost trade-off pick [LW-2B](https://huggingface.co/lightonai/LightOn-rerank-LW-2B); for maximum throughput on text-heavy workloads pick a PW model.