ColBERT 如何帮助开发人员克服 RAG 的限制

检索增强生成 (RAG) 现已成为生成人工智能 (AI) 应用的标准部分。使用从矢量数据库检索到的相关上下文来补充您的应用程序提示可以显着提高准确性并减少幻觉。这意味着矢量搜索结果中相关性的增加与 RAG 应用程序的质量直接相关。

即使大型语言模型 (LLM) 增加了上下文窗口， RAG 仍然流行且日益相关的原因有两个：

1. LLM 响应时间和价格均随上下文长度线性增加。

2. 法学硕士仍然在大规模上下文中的检索和推理方面遇到困难。

   
   

但RAG并不是一根魔杖。特别是，最常见的设计，即密集段落检索 (DPR)，将查询和段落表示为单个嵌入向量，并使用简单的余弦相似度来对相关性进行评分。这意味着 DPR 在很大程度上依赖于具有广泛训练范围的嵌入模型来识别所有相关搜索词。

不幸的是，现成的模型很难处理不常见的术语，包括名称，这些术语在训练数据中并不常见。 DPR 还往往对分块策略高度敏感，如果相关段落被大量不相关信息包围，可能会导致相关段落被遗漏。所有这些都给应用程序开发人员带来了“第一次就把事情做好”的负担，因为错误通常会导致需要从头开始重建索引。

使用 ColBERT 解决 DPR 的挑战

ColBERT是一种使用BERT语言模型对段落相关性进行评分的新方法，基本上解决了 DPR 的问题。第一篇 ColBERT 论文中的下图展示了它为何如此令人兴奋：

这将 ColBERT 的性能与 MS-MARCO 数据集的其他最先进的解决方案进行了比较。 （MS-MARCO 是一组 Bing 查询，微软手动对最相关的段落进行评分。它是更好的检索基准之一。）越低越好。

简而言之，ColBERT 轻松胜过大多数更为复杂的解决方案领域，但代价是延迟略有增加。

为了测试这一点，我创建了一个演示，并使用 ada002 DPR 和 ColBERT 对 1,000 多篇维基百科文章建立了索引。我发现 ColBERT 对于不常见的搜索词可提供明显更好的结果。

以下屏幕截图显示，DPR 无法识别亚伯拉罕·林肯的助手 William H. Herndon 的不寻常名字，而 ColBERT 在斯普林菲尔德的文章中找到了参考文献。另请注意，ColBERT 的第二名结果是针对不同的 William，而 DPR 的结果均不相关。

ColBERT 通常用密集的机器学习术语来描述，但实际上非常简单。我将展示如何仅使用几行 Python 和Cassandra 查询语言(CQL) 在DataStax Astra DB上实现 ColBERT 检索和评分。

大创意

ColBERT 不是传统的基于单向量的 DPR 将段落转换为单个“嵌入”向量，而是为段落中的每个标记生成一个受上下文影响的向量。 ColBERT 类似地为查询中的每个标记生成向量。

（标记化是指在法学硕士处理之前将输入分解为单词的片段。OpenAI团队的创始成员 Andrej Karpathy 刚刚发布了一段精彩视频，介绍了其工作原理。）

然后，每个文档的得分是每个查询嵌入与任何文档嵌入的最大相似度之和：

 def maxsim(qv, document_embeddings): return max(qv @ dv for dv in document_embeddings) def score(query_embeddings, document_embeddings): return sum(maxsim(qv, document_embeddings) for qv in query_embeddings)

（@ 是 PyTorch 点积运算符，是向量相似度的最常见度量。）

就是这样——您可以用四行 Python 代码实现 ColBERT 评分！现在，您比 X（以前称为 Twitter）上发布有关 ColBERT 的 99% 的人更了解 ColBERT。

ColBERT 论文的其余部分涉及：

1. 如何微调 BERT 模型以生成给定数据集的最佳嵌入？
2. 如何限制计算此处显示的（相对昂贵的）分数的文档集？

第一个问题是可选的，超出了本文的范围。我将使用预训练的 ColBERT 检查点。但第二个方法很简单，可以使用 DataStax Astra DB 等矢量数据库来完成。

Astra DB 上的 ColBERT

ColBERT 有一个流行的 Python 一体化库，名为RAGatouille ；然而，它假设一个静态数据集。 RAG 应用程序的强大功能之一是实时响应动态变化的数据。因此，我将使用 Astra 的向量索引来缩小我需要评分的文档集，以找到每个子向量的最佳候选者。

将 ColBERT 添加到 RAG 应用程序有两个步骤：摄取和检索。

食入

因为每个文档块都有多个与之关联的嵌入，所以我需要两个表：

 CREATE TABLE chunks ( title text, part int, body text, PRIMARY KEY (title, part) ); CREATE TABLE colbert_embeddings ( title text, part int, embedding_id int, bert_embedding vector<float, 128>, PRIMARY KEY (title, part, embedding_id) ); CREATE INDEX colbert_ann ON colbert_embeddings(bert_embedding) WITH OPTIONS = { 'similarity_function': 'DOT_PRODUCT' };

安装 ColBERT 库 ( pip install colbert-ai ) 并下载预训练的 BERT 检查点后，我可以将文档加载到这些表中：

 from colbert.infra.config import ColBERTConfig from colbert.modeling.checkpoint import Checkpoint from colbert.indexing.collection_encoder import CollectionEncoder from cassandra.concurrent import execute_concurrent_with_args from db import DB def encode_and_save(title, passages): db = DB() cf = ColBERTConfig(checkpoint='checkpoints/colbertv2.0') cp = Checkpoint(cf.checkpoint, colbert_config=cf) encoder = CollectionEncoder(cf, cp) # encode_passages returns a flat list of embeddings and a list of how many correspond to each passage embeddings_flat, counts = encoder.encode_passages(passages) # split up embeddings_flat into a nested list start_indices = [0] + list(itertools.accumulate(counts[:-1])) embeddings_by_part = [embeddings_flat[start:start+count] for start, count in zip(start_indices, counts)] # insert into the database for part, embeddings in enumerate(embeddings_by_part): execute_concurrent_with_args(db.session, db.insert_colbert_stmt, [(title, part, i, e) for i, e in enumerate(embeddings)])

（我喜欢将数据库逻辑封装在专用模块中；您可以在我的GitHub 存储库中访问完整源代码。）

恢复

然后检索看起来像这样：

 def retrieve_colbert(query): db = DB() cf = ColBERTConfig(checkpoint='checkpoints/colbertv2.0') cp = Checkpoint(cf.checkpoint, colbert_config=cf) encode = lambda q: cp.queryFromText([q])[0] query_encodings = encode(query) # find the most relevant documents for each query embedding. using a set # handles duplicates so we don't retrieve the same one more than once docparts = set() for qv in query_encodings: rows = db.session.execute(db.query_colbert_ann_stmt, [list(qv)]) docparts.update((row.title, row.part) for row in rows) # retrieve these relevant documents and score each one scores = {} for title, part in docparts: rows = db.session.execute(db.query_colbert_parts_stmt, [title, part]) embeddings_for_part = [tensor(row.bert_embedding) for row in rows] scores[(title, part)] = score(query_encodings, embeddings_for_part) # return the source chunk for the top 5 return sorted(scores, key=scores.get, reverse=True)[:5]

以下是针对最相关文档部分（ db.query_colbert_ann_stmt ）执行的查询：

 SELECT title, part FROM colbert_embeddings ORDER BY bert_embedding ANN OF ? LIMIT 5

超越基础：RAGStack

本文和链接的存储库简要介绍了 ColBERT 的工作原理。您今天就可以使用自己的数据实施此操作并立即看到结果。与人工智能领域的一切一样，最佳实践每天都在变化，新技术也在不断涌现。

为了更轻松地跟上最先进的技术，DataStax 正在将此功能和其他增强功能引入RAGStack ，这是我们利用 LangChain 和 LlamaIndex 的生产就绪型 RAG 库。我们的目标是为开发人员提供一致的 RAG 应用程序库，使他们能够控制新功能的升级。您不必跟上技术和库的无数变化，而是拥有单一流，因此您可以专注于构建应用程序。您现在就可以使用 RAGStack 来整合 LangChain 和 LlamaIndex 的最佳实践，开箱即用；像 ColBERT 这样的进步将在即将发布的版本中出现在 RAGstack 中。

作者：DataStax 的 Jonathan Ellis