🚀 GTE-ModernColBERT-v1

这是一个基于 Alibaba-NLP/gte-modernbert-base 的 PyLate 模型，可将句子和段落映射为128维的密集向量序列，适用于使用 MaxSim 算子进行语义文本相似度计算。

🚀 快速开始

首先安装 PyLate 库：

pip install -U pylate

检索

PyLate 提供了一个简化的接口，用于使用 ColBERT 模型对文档进行索引和检索。索引利用 Voyager HNSW 索引来高效处理文档嵌入，并实现快速检索。

文档索引

首先，加载 ColBERT 模型并初始化 Voyager 索引，然后对文档进行编码和索引：

from pylate import indexes, models, retrieve

# 步骤 1: 加载 ColBERT 模型
model = models.ColBERT(
    model_name_or_path=lightonai/GTE-ModernColBERT-v1,
)

# 步骤 2: 初始化 Voyager 索引
index = indexes.Voyager(
    index_folder="pylate-index",
    index_name="index",
    override=True,  # 如果存在现有索引，则覆盖它
)

# 步骤 3: 编码文档
documents_ids = ["1", "2", "3"]
documents = ["document 1 text", "document 2 text", "document 3 text"]

documents_embeddings = model.encode(
    documents,
    batch_size=32,
    is_query=False,  # 确保将其设置为 False 以表明这些是文档，而不是查询
    show_progress_bar=True,
)

# 步骤 4: 通过提供嵌入和相应的 ID 将文档嵌入添加到索引中
index.add_documents(
    documents_ids=documents_ids,
    documents_embeddings=documents_embeddings,
)

请注意，您不必每次都重新创建索引和编码文档。一旦创建了索引并添加了文档，您可以通过加载它来重复使用该索引：

# 要加载索引，只需使用正确的文件夹/名称实例化它，而不覆盖它
index = indexes.Voyager(
    index_folder="pylate-index",
    index_name="index",
)

查询前 k 个文档

一旦文档被索引，您可以为给定的查询集检索前 k 个最相关的文档。 为此，使用要搜索的索引初始化 ColBERT 检索器，对查询进行编码，然后检索前 k 个文档以获取前匹配项的 ID 和相关性分数：

# 步骤 1: 初始化 ColBERT 检索器
retriever = retrieve.ColBERT(index=index)

# 步骤 2: 编码查询
queries_embeddings = model.encode(
    ["query for document 3", "query for document 1"],
    batch_size=32,
    is_query=True,  # 确保将其设置为 True 以表明这些是查询
    show_progress_bar=True,
)

# 步骤 3: 检索前 k 个文档
scores = retriever.retrieve(
    queries_embeddings=queries_embeddings, 
    k=10,  # 为每个查询检索前 10 个匹配项
)

重排序

如果您只想使用 ColBERT 模型在第一阶段检索管道的基础上进行重排序，而不构建索引，您可以简单地使用 rank 函数并传递要重排序的查询和文档：

from pylate import rank, models

queries = [
    "query A",
    "query B",
]

documents = [
    ["document A", "document B"],
    ["document 1", "document C", "document B"],
]

documents_ids = [
    [1, 2],
    [1, 3, 2],
]

model = models.ColBERT(
    model_name_or_path=pylate_model_id,
)

queries_embeddings = model.encode(
    queries,
    is_query=True,
)

documents_embeddings = model.encode(
    documents,
    is_query=False,
)

reranked_documents = rank.rerank(
    documents_ids=documents_ids,
    queries_embeddings=queries_embeddings,
    documents_embeddings=documents_embeddings,
)

✨ 主要特性

• 语义相似度计算：能够将句子和段落映射为128维的密集向量序列，使用 MaxSim 算子进行语义文本相似度计算。
• 高效检索：利用 Voyager HNSW 索引，实现快速的文档检索。
• 长上下文处理：在长上下文嵌入基准测试中表现出色，能够处理超出训练长度的文档。

📦 安装指南

pip install -U pylate

💻 使用示例

基础用法

# 加载 ColBERT 模型
from pylate import models
model = models.ColBERT(
    model_name_or_path=lightonai/GTE-ModernColBERT-v1,
)

# 编码文档
documents = ["document 1 text", "document 2 text"]
documents_embeddings = model.encode(
    documents,
    is_query=False,
)

# 编码查询
queries = ["query for document 1"]
queries_embeddings = model.encode(
    queries,
    is_query=True,
)

高级用法

# 文档索引和检索
from pylate import indexes, retrieve

# 初始化 Voyager 索引
index = indexes.Voyager(
    index_folder="pylate-index",
    index_name="index",
    override=True,
)

# 添加文档嵌入到索引
documents_ids = ["1", "2"]
index.add_documents(
    documents_ids=documents_ids,
    documents_embeddings=documents_embeddings,
)

# 初始化 ColBERT 检索器
retriever = retrieve.ColBERT(index=index)

# 检索前 10 个文档
scores = retriever.retrieve(
    queries_embeddings=queries_embeddings, 
    k=10,
)

📚 详细文档

模型详情

文档长度

GTE-ModernColBERT 在 MS MARCO 上使用知识蒸馏进行训练，文档长度为 300 个标记，这解释了其文档长度的默认值。 然而，正如 ModernBERT 论文所示，ColBERT 模型可以推广到远远超出其训练长度的文档长度，并且 GTE-ModernColBERT 在长上下文嵌入基准测试中实际上取得了远远高于当前最优水平的结果，请参阅 LongEmbed 结果。 在加载模型时，只需根据需要调整文档长度参数：

model = models.ColBERT(
    model_name_or_path=lightonai/GTE-ModernColBERT-v1,
    document_length=8192,
)

ModernBERT 本身仅在 8K 上下文长度上进行了训练，但似乎 GTE-ModernColBERT 可以推广到更大的上下文大小，不过这并不保证，因此请自行进行测试！

模型来源

• 文档：PyLate 文档
• 仓库：GitHub 上的 PyLate
• Hugging Face：Hugging Face 上的 PyLate 模型

完整模型架构

ColBERT(
  (0): Transformer({'max_seq_length': 299, 'do_lower_case': False}) with Transformer model: ModernBertModel 
  (1): Dense({'in_features': 768, 'out_features': 128, 'bias': False, 'activation_function': 'torch.nn.modules.linear.Identity'})
)

🔧 技术细节

训练超参数

非默认超参数

• eval_strategy: steps
• per_device_train_batch_size: 16
• learning_rate: 3e-05
• bf16: True

所有超参数

训练日志