Colqwenstella 2b Multilingual

🚀 ColQwenStella-2b-multilingual: Qwen2 Visionとstella_en_1.5B_v5モデルを組み合わせた多言語ビジュアルリトリーバ

このモデルは、Qwen2モデルのビジョンコンポーネントとstella_en_1.5B_v5を埋め込みモデルとして組み合わせたアーキテクチャを持ち、多言語のビジュアルドキュメント検索を効率的に行うことができます。

🚀 クイックスタート

2025年2月11日現在、このモデルはパラメータが2B以下のモデルの中で1位、Vidoreベンチマーク全体で8位にランクインしています。報告されているスコアはcheckpoint-1800に対応しています。

これは、4xA100 80GBでper_device_batch_size=128、5エポックでトレーニングされたベースバージョンです。

ColQwenStella-2b-multilingualのアーキテクチャは、Qwen2モデルのビジョンコンポーネントとstella_en_1.5B_v5を埋め込みモデルとして組み合わせています。トレーニングはColPali: Efficient Document Retrieval with Vision Language Modelsのレシピに従って行われています。

✨ 主な機能

• 多言語対応のビジュアルドキュメント検索
• Qwen2モデルのビジョンコンポーネントとstella_en_1.5B_v5を組み合わせたアーキテクチャ

📦 インストール

pip install transformers>=4.46.3

💻 使用例

基本的な使用法

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)

📚 ドキュメント

データ

• 合成データ：openbmb/VisRAG-Ret-Train-Synthetic-dataデータセットから選択され、前処理されました。
• ドメイン内VQAデータセット：openbmb/VisRAG-Ret-Train-In-domain-dataから取得されました。
• Docmatixデータセット：Metric-AI/rag_docmatix_100kデータセットから抽出されました。
• Colpaliデータセット：vidore/colpali_train_setから取得されました。
• 多言語データセット：llamaindex/vdr-multilingual-trainから取得されました。

モデルトレーニング

パラメータ

モデルのトレーニングには、低ランクアダプター(LoRA)を使用しています。言語モデルのトランスフォーマーレイヤー、vison_model.mergerのmlpレイヤー、および最後のランダムに初期化された投影レイヤーでalpha=128、r=128を使用し、adamwオプティマイザーを使用しています。

4xA100 GPUセットアップで分散データ並列（accelerate経由）を使用し、学習率5e-4、コサイン減衰、100ウォームアップステップ、デバイスごとのバッチサイズ128、bfloat16形式でトレーニングしています。

📄 ライセンス

モデルに付属するアダプターはMITライセンスの下にあります。

• 開発元: Metric AI Research Lab