Views
No views yet
1import torch
2from transformers import AutoTokenizer, AutoModelForSequenceClassification
3from peft import PeftModel, PeftConfig
4
5# Load LoRA adapter
6adapter_path = "abdoelsayed/dear-3b-reranker-ranknet-lora-v1"
7config = PeftConfig.from_pretrained(adapter_path)
8
9# Load tokenizer
10tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)
11if tokenizer.pad_token is None:
12 tokenizer.pad_token = tokenizer.eos_token
13
14# Load base model
15base_model = AutoModelForSequenceClassification.from_pretrained(
16 config.base_model_name_or_path,
17 num_labels=1,
18 torch_dtype=torch.bfloat16
19)
20
21# Load and merge LoRA
22model = PeftModel.from_pretrained(base_model, adapter_path)
23model = model.merge_and_unload()
24model.eval().cuda()
25
26# Use model
27query = "What is machine learning?"
28document = "Machine learning is a subset of artificial intelligence..."
29
30inputs = tokenizer(
31 f"query: {query}",
32 f"document: {document}",
33 return_tensors="pt",
34 truncation=True,
35 max_length=228,
36 padding="max_length"
37)
38inputs = {k: v.cuda() for k, v in inputs.items()}
39
40with torch.no_grad():
41 score = model(**inputs).logits.squeeze().item()
42print(f"Relevance score: {score}")1from typing import List, Tuple
2
3def load_3b_lora_ranker(adapter_path: str):
4 """Load 3B LoRA adapter efficiently."""
5 config = PeftConfig.from_pretrained(adapter_path)
6
7 tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)
8 if tokenizer.pad_token is None:
9 tokenizer.pad_token = tokenizer.eos_token
10
11 base = AutoModelForSequenceClassification.from_pretrained(
12 config.base_model_name_or_path,
13 num_labels=1,
14 torch_dtype=torch.bfloat16
15 )
16
17 model = PeftModel.from_pretrained(base, adapter_path)
18 model = model.merge_and_unload()
19 model.eval().cuda()
20
21 return tokenizer, model
22
23# Load once
24tokenizer, model = load_3b_lora_ranker("abdoelsayed/dear-3b-reranker-ranknet-lora-v1")
25
26# Rerank function
27@torch.inference_mode()
28def rerank(tokenizer, model, query: str, docs: List[Tuple[str, str]], batch_size=128):
29 scores = []
30 device = next(model.parameters()).device
31
32 for i in range(0, len(docs), batch_size):
33 batch = docs[i:i + batch_size]
34 queries = [f"query: {query}"] * len(batch)
35 documents = [f"document: {t} {p}" for t, p in batch]
36
37 inputs = tokenizer(queries, documents, return_tensors="pt",
38 truncation=True, max_length=228, padding=True)
39 inputs = {k: v.to(device) for k, v in inputs.items()}
40
41 logits = model(**inputs).logits.squeeze(-1)
42 scores.extend(logits.cpu().tolist())
43
44 return sorted(enumerate(scores), key=lambda x: x[1], reverse=True)1{
2 "r": 16,
3 "lora_alpha": 32,
4 "target_modules": [
5 "q_proj", "v_proj", "k_proj", "o_proj",
6 "gate_proj", "up_proj", "down_proj"
7 ],
8 "lora_dropout": 0.05,
9 "bias": "none",
10 "task_type": "SEQ_CLS"
11}LoRA Adapter: 40MB
Full 3B Model: 6GB
Full 8B Model: 16GB
Ratio: 0.67% of 3B, 0.25% of 8B3B LoRA: 1.5s (100 docs)
8B Full: 2.2s (100 docs)
Speedup: 1.47x faster than 8B3B LoRA: 10GB GPU
3B Full: 12GB GPU
8B Full: 18GB GPU1# Minimal memory deployment
2import torch
3from transformers import AutoModelForSequenceClassification
4from peft import PeftModel
5
6adapter_path = "abdoelsayed/dear-3b-reranker-ranknet-lora-v1"
7
8# Load with memory optimization
9model = AutoModelForSequenceClassification.from_pretrained(
10 "meta-llama/Llama-3.2-3B",
11 num_labels=1,
12 torch_dtype=torch.bfloat16,
13 device_map="auto",
14 low_cpu_mem_usage=True
15)
16
17# Load adapter
18model = PeftModel.from_pretrained(model, adapter_path)
19model = model.merge_and_unload()
20model.eval()
21
22# Optional: Compile for speedup
23if hasattr(torch, 'compile'):
24 model = torch.compile(model, mode="max-autotune")Model Size vs NDCG@10 (TREC DL19):
├─ Teacher-13B: 73.8 (26GB)
├─ DeAR-8B-Full: 74.5 (16GB)
├─ DeAR-8B-LoRA: 74.2 (100MB + base)
├─ DeAR-3B-Full: 71.2 (6GB)
└─ DeAR-3B-LoRA: 70.9 (40MB + base) ← This model
Best Efficiency: 95% accuracy at 0.25% size of 8B!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}