GTE ModernColBERT V1

🚀 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

所有超參數

訓練日誌