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-ce-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# Score query-document pair
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}")1@torch.inference_mode()
2def rerank(tokenizer, model, query: str, documents, batch_size=64):
3 scores = []
4 device = next(model.parameters()).device
5
6 for i in range(0, len(documents), batch_size):
7 batch = documents[i:i + batch_size]
8 queries = [f"query: {query}"] * len(batch)
9 docs = [f"document: {title} {text}" for title, text in batch]
10
11 inputs = tokenizer(queries, docs, return_tensors="pt",
12 truncation=True, max_length=228, padding=True)
13 inputs = {k: v.to(device) for k, v in inputs.items()}
14
15 logits = model(**inputs).logits.squeeze(-1)
16 scores.extend(logits.cpu().tolist())
17
18 return sorted(enumerate(scores), key=lambda x: x[1], reverse=True)1{
2 "r": 16,
3 "lora_alpha": 32,
4 "target_modules": ["q_proj", "v_proj", "k_proj", "o_proj",
5 "gate_proj", "up_proj", "down_proj"],
6 "lora_dropout": 0.05,
7 "bias": "none",
8 "task_type": "SEQ_CLS"
9}| Feature | LoRA | Full Model |
|---|---|---|
| Storage | 100MB | 16GB |
| Training Time | 12h | 34h |
| Performance | 98% | 100% |
| Memory | 28GB | 38GB |
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}