Colqwenstella 2b Multilingual

🚀 ColQwenStella-2b-multilingual: Multilingual Visual Retriever

ColQwenStella-2b-multilingual is a multilingual visual retriever that combines the Vision component of the Qwen2 model with stella_en_1.5B_v5 as its embedding model. It offers efficient and accurate visual document retrieval across multiple languages.

🚀 Quick Start

This model ranked #1 among models <= 2B parameters and #8 overall on the Vidore benchmark (as of February 11, 2025). The reported scores on the Vidore Leaderboard correspond to checkpoint-1800. It's the base version trained on 4xA100 80GB with per_device_batch_size = 128 for 5 epochs.

✨ Features

• Combines the Vision component of the Qwen2 model with stella_en_1.5B_v5 as the embedding model.
• Supports multiple languages including English, French, Spanish, Italian, and German.
• Achieves high performance on the Vidore benchmark.

📦 Installation

pip install transformers>=4.46.3

💻 Usage Examples

Basic Usage

import torch
from PIL import Image

from transformers import AutoModel, AutoProcessor

model = AutoModel.from_pretrained(
        "Metric-AI/ColQwenStella-2b-multilingual",
        torch_dtype=torch.bfloat16,
        device_map="cuda:0",  # or "mps" if on Apple Silicon
        trust_remote_code=True
    ).eval()
processor = AutoProcessor.from_pretrained("Metric-AI/ColQwenStella-2b-multilingual", trust_remote_code=True)

# 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)

📚 Documentation

Data

• Synthetic data: Selected and preprocessed from the openbmb/VisRAG-Ret-Train-Synthetic-data dataset.
• In-domain VQA dataset: Drawn from openbmb/VisRAG-Ret-Train-In-domain-data.
• Docmatix dataset: Extracted from the Metric-AI/rag_docmatix_100k dataset.
• Colpali dataset: Taken from vidore/colpali_train_set.
• Multilingual dataset: Taken from llamaindex/vdr-multilingual-train.

Model Training

Parameters

We train models using low-rank adapters (LoRA) with alpha = 128 and r = 128 on the transformer layers from the language model, and mlp layers of the vison_model.merger as well as the final randomly initialized projection layer, and use an adamw optimizer. We train on a 4xA100 GPU setup with distributed data parallelism (via accelerate), a learning rate of 5e - 4 with cosine decay with 100 warmup steps, batch size per device is 128, in bfloat16 format.

📄 License

The adapters attached to the model are under MIT license.

• Developed by: Metric AI Research Lab