FLMR

The huggingface-transformers implementation of Fine-grained Late-interaction Multi-modal Retriever.

The official implementation is at here.

The details of the model and checkpoints can be found here.

The details for reproducing the datasets and evaluation in the paper can be found here.

Updates

• [26/1/2025] 🔥🔥🔥 The M2KR benchmark is now featured in the Efficient Representation Learning for Multimodal Information Retrieval Workshop at WWW 2025! Check it out here: Multimodal Document Retrieval Challenge
• [19/12/2024] 🔥🔥🔥 We released the multilingual version (Chinese + English) of PreFLMR, you can download PreFLMR ENCN model here.
• [03/09/2024] We have uploaded the images used in the M2KR benchmark here .
• [10/08/2024] We received many requests regarding adding multilingual abilities to PreFLMR. We announce that we are now training the Chinese version of PreFLMR and will release it very soon. Stay tuned!
• [05/06/2024] 🔥🔥🔥We made some updates to the implementation
  • Added an evaluation script that reproduces the results in the PreFLMR paper here
  • Added the updated benchmark results with the transformer implementation here
  • Added an example script to fine-tune PreFLMR on a custom retrieval dataset here
  • IMPORTANT: fixed the OVEN data splits in the M2KR benchmark, and updated each entry with a fixed instruction to ensure the evaluation result is not affected by random sampling of instructions. Please delete your local cache and download the dataset again.

Table of Contents

• FLMR
  • Updates
  • Table of Contents
  • Models and Benchmark Results
  • Models and Benchmark Results for the M2KR Amd M2KR-CN (the Chinese split of the M2KR)
  • How to use this package
    • Environment
    • Index a custom document collection
    • Search a custom document collection
    • Training with contrastive learning
  • Alternative: use transformers.AutoModel to load pre-trained models
  • Use example scripts
    • Use FLMR
    • [NEW!] Use PreFLMR_ENCN
    • [NEW!] Use PreFLMR
    • [NEW!] Evaluate the PreFLMR models on all M2KR benchmarks
    • [NEW!] Finetune the PreFLMR model on downstream datasets
      • Run finetuning
      • Run Testing
      • Example finetuning results
  • Note
  • Citation

Models and Benchmark Results

| Model | WIT Recall@10 | IGLUE Recall@1 | KVQA Recall@5 | MSMARCO Recall@5 | OVEN Recall@5 | LLaVA Recall@1 | EVQA Recall@5 | EVQA Pseudo Recall@5 | OKVQA Recall@5 | OKVQA Pseudo Recall@5 | Infoseek Recall@5 | Infoseek Pseudo Recall@5 | |---------------|---------------|----------------|---------------|------------------|---------------|----------------|---------------|----------------------|----------------|-----------------------|-------------------|--------------------------| | LinWeizheDragon/PreFLMR_ViT-G🤗 | 0.619 | 0.718 | 0.419 | 0.783 | 0.643 | 0.726 | 0.625 | 0.721 | 0.302 | 0.674 | 0.392 | 0.577 | | LinWeizheDragon/PreFLMR_ViT-L🤗 | 0.605 | 0.699 | 0.440 | 0.779 | 0.608 | 0.729 | 0.609 | 0.708 | 0.314 | 0.690 | 0.374 | 0.578 | | LinWeizheDragon/PreFLMR_ViT-B🤗 | 0.427 | 0.574 | 0.294 | 0.786 | 0.468 | 0.673 | 0.550 | 0.663 | 0.272 | 0.658 | 0.260 | 0.496 |

Note: We converted the checkpoints from PyTorch to Huggingface-transformers, whose benchmark results differ from the numbers reported in the original paper slightly. You can reproduce the results in the above paper by referring to the instructions in this document.

Models and Benchmark Results for the M2KR and M2KR-CN (the Chinese split of the M2KR)

| Model | WIT(EN) Recall@10 | WIT(CN) Recall@10 | KVQA(EN) Recall@5 | KVQA(EN) Recall@5 | MSMARCO(EN) Recall@5 | MSMARCO(CN) Recall@5 | OVEN(EN) Recall@5 | OVEN(CN) Recall@5 | LLaVA(EN) Recall@1 | LLaVA(CN) Recall@1 | EVQA(EN) Recall@5 | EVQA(CN) Recall@5 | OKVQA(EN) Recall@5 | OKVQA(CN) Recall@5 | Infoseek(EN) Recall@5 | Infoseek(CN) Recall@5 | | :----------------------------------------------------------: | :---------------: | :---------------: | :---------------: | :---------------: | :------------------: | :------------------: | :---------------: | :---------------: | :----------------: | :----------------: | :---------------: | :---------------: | :----------------: | :----------------: | :-------------------: | :-------------------: | | LinWeizheDragon/PreFLMR_ViT-L🤗 | 60.5 | 10.9 | 43.6 | 3.2 | 78.7 | 10.3 | 59.8 | 6.6 | 71.8 | 3.2 | 70.8 | 2.8 | 68.5 | 2.1 | 57.9 | 7.9 | | LinWeizheDragon/PreFLMR_ViT-L_ENCN🤗 | 60.8 | 83.4 | 41.1 | 37.3 | 82.6 | 82.3 | 60.8 | 58.8 | 71.1 | 58.9 | 58. 0 | 46.6 | 13.9 | 13.3 | 41.9 | 39.7 |

How to use this package

Environment

Create virtualenv:

conda create -n FLMR python=3.10 -y
conda activate FLMR

Install Pytorch:

pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

Install faiss

conda install -c pytorch -c nvidia faiss-gpu=1.7.4 mkl=2021 blas=1.0=mkl

Test if faiss generate error

python -c "import faiss"

Install FLMR

git clone https://github.com/LinWeizheDragon/FLMR.git
cd FLMR
pip install -e .

Install ColBERT engine

cd third_party/ColBERT
pip install -e .

Install other dependencies

pip install ujson gitpython easydict ninja datasets transformers==4.49

Index a custom document collection

1. Load pre-trained models

   import os
   import torch
   import pandas as pd
   import numpy as np
   from torchvision.transforms import ToPILImage
   from transformers import AutoImageProcessor
   
   from flmr import index_custom_collection
   from flmr import FLMRQueryEncoderTokenizer, FLMRContextEncoderTokenizer, FLMRModelForRetrieval, FLMRConfig
   
   # load models
   checkpoint_path = "LinWeizheDragon/PreFLMR_ViT-G"
   image_processor_name = "laion/CLIP-ViT-bigG-14-laion2B-39B-b160k"
   flmr_config = FLMRConfig.from_pretrained(checkpoint_path)
   
   query_tokenizer = FLMRQueryEncoderTokenizer.from_pretrained(checkpoint_path, text_config=flmr_config.text_config, subfolder="query_tokenizer")
   context_tokenizer = FLMRContextEncoderTokenizer.from_pretrained(
       checkpoint_path, text_config=flmr_config.text_config, subfolder="context_tokenizer"
   )
   
   model = FLMRModelForRetrieval.from_pretrained(
       checkpoint_path,
       query_tokenizer=query_tokenizer,
       context_tokenizer=context_tokenizer,
   )
   image_processor = AutoImageProcessor.from_pretrained(image_processor_name)
   

2. Create document collections

   num_items = 100
   feature_dim = 1664
   passage_contents = [f"This is test sentence {i}" for i in range(num_items)]
   # Option 1. text-only documents
   custom_collection = passage_contents
   # Option 2. multi-modal documents with pre-extracted image features
   # passage_image_features = np.random.rand(num_items, feature_dim)
   # custom_collection = [
   #     (passage_content, passage_image_feature, None) for passage_content, passage_image_feature in zip(passage_contents, passage_image_features)
   # ]
   # Option 3. multi-modal documents with images
   # random_images = torch.randn(num_items, 3, 224, 224)
   # to_img = ToPILImage()
   # if not os.path.exists("./test_images"):
   #     os.makedirs("./test_images")
   # for i, image in enumerate(random_images):
   #     image = to_img(image)
   #     image.save(os.path.join("./test_images", "{}.jpg".format(i)))
   
   # image_paths = [os.path.join("./test_images", "{}.jpg".format(i)) for i in range(num_items)]
   
   # custom_collection = [
   #     (passage_content, None, image_path)
   #     for pa