ColFlor: Towards BERT-Size Vision-Language Document Retrieval Models

In June 2024, ColPali was introduced as an OCR-free document retrieval model, built over PaliGemma, shifting the paradigm of PDF document retrieval by directly processing images instead of using error-prone and resource-heavy OCR pipelines. However, with three billion parameters, ColPali might be computationally expensive, especially for large document databases. In contrast, text retrieval models like ColBERT are more efficient with just a few hundred million parameters, but they require error-prone and expensive OCR pipelines to. To bridge this gap, we introduce ColFlor, an OCR-free visual document retrieval model with only 174 million parameters. ColFlor is 17 times smaller than ColPali, 9.8 times faster in encoding queries and 5.25 faster in encoding images, with only a 1.8% drop in performance on text-rich English documents.

More details about the model can be found in the ColFlor blogpost

Usage

First, you need to clone the github repo and install the dependencies as follows

git clone https://github.com/AhmedMasryKU/colflor
cd colflor
pip install . -e

Then, you can run the following inference code:

import pprint
from typing import List, cast

import torch
from datasets import Dataset, load_dataset
from torch.utils.data import DataLoader
from tqdm import tqdm

from colpali_engine.models import ColFlor
from colpali_engine.models import ColFlorProcessor
from colpali_engine.utils.processing_utils import BaseVisualRetrieverProcessor
from colpali_engine.utils.torch_utils import ListDataset, get_torch_device


def main():
    """
    Example script to run inference with ColFlor.
    """

    device = get_torch_device("auto")
    print(f"Device used: {device}")

    # Model name
    model_name = "ahmed-masry/ColFlor"

    # Load model
    model = ColFlor.from_pretrained(
        model_name,
        #torch_dtype=torch.bfloat16,
        device_map=device,
    ).eval()

    # Load processor
    processor = cast(ColFlorProcessor, ColFlorProcessor.from_pretrained(model_name))

    if not isinstance(processor, BaseVisualRetrieverProcessor):
        raise ValueError("Processor should be a BaseVisualRetrieverProcessor")

    # NOTE: Only the first 16 images are used for demonstration purposes
    dataset = cast(Dataset, load_dataset("vidore/docvqa_test_subsampled", split="test[:16]"))
    images = dataset["image"]

    # Select a few queries for demonstration purposes
    query_indices = [12, 15]
    queries = [dataset[idx]["query"] for idx in query_indices]
    print("Selected queries:")
    pprint.pprint(dict(zip(query_indices, queries)))

    # Run inference - docs
    dataloader = DataLoader(
        dataset=ListDataset[str](images),
        batch_size=4,
        shuffle=False,
        collate_fn=lambda x: processor.process_images(x),
    )
    ds: List[torch.Tensor] = []
    for batch_doc in tqdm(dataloader):
        with torch.no_grad():
            batch_doc = {k: v.to(model.device) for k, v in batch_doc.items()}
            embeddings_doc = model(**batch_doc)
        ds.extend(list(torch.unbind(embeddings_doc.to("cpu"))))

    # Run inference - queries
    dataloader = DataLoader(
        dataset=ListDataset[str](queries),
        batch_size=4,
        shuffle=False,
        collate_fn=lambda x: processor.process_queries(x),
    )

    qs: List[torch.Tensor] = []
    for batch_query in dataloader:
        with torch.no_grad():
            batch_query = {k: v.to(model.device) for k, v in batch_query.items()}
            embeddings_query = model(**batch_query)
        qs.extend(list(torch.unbind(embeddings_query.to("cpu"))))

    # Run scoring
    scores = processor.score(qs, ds).cpu().numpy()
    idx_top_1 = scores.argmax(axis=1)
    print("Indices of the top-1 retrieved documents for each query:", idx_top_1)

    # Sanity check
    if idx_top_1.tolist() == query_indices:
        print("The top-1 retrieved documents are correct.")
    else:
        print("The top-1 retrieved documents are incorrect.")

    return


if __name__ == "__main__":
    typer.run(main)

Limitations

Figures: While ColFlor exhibits reasonable performance on figures, there's a relatively large gap in performance between it and larger models such as ColPali.
Multilinguality: The current version of the model only supports the Engligh language and performs poorly on other languages.

License

We release this model under the MIT license.

Contact

If you have any questions about this work, feel free to reach out to Ahmed Masry at [email protected] or [email protected].

Acknowledgement

This work was carried out at the Intelligent Visualization Lab at York University in Canada. It was supported by the Natural Sciences Engineering Research Council (NSERC) of Canada and Canada Foundation for Innovation (CFI). Additionally, it received support through a GCP credits award from Google's PaliGemma Academic Program.

We appreciate the well-documented training and evaluation GitHub repositories provided by the ColPali team, which were essential in our model development. This model card is adapted from ColPali Model Card

Citation

If you plan to use ColFlor in your research, please consider citing us as follows:

@misc{masry2024colflor,
    title={ColFlor: Towards BERT-Size Vision-Language Document Retrieval Models},
    url={https://huggingface.co/blog/ahmed-masry/colflor},
    author={Masry, Ahmed},
    month={October},
    year={2024}
}

ahmed-masry
/

ColFlor