|
|
| --- |
| language: en |
| tags: |
| - reranker |
| - RAG |
| - multimodal |
| - vision-language |
| - Qwen |
| license: cc-by-4.0 |
| pipeline_tag: visual-document-retrieval |
| --- |
| |
| # DocReRank: Multi-Modal Reranker |
|
|
| This is the official model from the paper: |
|
|
| π **[DocReRank: Single-Page Hard Negative Query Generation for Training Multi-Modal RAG Rerankers](https://arxiv.org/abs/2505.22584)** |
|
|
| See [Project Page](https://navvewas.github.io/DocReRank/) for more information. |
|
|
| --- |
|
|
| ## β
Model Overview |
| - **Base model:** [Qwen/Qwen2-VL-2B-Instruct](https://huggingface.co/Qwen/Qwen2-VL-2B-Instruct) |
| - **Architecture:** Vision-Language reranker |
| - **Fine-tuning method:** PEFT (LoRA) |
| - **Training data:** Generated by **Single-Page Hard Negative Query Generation** Pipeline. |
| - **Purpose:** Improves second-stage reranking for Retrieval-Augmented Generation (RAG) in multimodal scenarios. |
|
|
| --- |
|
|
| ## β
How to Use |
|
|
| This adapter requires the base Qwen2-VL model. |
|
|
| ```python |
| from transformers import AutoProcessor, Qwen2VLForConditionalGeneration |
| from peft import PeftModel |
| import torch |
| from PIL import Image |
| from huggingface_hub import hf_hub_download |
| import os |
| |
| # β
Load base model |
| base_model = Qwen2VLForConditionalGeneration.from_pretrained( |
| "Qwen/Qwen2-VL-2B-Instruct", |
| torch_dtype=torch.bfloat16, |
| device_map="cuda" |
| ) |
| |
| # β
Load DocReRank adapter |
| model = PeftModel.from_pretrained(base_model, "DocReRank/DocReRank-Reranker").eval() |
| |
| # β
Load processor |
| processor = AutoProcessor.from_pretrained("Qwen/Qwen2-VL-2B-Instruct") |
| processor.image_processor.min_pixels = 200704 |
| processor.image_processor.max_pixels = 589824 |
| |
| # β
Define query and images |
| query_text = "What are the performances of the DocReRank model on restaurant and biomedical benchmarks?" |
| # query_text = "Are there ablation results for the DocReRank model?" |
| |
| # Downloading Pages for Demo |
| save_dir = os.path.join(os.getcwd(), "paper_pages") |
| os.makedirs(save_dir, exist_ok=True) |
| image_files = ["DocReRank_paper_page_2.png","DocReRank_paper_page_4.png","DocReRank_paper_page_6.png","DocReRank_paper_page_8.png"] |
| local_paths = [] |
| for f in image_files: |
| local_path = hf_hub_download(repo_id="DocReRank/DocReRank-Reranker",filename=f,local_dir=save_dir) |
| local_paths.append(local_path) |
| print("β
Files downloaded to:", local_paths) |
| |
| image_paths = [ "paper_pages/DocReRank_paper_page_2.png", "paper_pages/DocReRank_paper_page_4.png", "paper_pages/DocReRank_paper_page_6.png", "paper_pages/DocReRank_paper_page_8.png"] |
| # β
Reranking prompt template |
| def compute_score(image_path, query_text): |
| image = Image.open(image_path) |
| prompt = f"Assert the relevance of the previous image document to the following query, answer True or False. The query is: {query_text}" |
| messages = [{"role": "user", "content": [{"type": "image", "image": image}, {"type": "text", "text": prompt}]}] |
| |
| # Tokenize |
| text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) |
| inputs = processor(text=text, images=image, return_tensors="pt").to(model.device, torch.bfloat16) |
| |
| # Compute logits |
| with torch.no_grad(): |
| outputs = model(**inputs) |
| logits = outputs.logits[:, -1, :] |
| true_id = processor.tokenizer.convert_tokens_to_ids("True") |
| false_id = processor.tokenizer.convert_tokens_to_ids("False") |
| probs = torch.softmax(logits[:, [true_id, false_id]], dim=-1) |
| relevance_score = probs[0, 0].item() # Probability of "True" |
| |
| return relevance_score |
| |
| # β
Compute scores for both images |
| scores = [(img, compute_score(img, query_text)) for img in image_paths] |
| |
| # β
Print results |
| for img, score in scores: |
| print(f"Image: {img} | Relevance Score: {score:.4f}") |
| ``` |
|
|
| ## Citation |
| If you use this dataset, please cite: |
| ```bibtex |
| @article{wasserman2025docrerank, |
| title={DocReRank: Single-Page Hard Negative Query Generation for Training Multi-Modal RAG Rerankers}, |
| author={Wasserman, Navve and Heinimann, Oliver and Golbari, Yuval and Zimbalist, Tal and Schwartz, Eli and Irani, Michal}, |
| journal={arXiv preprint arXiv:2505.22584}, |
| year={2025} |
| } |
| ``` |
|
|
|
|
|
|