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-8b-reranker-ranknet-lora-v1"
7
8# Get base model from adapter config
9config = PeftConfig.from_pretrained(adapter_path)
10base_model_name = config.base_model_name_or_path
11
12# Load tokenizer
13tokenizer = AutoTokenizer.from_pretrained(base_model_name)
14if tokenizer.pad_token is None:
15 tokenizer.pad_token = tokenizer.eos_token
16 tokenizer.pad_token_id = tokenizer.eos_token_id
17
18# Load base model
19base_model = AutoModelForSequenceClassification.from_pretrained(
20 base_model_name,
21 num_labels=1,
22 torch_dtype=torch.bfloat16
23)
24
25# Load and merge LoRA adapter
26model = PeftModel.from_pretrained(base_model, adapter_path)
27model = model.merge_and_unload() # Merge adapter into base model
28
29model.eval().cuda()
30
31# Use the model
32query = "What is machine learning?"
33document = "Machine learning is a subset of artificial intelligence..."
34
35inputs = tokenizer(
36 f"query: {query}",
37 f"document: {document}",
38 return_tensors="pt",
39 truncation=True,
40 max_length=228,
41 padding="max_length"
42)
43inputs = {k: v.cuda() for k, v in inputs.items()}
44
45with torch.no_grad():
46 score = model(**inputs).logits.squeeze().item()
47
48print(f"Relevance score: {score}")1import torch
2from typing import List, Tuple
3from transformers import AutoTokenizer, AutoModelForSequenceClassification
4from peft import PeftModel, PeftConfig
5
6def load_lora_ranker(adapter_path: str, device: str = "cuda"):
7 """Load LoRA adapter and merge with base model."""
8 # Get base model path from adapter config
9 peft_config = PeftConfig.from_pretrained(adapter_path)
10 base_model_name = peft_config.base_model_name_or_path
11
12 # Load tokenizer
13 tokenizer = AutoTokenizer.from_pretrained(base_model_name)
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 # Load base model
20 base_model = AutoModelForSequenceClassification.from_pretrained(
21 base_model_name,
22 num_labels=1,
23 torch_dtype=torch.bfloat16
24 )
25
26 # Load LoRA adapter and merge
27 model = PeftModel.from_pretrained(base_model, adapter_path)
28 model = model.merge_and_unload()
29
30 model.eval().to(device)
31 return tokenizer, model
32
33# Load model
34tokenizer, model = load_lora_ranker("abdoelsayed/dear-8b-reranker-ranknet-lora-v1")
35
36# Rerank documents
37@torch.inference_mode()
38def rerank(tokenizer, model, query: str, docs: List[Tuple[str, str]], batch_size: int = 64):
39 """Rerank documents for a query."""
40 device = next(model.parameters()).device
41 scores = []
42
43 for i in range(0, len(docs), batch_size):
44 batch = docs[i:i + batch_size]
45 queries = [f"query: {query}"] * len(batch)
46 documents = [f"document: {title} {text}" for title, text in batch]
47
48 inputs = tokenizer(
49 queries,
50 documents,
51 return_tensors="pt",
52 truncation=True,
53 max_length=228,
54 padding=True
55 )
56 inputs = {k: v.to(device) for k, v in inputs.items()}
57
58 logits = model(**inputs).logits.squeeze(-1)
59 scores.extend(logits.cpu().tolist())
60
61 return sorted(enumerate(scores), key=lambda x: x[1], reverse=True)
62
63# Example
64query = "When did Thomas Edison invent the light bulb?"
65docs = [
66 ("", "Thomas Edison invented the light bulb in 1879"),
67 ("", "Coffee is good for diet"),
68 ("", "Lightning strike at Seoul"),
69]
70
71ranking = rerank(tokenizer, model, query, docs)
72print(ranking) # [(0, 5.2), (2, -3.1), (1, -4.8)]1from peft import PeftModel, PeftConfig
2from transformers import AutoModelForSequenceClassification
3
4adapter_path = "abdoelsayed/dear-8b-reranker-ranknet-lora-v1"
5config = PeftConfig.from_pretrained(adapter_path)
6
7# Load base model
8base_model = AutoModelForSequenceClassification.from_pretrained(
9 config.base_model_name_or_path,
10 num_labels=1,
11 torch_dtype=torch.bfloat16,
12 device_map="auto"
13)
14
15# Load adapter (without merging)
16model = PeftModel.from_pretrained(base_model, adapter_path)
17model.eval()
18
19# Use model (adapter layers will be applied automatically)
20# ... same inference code as above ...| Benchmark | LoRA | Full Model | Difference |
|---|---|---|---|
| TREC DL19 | 74.2 | 74.5 | -0.3 |
| TREC DL20 | 72.5 | 72.8 | -0.3 |
| BEIR (Avg) | 44.9 | 45.2 | -0.3 |
| MS MARCO | 68.6 | 68.9 | -0.3 |
1lora_config = {
2 "r": 16, # LoRA rank
3 "lora_alpha": 32, # Scaling factor
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}1training_args = {
2 "learning_rate": 1e-4, # Higher than full fine-tuning
3 "batch_size": 4, # Larger batch possible due to lower memory
4 "gradient_accumulation": 2,
5 "epochs": 2,
6 "warmup_ratio": 0.1,
7 "weight_decay": 0.01,
8 "max_length": 228,
9 "bf16": True
10}| Aspect | LoRA | Full Model |
|---|---|---|
| Storage | 100MB | 16GB |
| Training Time | 12h | 36h |
| Training Memory | 28GB | 38GB |
| Performance | 98% | 100% |
| Loading Time | Fast | Slow |
| Easy Updates | ✅ Yes | ❌ No |
1from peft import LoraConfig, get_peft_model, TaskType
2from transformers import AutoModelForSequenceClassification, TrainingArguments, Trainer
3
4# Load base model
5base_model = AutoModelForSequenceClassification.from_pretrained(
6 "meta-llama/Llama-3.1-8B",
7 num_labels=1
8)
9
10# Configure LoRA
11lora_config = LoraConfig(
12 task_type=TaskType.SEQ_CLS,
13 r=16,
14 lora_alpha=32,
15 target_modules=["q_proj", "v_proj", "k_proj", "o_proj"],
16 lora_dropout=0.05,
17 bias="none",
18)
19
20# Apply LoRA
21model = get_peft_model(base_model, lora_config)
22model.print_trainable_parameters()
23# Output: trainable params: 67M || all params: 8B || trainable%: 0.8%
24
25# Train
26training_args = TrainingArguments(
27 output_dir="./lora-finetuned",
28 learning_rate=1e-4,
29 per_device_train_batch_size=8,
30 num_train_epochs=3,
31 bf16=True,
32)
33
34trainer = Trainer(
35 model=model,
36 args=training_args,
37 train_dataset=your_dataset,
38)
39
40trainer.train()
41
42# Save only the LoRA adapter
43model.save_pretrained("./lora-adapter")adapter_config.json - LoRA configurationadapter_model.safetensors or adapter_model.bin - Adapter weights (~100MB)README.md - This documentation1@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}