| Benchmark | NDCG@10 |
|---|---|
| TREC DL19 | 73.9 |
| TREC DL20 | 72.1 |
| BEIR (Avg) | 44.8 |
| MS MARCO Dev | 68.5 |
1import torch
2from transformers import AutoTokenizer, AutoModelForSequenceClassification
3
4# Load model
5model_path = "abdoelsayed/dear-8b-reranker-ce-v1"
6tokenizer = AutoTokenizer.from_pretrained(model_path)
7model = AutoModelForSequenceClassification.from_pretrained(
8 model_path,
9 torch_dtype=torch.bfloat16
10)
11model.eval().cuda()
12
13# Score a query-document pair
14query = "What is llama?"
15document = "The llama is a domesticated South American camelid..."
16
17inputs = tokenizer(
18 f"query: {query}",
19 f"document: {document}",
20 return_tensors="pt",
21 truncation=True,
22 max_length=228,
23 padding="max_length"
24)
25inputs = {k: v.cuda() for k, v in inputs.items()}
26
27with torch.no_grad():
28 score = model(**inputs).logits.squeeze().item()
29
30print(f"Relevance score: {score}")1import torch
2from typing import List, Tuple
3from transformers import AutoTokenizer, AutoModelForSequenceClassification
4
5def load_reranker(model_path: str, device: str = "cuda"):
6 """Load the reranker model and tokenizer."""
7 tokenizer = AutoTokenizer.from_pretrained(model_path)
8 model = AutoModelForSequenceClassification.from_pretrained(
9 model_path,
10 torch_dtype=torch.bfloat16
11 )
12
13 # Configure padding token
14 if tokenizer.pad_token is None:
15 tokenizer.pad_token = tokenizer.eos_token
16 tokenizer.pad_token_id = tokenizer.eos_token_id
17 tokenizer.padding_side = "right"
18
19 model.eval()
20 model.to(device)
21 return tokenizer, model
22
23@torch.inference_mode()
24def rerank(
25 tokenizer,
26 model,
27 query: str,
28 documents: List[Tuple[str, str]], # (title, text)
29 batch_size: int = 64
30) -> List[Tuple[int, float]]:
31 """
32 Rerank documents for a query.
33
34 Returns:
35 List of (doc_index, score) sorted by relevance (descending)
36 """
37 device = next(model.parameters()).device
38 scores = []
39
40 for i in range(0, len(documents), batch_size):
41 batch = documents[i:i + batch_size]
42
43 # Prepare batch
44 queries = [f"query: {query}"] * len(batch)
45 docs = [f"document: {title} {text}" for title, text in batch]
46
47 inputs = tokenizer(
48 queries,
49 docs,
50 return_tensors="pt",
51 truncation=True,
52 max_length=228,
53 padding=True,
54 return_attention_mask=True
55 )
56 inputs = {k: v.to(device) for k, v in inputs.items()}
57
58 # Score batch
59 logits = model(**inputs).logits.squeeze(-1)
60 scores.extend(logits.cpu().tolist())
61
62 # Rank by score
63 ranked = sorted(enumerate(scores), key=lambda x: x[1], reverse=True)
64 return ranked
65
66
67# Example
68tokenizer, model = load_reranker("abdoelsayed/dear-8b-reranker-ce-v1")
69
70query = "When did Thomas Edison invent the light bulb?"
71documents = [
72 ("", "Lightning strike at Seoul National University"),
73 ("", "Thomas Edison tried to invent a device for car but failed"),
74 ("", "Coffee is good for diet"),
75 ("", "KEPCO fixes light problems"),
76 ("", "Thomas Edison invented the light bulb in 1879"),
77]
78
79ranking = rerank(tokenizer, model, query, documents)
80print(ranking)
81# Output: [(4, -2.015625), (1, -5.6875), (2, -6.375), (0, -6.5), (3, -6.78125)]
82# Document at index 4 is most relevant1{
2 "base_model": "meta-llama/Llama-3.1-8B",
3 "teacher_model": "abdoelsayed/llama2-13b-rankllama-teacher",
4 "loss": "Binary Cross-Entropy",
5 "distillation": {
6 "temperature": 2.0,
7 "alpha": 0.1
8 },
9 "optimizer": "AdamW",
10 "learning_rate": 1e-4,
11 "batch_size": 2,
12 "gradient_accumulation": 2,
13 "epochs": 2,
14 "max_length": 228,
15 "q_max_len": 32,
16 "p_max_len": 196,
17 "warmup_ratio": 0.1,
18 "weight_decay": 0.01,
19 "bf16": true
20}1L_total = (1 - α) * BCE(y_pred, y_true) + α * KL(σ(z_s/T), σ(z_t/T))
2
3where:
4- BCE: Binary cross-entropy loss
5- KL: KL divergence
6- z_s: Student logits
7- z_t: Teacher logits
8- T: Temperature (2.0)
9- α: Distillation weight (0.1)
10- σ: Sigmoid function| Dataset | NDCG@10 | NDCG@20 | MRR@10 | MAP |
|---|---|---|---|---|
| DL19 | 73.90 | 69.82 | 87.3 | 44.92 |
| DL20 | 72.10 | 68.45 | 85.1 | 42.67 |
| Dataset | NDCG@10 | NDCG@100 |
|---|---|---|
| MS MARCO | 68.5 | 75.2 |
| NQ | 51.8 | 69.4 |
| HotpotQA | 61.2 | 74.8 |
| FiQA | 46.8 | 62.3 |
| ArguAna | 58.9 | 71.5 |
| SciFact | 73.1 | 82.6 |
| TREC-COVID | 84.7 | 88.3 |
| NFCorpus | 39.4 | 51.7 |
| Average | 44.8 | 68.2 |
| Metric | Value |
|---|---|
| Inference Time (batch=64) | 2.2s |
| Throughput | ~45 docs/sec |
| GPU Memory (inference) | 18GB |
| Model Size (BF16) | 16GB |
| Model | Loss | DL19 | DL20 | BEIR Avg | Speed (s) |
|---|---|---|---|---|---|
| DeAR-8B-CE | BCE | 73.9 | 72.1 | 44.8 | 2.2 |
| DeAR-8B-RankNet | RankNet | 74.5 | 72.8 | 45.2 | 2.2 |
| MonoT5-3B | - | 71.8 | 68.9 | 43.5 | 3.5 |
| Teacher-13B | - | 73.8 | 71.2 | 44.8 | 5.8 |
Input Format: "query: [QUERY] document: [TITLE] [TEXT]"
↓
Tokenization (max_length=228)
↓
LLaMA-3.1-8B Transformer
↓
[CLS] Token Pooling
↓
Linear(hidden_size → 1)
↓
Sigmoid (optional)
↓
Relevance Score1from transformers import AutoModelForSequenceClassification, Trainer, TrainingArguments
2
3model = AutoModelForSequenceClassification.from_pretrained(
4 "abdoelsayed/dear-8b-reranker-ce-v1",
5 num_labels=1
6)
7
8training_args = TrainingArguments(
9 output_dir="./finetuned-model",
10 learning_rate=5e-6, # Lower LR for fine-tuning
11 per_device_train_batch_size=4,
12 num_train_epochs=1,
13 bf16=True,
14 logging_steps=100,
15)
16
17trainer = Trainer(
18 model=model,
19 args=training_args,
20 train_dataset=your_dataset,
21)
22
23trainer.train()1@article{abdallah2025dear,
2 title={DeAR: Dual-Stage Document Reranking with Reasoning Agents via LLM Distillation},
3 author={Abdallah, Abdelrahman and Mozafari, Jamshid and Piryani, Bhawna and Jatowt, Adam},
4 journal={arXiv preprint arXiv:2508.16998},
5 year={2025}
6}