|
--- |
|
license: mit |
|
--- |
|
|
|
# ColFlor: Towards BERT-Size Vision-Language Document Retrieval Models |
|
|
|
In June 2024, [ColPali](https://arxiv.org/abs/2407.01449) was introduced as an OCR-free document retrieval model, built over [PaliGemma](https://arxiv.org/abs/2407.07726), 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](https://arxiv.org/abs/2004.12832) 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 130 million parameters. |
|
|
|
<p align="center"><img width=800 src="https://github.com/AhmedMasryKU/colflor/blob/main/assets/colflor.png?raw=true"/></p> |
|
|
|
More details about the model can be found in the [ColFlor blogpost](https://huggingface.co/blog/ahmed-masry/colflor) |
|
|
|
## Usage |
|
|
|
First, you need to clone the github repo and install the dependencies as follows |
|
|
|
|
|
```bash |
|
git clone https://github.com/AhmedMasryKU/colflor |
|
cd colflor |
|
pip install . -e |
|
``` |
|
|
|
Then, you can run the following inference code: |
|
|
|
```python |
|
import torch |
|
import typer |
|
from torch.utils.data import DataLoader |
|
from tqdm import tqdm |
|
from transformers import AutoProcessor |
|
from PIL import Image |
|
|
|
from colpali_engine.models.paligemma_colbert_architecture import ColPali |
|
from colpali_engine.trainer.retrieval_evaluator import CustomEvaluator |
|
from colpali_engine.utils.colpali_processing_utils import process_images, process_queries |
|
from colpali_engine.utils.image_from_page_utils import load_from_dataset |
|
|
|
|
|
def main() -> None: |
|
"""Example script to run inference with ColPali""" |
|
|
|
# Load model |
|
model_name = "vidore/colpali" |
|
model = ColPali.from_pretrained("vidore/colpaligemma-3b-mix-448-base", torch_dtype=torch.bfloat16, device_map="cuda").eval() |
|
model.load_adapter(model_name) |
|
processor = AutoProcessor.from_pretrained(model_name) |
|
|
|
# select images -> load_from_pdf(<pdf_path>), load_from_image_urls(["<url_1>"]), load_from_dataset(<path>) |
|
images = load_from_dataset("vidore/docvqa_test_subsampled") |
|
queries = ["From which university does James V. Fiorca come ?", "Who is the japanese prime minister?"] |
|
|
|
# run inference - docs |
|
dataloader = DataLoader( |
|
images, |
|
batch_size=4, |
|
shuffle=False, |
|
collate_fn=lambda x: process_images(processor, x), |
|
) |
|
ds = [] |
|
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( |
|
queries, |
|
batch_size=4, |
|
shuffle=False, |
|
collate_fn=lambda x: process_queries(processor, x, Image.new("RGB", (448, 448), (255, 255, 255))), |
|
) |
|
|
|
qs = [] |
|
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 evaluation |
|
retriever_evaluator = CustomEvaluator(is_multi_vector=True) |
|
scores = retriever_evaluator.evaluate(qs, ds) |
|
print(scores.argmax(axis=1)) |
|
|
|
|
|
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 **masry20@yorku.ca** or **ahmed.elmasry24653@gmail.com**. |
|
|
|
## 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 instrumental in our model development. |
|
This model card is adapted from [ColPali Model Card](https://huggingface.co/vidore/colpali) |
|
|
|
## Citation |
|
|
|
If you plan to use ColFlor in your research, please consider citing us as follows: |
|
|
|
```bibtex |
|
@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} |
|
} |
|
``` |
