colSmol-256M

ColSmolVLM-Instruct-256M: Visual Retriever based on SmolVLM-Instruct-250M with ColBERT strategy

This is a version trained with batch_size 32 for 3 epochs

ColSmolVLM is a model based on a novel model architecture and training strategy based on Vision Language Models (VLMs) to efficiently index documents from their visual features. It is a SmolVLM extension that generates ColBERT- style multi-vector representations of text and images. It was introduced in the paper ColPali: Efficient Document Retrieval with Vision Language Models and first released in this repository

<p align="center"><img width=800 src="https://github.com/illuin-tech/colpali/blob/main/assets/colpali_architecture.webp?raw=true"/></p>

Version specificity

This version is trained with the commit b983e40 of the Colpali repository. (main branch from the repo)

Data is the same as the ColPali data described in the paper.

Model Training

Dataset

Our training dataset of 127,460 query-page pairs is comprised of train sets of openly available academic datasets (63%) and a synthetic dataset made up of pages from web-crawled PDF documents and augmented with VLM-generated (Claude-3 Sonnet) pseudo-questions (37%). Our training set is fully English by design, enabling us to study zero-shot generalization to non-English languages. We explicitly verify no multi-page PDF document is used both ViDoRe and in the train set to prevent evaluation contamination. A validation set is created with 2% of the samples to tune hyperparameters.

Note: Multilingual data is present in the pretraining corpus of the language model and most probably in the multimodal training.

Parameters

Unless specified otherwise, we train models in bfloat16 format, use low-rank adapters (LoRA) with alpha=32 and r=32 on the transformer layers from the language model, as well as the final randomly initialized projection layer, and use a paged_adamw_8bit optimizer. We train on a 4 GPU setup with data parallelism, a learning rate of 5e-4 with linear decay with 2.5% warmup steps, and a batch size of 8.

Usage

Using Sentence Transformers

ColSmolVLM can be used as a multi-vector (ColBERT-style late interaction) retriever directly with Sentence Transformers via the MultiVectorEncoder:

pip install "sentence-transformers[image]>=6.0.0"

from sentence_transformers import MultiVectorEncoder

model = MultiVectorEncoder("vidore/colSmol-256M")

queries = [
    "What is the variable represented on the y-axis of the graph?",
    "Total outlay is maximum in which year?",
]
images = [
    "https://huggingface.co/datasets/sentence-transformers/example-documents/resolve/main/doc1.jpg",
    "https://huggingface.co/datasets/sentence-transformers/example-documents/resolve/main/doc2.jpg",
    "https://huggingface.co/datasets/sentence-transformers/example-documents/resolve/main/doc3.jpg",
    "https://huggingface.co/datasets/sentence-transformers/example-documents/resolve/main/doc4.jpg",
]

query_embeddings = model.encode_query(queries)
document_embeddings = model.encode_document(images)
print(f"Query 0 shape:    {tuple(query_embeddings[0].shape)}")
print(f"Document 0 shape: {tuple(document_embeddings[0].shape)}")
# Query 0 shape:    (27, 128)
# Document 0 shape: (1135, 128)

# MaxSim late-interaction scoring (rows = queries, columns = images)
scores = model.similarity(query_embeddings, document_embeddings)
print(scores)
# tensor([[18.1855, 16.2119, 11.7363,  9.7974],
#         [ 9.2637, 15.1357, 10.4395,  8.0791]])

Using ColPali Engine

[!WARNING] Note: current colpali-engine no longer sends the query prefix and trailing newline that this checkpoint was trained with. The trailing newline went in 0.3.11 (illuin-tech/colpali#280) and the "Query: " prefix in 0.3.13 (illuin-tech/colpali#339), and the image document prompt was rewritten in 0.3.9 and again in 0.3.11. The Sentence Transformers configuration in this repository reproduces the original training-time format, so its embeddings differ from current colpali-engine output.

Make sure colpali-engine is installed from source or with a version superior to 0.3.5 (main branch from the repo currently). transformers version must be > 4.46.2.

pip install git+https://github.com/illuin-tech/colpali

import torch
from PIL import Image

from colpali_engine.models import ColIdefics3, ColIdefics3Processor

model = ColIdefics3.from_pretrained(
        "vidore/colSmol-256M",
        torch_dtype=torch.bfloat16,
        device_map="cuda:0",
        attn_implementation="flash_attention_2" # or eager
    ).eval()
processor = ColIdefics3Processor.from_pretrained("vidore/colSmol-256M")

# Your inputs
images = [
    Image.new("RGB", (32, 32), color="white"),
    Image.new("RGB", (16, 16), color="black"),
]
queries = [
    "Is attention really all you need?",
    "What is the amount of bananas farmed in Salvador?",
]

# Process the inputs
batch_images = processor.process_images(images).to(model.device)
batch_queries = processor.process_queries(queries).to(model.device)

# Forward pass
with torch.no_grad():
    image_embeddings = model(**batch_images)
    query_embeddings = model(**batch_queries)

scores = processor.score_multi_vector(query_embeddings, image_embeddings)

Limitations

• Focus: The model primarily focuses on PDF-type documents and high-ressources languages, potentially limiting its generalization to other document types or less represented languages.
• Support: The model relies on multi-vector retreiving derived from the ColBERT late interaction mechanism, which may require engineering efforts to adapt to widely used vector retrieval frameworks that lack native multi-vector support.

License

ColQwen2's vision language backbone model (Qwen2-VL) is under apache2.0 license. The adapters attached to the model are under MIT license.

Contact

• Manuel Faysse: [email protected]
• Hugues Sibille: [email protected]
• Tony Wu: [email protected]

Citation

If you use any datasets or models from this organization in your research, please cite the original dataset as follows:

@misc{faysse2024colpaliefficientdocumentretrieval,
  title={ColPali: Efficient Document Retrieval with Vision Language Models}, 
  author={Manuel Faysse and Hugues Sibille and Tony Wu and Bilel Omrani and Gautier Viaud and Céline Hudelot and Pierre Colombo},
  year={2024},
  eprint={2407.01449},
  archivePrefix={arXiv},
  primaryClass={cs.IR},
  url={https://arxiv.org/abs/2407.01449}, 
}