Colqwenstella 2b Multilingual

🚀 ColQwenStella-2b-多语言版：基于Qwen2视觉模型和stella_en_1.5B_v5模型组合的多语言视觉检索器

ColQwenStella-2b-多语言版模型结合了Qwen2模型的视觉组件和stella_en_1.5B_v5作为嵌入模型。它在Vidore基准测试中表现出色，在参数不超过20亿的模型中排名第1，总体排名第8。该模型按照特定的训练方法进行训练，为视觉文档检索任务提供了强大的支持。

🚀 快速开始

ColQwenStella-2b-多语言版在Vidore基准测试（截至2025年2月11日）中，参数不超过20亿的模型里排名第1，总体排名第8。在Vidore排行榜上报告的分数对应于检查点1800。这是在4xA100 80GB上训练的基础版本，每个设备的批量大小为128，训练了5个周期。

✨ 主要特性

• 架构创新：将Qwen2模型的视觉组件与stella_en_1.5B_v5作为嵌入模型相结合。
• 多语言支持：支持英语、法语、西班牙语、意大利语和德语等多种语言。
• 优秀性能：在Vidore基准测试中取得了优异的成绩。

📦 安装指南

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）训练模型，在语言模型的Transformer层、vison_model.merger的mlp层以及最终随机初始化的投影层上设置alpha=128和r=128，并使用adamw优化器。我们在4xA100 GPU设置上进行训练，采用分布式数据并行（通过accelerate），学习率为5e-4，采用余弦衰减，有100个热身步骤，每个设备的批量大小为128，使用bfloat16格式。

📄 许可证

模型附带的适配器遵循MIT许可证。

开发信息

• 开发者：Metric AI Research Lab