基础模型: answerdotai/ModernBERT-base
数据集:

• lightonai/ms-marco-en-bge
  语言:
• en
  库名称: PyLate
  管道标签: 句子相似度
  标签:
• ColBERT
• PyLate
• sentence-transformers
• 句子相似度
• 特征提取
• 训练器生成
• 数据集大小:808728
• 损失函数:蒸馏

基于answerdotai/ModernBERT-base的PyLate模型

这是一个从answerdotai/ModernBERT-base在train数据集上微调的PyLate模型。它将句子和段落映射为128维密集向量序列，可使用MaxSim算子进行语义文本相似度计算。

模型详情

模型描述

• 模型类型: PyLate模型
• 基础模型: answerdotai/ModernBERT-base
• 文档长度: 180个token
• 查询长度: 32个token
• 输出维度: 128个token
• 相似度函数: MaxSim
• 训练数据集:
  • train
• 语言: en

模型来源

• 文档: PyLate文档
• 代码库: GitHub上的PyLate
• Hugging Face: Hugging Face上的PyLate模型

完整模型架构

ColBERT(
  (0): Transformer({'max_seq_length': 179, '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'})
)

使用方法

首先安装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=pylate_model_id,
)

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

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

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",
)

检索查询的top-k文档

文档索引后，您可以检索给定查询集的最相关top-k文档。为此，使用您要搜索的索引初始化ColBERT检索器，编码查询然后检索top-k文档以获取最匹配的ID和相关分数:

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

# 步骤2: 编码查询
queries_embeddings = model.encode(
    ["查询文档3", "查询文档1"],
    batch_size=32,
    is_query=True,  # 确保设置为True表示这些是查询
    show_progress_bar=True,
)

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

重排序

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

from pylate import rank, models

queries = [
    "查询A",
    "查询B",
]

documents = [
    ["文档A", "文档B"],
    ["文档1", "文档C", "文档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,
)

训练详情

训练数据集

train

• 数据集: train at 11e6ffa
• 大小: 808,728个训练样本
• 列: query_id, document_ids, 和 scores
• 基于前1000个样本的近似统计:

  	query_id	document_ids	scores
  类型	string	list	list
  详情	最小: 5个token 平均: 5.59个token 最大: 6个token	大小: 32个元素	大小: 32个元素

• 样本:

  query_id	document_ids	scores
  121352	['2259784', '4923159', '40211', '1545154', '8527175', ...]	[0.2343463897705078, 0.639204204082489, 0.3806908428668976, 0.5623092651367188, 0.8051995635032654, ...]
  634306	['7723525', '1874779', '379307', '2738583', '7599583', ...]	[0.7124203443527222, 0.7379189729690552, 0.5786551237106323, 0.6142299175262451, 0.6755089163780212, ...]
  920825	['5976297', '2866112', '3560294', '3285659', '4706740', ...]	[0.6462352871894836, 0.7880821228027344, 0.791019856929779, 0.7709633111953735, 0.8284491300582886, ...]

• 损失函数: pylate.losses.distillation.Distillation

评估结果

多向量检索模型的nDCG@10分数

训练超参数

非默认超参数

• per_device_train_batch_size: 16
• learning_rate: 8e-05
• num_train_epochs: 1
• warmup_ratio: 0.05
• bf16: True

所有超参数