neural-cherche - 专注于神经搜索模型微调和快速推理的工具库

Neural-Cherche 项目介绍

Neural-Cherche 是一个专为精调神经搜索模型而设计的库，主要用于模型如 Splade、ColBERT 和 SparseEmbed。在处理特定数据集时，它能够提供高效的推理方法，并支持对精调后的检索器或排序器进行操作。Neural-Cherche 的设计目标是提供一种简便有效的方法，让用户能够在离线和在线环境中利用神经搜索模型。此外，它还支持保存所有计算的嵌入，以避免冗余计算。

兼容性

Neural-Cherche 可以在多种设备上运行，包括 CPU、GPU 和 MPS 等。在精调过程中，用户可以用任何预训练的 Sentence Transformer 检查点来精调 ColBERT，而对于 Splade 和 SparseEmbed 的精调则需要一个基于语言模型 (MLM) 的预训练模型。

安装指南

要安装 Neural-Cherche，可以使用以下命令：

pip install neural-cherche

如果您计划在训练期间对模型进行评价，请使用：

pip install "neural-cherche[eval]"

快速开始

训练数据集需要由三元组 (anchor, positive, negative) 组成，其中 anchor 是查询，positive 是与 anchor 直接相关的文档，negative 是与 anchor 无关的文档。例如：

X = [
    ("anchor 1", "positive 1", "negative 1"),
    ("anchor 2", "positive 2", "negative 2"),
    ("anchor 3", "positive 3", "negative 3"),
]

以下是使用 Neural-Cherche 从 Sentence Transformer 预训练检查点对 ColBERT 模型进行精调的示例：

import torch
from neural_cherche import models, utils, train

model = models.ColBERT(
    model_name_or_path="raphaelsty/neural-cherche-colbert",
    device="cuda" if torch.cuda.is_available() else "cpu"  # 或 mps
)

optimizer = torch.optim.AdamW(model.parameters(), lr=3e-6)

X = [
    ("query", "positive document", "negative document"),
    ("query", "positive document", "negative document"),
    ("query", "positive document", "negative document"),
]

for step, (anchor, positive, negative) in enumerate(utils.iter(
        X,
        epochs=1,  # 训练轮数
        batch_size=8,  # 每批次三元组数量
        shuffle=True
    )):

    loss = train.train_colbert(
        model=model,
        optimizer=optimizer,
        anchor=anchor,
        positive=positive,
        negative=negative,
        step=step,
        gradient_accumulation_steps=50,
    )

    if (step + 1) % 1000 == 0:
        model.save_pretrained("checkpoint")

文档检索

使用精调后的 ColBERT 模型可以对文档进行重新排序。例如：

import torch
from lenlp import sparse
from neural_cherche import models, rank, retrieve

documents = [
    {"id": "doc1", "title": "Paris", "text": "Paris is the capital of France."},
    {"id": "doc2", "title": "Montreal", "text": "Montreal is the largest city in Quebec."},
    {"id": "doc3", "title": "Bordeaux", "text": "Bordeaux in Southwestern France."},
]

retriever = retrieve.BM25(
    key="id",
    on=["title", "text"],
    count_vectorizer=sparse.CountVectorizer(
        normalize=True, ngram_range=(3, 5), analyzer="char_wb", stop_words=[]
    ),
    k1=1.5,
    b=0.75,
    epsilon=0.0,
)

model = models.ColBERT(
    model_name_or_path="raphaelsty/neural-cherche-colbert",
    device="cuda" if torch.cuda.is_available() else "cpu"  # 或 mps
)

ranker = rank.ColBERT(
    key="id",
    on=["title", "text"],
    model=model,
)

documents_embeddings = retriever.encode_documents(documents=documents)

retriever.add(documents_embeddings=documents_embeddings)

queries = ["Paris", "Montreal", "Bordeaux"]
queries_embeddings = retriever.encode_queries(queries=queries)
ranker_queries_embeddings = ranker.encode_queries(queries=queries)

candidates = retriever(
    queries_embeddings=queries_embeddings,
    batch_size=32,
    k=100,  # 检索文档数量
)

ranker_documents_embeddings = ranker.encode_candidates_documents(
    candidates=candidates,
    documents=documents,
    batch_size=32,
)

scores = ranker(
    queries_embeddings=ranker_queries_embeddings,
    documents_embeddings=ranker_documents_embeddings,
    documents=candidates,
    batch_size=32,
)

scores

预训练模型

Neural-Cherche 提供了专门为其设计的预训练检查点，例如 [raphaelsty/neural-cherche-sparse-embed] 和 [raphaelsty/neural-cherche-colbert]。这些检查点在 MS-MARCO 数据集的一个子集上进行了精调，可以进一步在特定数据集上进行精调以更好地适应特定语言。

项目贡献者

Benjamin Clavié
Arthur Satouf

参考文献

Neural-Cherche 项目使用了一些知名的学术论文作为其理论基础，包括 SPLADE、SparseEmbed、ColBERT 等项目。

许可证

本项目基于 MIT 开源许可证发布，SPLADE 模型仅可用于非商业目的，而 SparseEmbed 和 ColBERT 完全开源，包括商业用途。