Modernbert Base ColBERT

🚀 基於answerdotai/ModernBERT-base的PyLate模型

本模型基於answerdotai/ModernBERT-base，在訓練數據集上進行微調。它能將句子和段落映射為128維的密集向量序列，可通過MaxSim算子用於語義文本相似度計算。

🚀 快速開始

本模型是基於answerdotai/ModernBERT-base微調的PyLate模型，可用於語義文本相似度任務。以下是使用該模型的基本步驟：

1. 安裝PyLate庫。
2. 加載模型並進行文檔索引和檢索。
3. 可根據需求進行重排序操作。

✨ 主要特性

• 高效映射：能夠將句子和段落映射為128維的密集向量，便於進行語義文本相似度計算。
• 多場景應用：支持文檔索引、檢索和重排序等多種應用場景。
• 集成Voyager索引：在檢索過程中，利用Voyager HNSW索引高效處理文檔嵌入，實現快速檢索。

📦 安裝指南

首先，安裝PyLate庫：

pip install -U pylate

💻 使用示例

基礎用法

文檔索引

from pylate import indexes, models, retrieve

# Step 1: Load the ColBERT model
model = models.ColBERT(
    model_name_or_path=pylate_model_id,
)

# Step 2: Initialize the Voyager index
index = indexes.Voyager(
    index_folder="pylate-index",
    index_name="index",
    override=True,  # This overwrites the existing index if any
)

# Step 3: Encode the documents
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,  # Ensure that it is set to False to indicate that these are documents, not queries
    show_progress_bar=True,
)

# Step 4: Add document embeddings to the index by providing embeddings and corresponding ids
index.add_documents(
    documents_ids=documents_ids,
    documents_embeddings=documents_embeddings,
)

注意，無需每次都重新創建索引和編碼文檔。創建索引並添加文檔後，可通過以下方式加載並重復使用索引：

# To load an index, simply instantiate it with the correct folder/name and without overriding it
index = indexes.Voyager(
    index_folder="pylate-index",
    index_name="index",
)

檢索前k個相關文檔

# Step 1: Initialize the ColBERT retriever
retriever = retrieve.ColBERT(index=index)

# Step 2: Encode the queries
queries_embeddings = model.encode(
    ["query for document 3", "query for document 1"],
    batch_size=32,
    is_query=True,  #  # Ensure that it is set to False to indicate that these are queries
    show_progress_bar=True,
)

# Step 3: Retrieve top-k documents
scores = retriever.retrieve(
    queries_embeddings=queries_embeddings, 
    k=10,  # Retrieve the top 10 matches for each query
)

高級用法

重排序

如果僅想使用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,
)

📚 詳細文檔

模型詳情

模型描述

模型來源

• 文檔：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'})
)

訓練詳情

訓練數據集

訓練集

• 數據集：train，版本為11e6ffa
• 大小：808,728個訓練樣本
• 列：query_id、document_ids和scores
• 基於前1000個樣本的近似統計信息：

  	query_id	document_ids	scores
  類型	字符串	列表	列表
  詳情	最小：5個詞元 平均：5.59個詞元 最大：6個詞元	大小：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

所有超參數

框架版本

• Python：3.11.9
• Sentence Transformers：3.3.0
• PyLate：1.1.4
• Transformers：4.48.0.dev0
• PyTorch：2.4.0
• Accelerate：1.2.1
• Datasets：2.21.0
• Tokenizers：0.21.0

📄 許可證

文檔未提及相關信息。

🔧 技術細節

本模型基於answerdotai/ModernBERT-base進行微調，使用ColBERT架構。在訓練過程中，利用Voyager HNSW索引對文檔嵌入進行高效處理，以實現快速檢索。模型將句子和段落映射為128維的密集向量，通過MaxSim算子計算語義文本相似度。訓練數據集為lightonai/ms-marco-en-bge，包含808,728個樣本。訓練過程中使用了Distillation損失函數，並設置了一系列超參數，如per_device_train_batch_size為16，learning_rate為8e-05等。

📄 引用

BibTeX

Sentence Transformers

@inproceedings{reimers-2019-sentence-bert,
    title = "Sentence-BERT: Sentence Embeddings using Siamese BERT-Networks",
    author = "Reimers, Nils and Gurevych, Iryna",
    booktitle = "Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing",
    month = "11",
    year = "2019",
    publisher = "Association for Computational Linguistics",
    url = "https://arxiv.org/abs/1908.10084"
}

PyLate

@misc{PyLate,
title={PyLate: Flexible Training and Retrieval for Late Interaction Models},
author={Chaffin, Antoine and Sourty, Raphaël},
url={https://github.com/lightonai/pylate},
year={2024}
}