使用 DuckDB + Ibis 进行 RAG
原文：Using DuckDB + Ibis for RAG[1]
翻译：Gemini
校对：alitrack
author: Cody Peterson
date: 2024-02-22
前言
昨天刚讨论过 【DuckDB 也要加入向量数据库的战局?】， 今天便发现了这篇文章，严格说，这篇文章可以和 DuckDB 的向量数据库能力无关，但仍不失是一篇不错的文章，另外我没有使用 OpenAI 的 embedding，而是使用了 ollama 的 embedding 这里把我的代码附上。
import ollama
def _embed(text: str) -> list[float]:
"""Text to fixed-length array embedding."""
model = "nomic-embed-text"#"text-embedding-3-small"
text = text.replace("\n", " ")
# client = OpenAI()
try:
return (
# client.embeddings.create(input=[text], model=model).data[0].embedding
ollama.embeddings(model=model, prompt=text)['embedding']
)
except Exception as e:
print(e)
return None
@ibis.udf.scalar.python
def embed(text: str, tokens_estimate: float) -> list[float]:
"""Text to fixed-length array embedding."""
if 0 < tokens_estimate < 8191:
return _embed(text)
return None
💡：
使用 ollama 替代了 OpenAI，速度上去了，成本下来了，对后面的结果排序略有不同，有兴趣的自己测试
另外代码有 bug，
def embed(text: str, tokens_estimate: int) -> list[float]:
改为
def embed(text: str, tokens_estimate: float) -> list[float]:
否则报错，原因不明（tokens_estimate 按说是 int64 数据类型）
概述
在这篇文章中，我们将通过 Ibis 演示传统的检索增强生成 (RAG) ， DuckDB 和 OpenAI，并讨论其优缺点。请注意，由于 Ibis 是可移植的，因此您可以使用 Ibis 支持的（几乎）任何数据库进行 RAG！
💡.我可以使用哪些数据库
数据库必须支持数组类型，并且在数字数组之间具有一定的相似性度量。或者，可以将自定义用户定义函数 (UDF) 用于相似性度量。
如果数据库支持 固定大小数组[2]（如 DuckDB 最近在 0.10.0 版本中发布），则计算相似性的性能也会快得多。我们仍在本文中使用 0.9.2，但升级很容易（我用了 0.10.0）。
DuckDB 是 Ibis 的默认后端[3]，并且使本地快速原型设计变得非常容易。您可以通过更改连接字符串将其部署到 MotherDuck[4]，或将 DuckDB 换成另一个 Ibis 后端。
什么是 RAG？
检索增强生成 (RAG) 是一种语言模型系统从知识库中检索相关信息、用该知识增强模型的输入，然后生成响应的技术。它有助于提供有用的上下文，因为语言模型的上下文长度有限，并提高响应的质量。
💡.RAG 是否已死？
随着语言模型上下文长度的增加，RAG 是否已死？
否。尽管正如我们稍后将讨论的，尚不清楚传统的向量相似性搜索 RAG 是否是最佳方法。无论如何，仅仅因为语言模型可以使用教科书的全部内容并不意味着您会获得比首先检索教科书的相关子集更好的性能（准确性或效率）。
传统的 RAG 将文本嵌入到数字向量中，然后计算输入文本与知识库之间的相似性度量，以确定最相关的上下文。
让我们进入代码！
在 Ibis 文档上进行 RAG
我们将使用 Ibis 文档作为我们的知识库。
设置代码
首先，我们需要安装并导入必要的软件包：
pip install 'ibis-framework[duckdb]' openai python-dotenv
import os # <1>
import ibis
from openai import OpenAI
from dotenv import load_dotenv
from pathlib import Path # <1>
1. 导入我们将使用的库。
我们需要将 OPENAI_API_KEY
环境变量设置为我们的 OpenAI API 密钥。通过在代码所在的目录中创建一个 .env
文件并添加以下行来执行此操作：
OPENAI_API_KEY=<your-api-key> # <1>
1. 替换为您的 OpenAI API 密钥。
然后加载环境变量：
load_dotenv() # <1>
1. 从
.env
文件加载环境变量。
现在我们准备调用 OpenAI。
我们还为交互式使用设置 Ibis：
ibis.options.interactive = True # <1>
1. 打开 Ibis 交互模式。
将数据导入 DuckDB
让我们准备将所有 Quarto markdown 文件导入名为 rag.ddb
的 DuckDB 数据库：
rag_con = ibis.connect("duckdb://rag.ddb") # <1>
for table in rag_con.list_tables(): # <2>
rag_con.drop_table(table) # <2>
rag_con.list_tables() # <3>
1. 创建到 DuckDB 数据库的 Ibis 连接。
2. 删除以前运行的任何现有表。
3. 列出数据库中的表以显示没有表。
我们将遍历 Ibis docs/
目录。对于 .qmd
文件，读取内容并将它们插入 ibis
表：
💡 我们在 Ibis 存储库中运行此博客。要在您自己的数据上运行此操作，请相应地调整代码！
table_name = "ibis" # <1>
for filepath in Path(ibis.__file__).parents[1].joinpath("docs").rglob("*.qmd"): # <2>
contents = filepath.read_text() # <3>
data = { # <4>
"filepath": [str(filepath).split("ibis/")[1]],
"contents": [contents],
} # <4>
t = ibis.memtable(data) # <5>
if "ibis" not in rag_con.list_tables(): # <6>
rag_con.create_table("ibis", t) # <6>
else: # <7>
rag_con.insert("ibis", t) # <7>
for table in rag_con.list_tables(): # <8>
if "memtable" in table: # <8>
rag_con.drop_view(table) # <8>
rag_con.list_tables() # <9>
1. 定义表名。
2. 遍历
docs/
目录并查找所有.qmd
文件。3. 读取每个文件的内容。
4. 使用文件路径和内容创建一个字典。
5. 从字典创建一个 Ibis 内存表。
6. 如果表不存在，则创建它并插入数据。
7. 否则，将数据插入现有表。
8. 删除创建的任何内存表。
9. 列出数据库中的表以显示已创建
ibis
表。
让我们检查一下表：
t = rag_con.table(table_name) # <1>
t
1. 从 DuckDB 数据库获取
ibis
表。
现在我们有一个 DuckDB 数据库，其中包含 ibis
表，该表包含 ibis
存储库中 Quarto markdown 文件的所有文件路径和内容。您可以调整此代码以使其适用于任何存储库和任何文件类型！例如，我们可能还希望导入 Ibis repsoitory 中的所有 .py
文件。
将文本嵌入到向量
要以传统的 RAG 方式搜索相似文本，我们需要将文本嵌入到向量中。我们可以使用 OpenAI 的嵌入模型[5] 来执行此步骤。
请注意，输入标记的最大长度为 8,191。一个单词通常由 1-3 个标记组成。为了估计 Python 中字符串的标记数，我们可以将该字符串的长度除以 4。这并不完美，但现在可以工作。让我们使用估计的标记数来扩充我们的数据并从大到小排序：
t = (
t.mutate(tokens_estimate=t["contents"].length() // 4) # <1>
.order_by(ibis._["tokens_estimate"].desc()) # <2>
.relocate("filepath", "tokens_estimate") # <3>
)
t
1. 使用估计的标记数扩充表。
2. 按估计的标记数从大到小对表进行排序。
3. 将
filepath
和tokens_estimate
列重新定位到表的前面。
最长的文本肯定超过了我们的嵌入模型的标记限制，第二个最长的文本也可能超过了限制。在实践中，我们希望将文本分成更小的块。这是一个挑战，但出于演示目的，我们将忽略太长的文本。
让我们定义我们的嵌入函数：
def _embed(text: str) -> list[float]: # <1>
"""Text to fixed-length array embedding.""" # <1>
model = "text-embedding-3-small" # <2>
text = text.replace("\n", " ") # <3>
client = OpenAI() # <4>
try: # <5>
return ( # <6>
client.embeddings.create(input=[text], model=model).data[0].embedding
) # <6>
except Exception as e: # <7>
print(e) # <7>
return None # <7>
@ibis.udf.scalar.python # <8>
def embed(text: str, tokens_estimate: int) -> list[float]: # <8>
"""Text to fixed-length array embedding.""" # <8>
if 0 < tokens_estimate < 8191: # <9>
return _embed(text) # <10>
return None # <11>
1. 定义一个常规的 Python 函数来嵌入文本。
2. 定义要使用的嵌入模型。
3. 用空格替换换行符。
4. 创建一个 OpenAI 客户端。
5. 由于这是一个外部 API 调用，我们需要处理异常。
6. 调用 OpenAI API 来嵌入文本。
7. 处理异常，如果发生错误，则返回
None
。8. 使用 Python 函数定义一个 Ibis UDF。
9. 如果估计令牌数在限制范围内...
10. ...调用 Python 函数来嵌入文本。
11. 否则，返回
None
。
让我们嵌入文本：
t = t.mutate( # <1>
embedding=embed(t["contents"], t["tokens_estimate"]) # <1>
).cache() # <1>
t
1. 使用嵌入文本扩充表格并缓存结果。
注意有一个错误——第二行的内容确实太长了！我们的标记估计有点偏差，但总体上很有用。我们在上面的 Python UDF 中考虑了这种情况，当文本太长或发生错误时返回 None
（在 DuckDB 中映射为 NULL
）。
使用余弦相似性进行搜索
向量的余弦相似性是衡量 RAG 应用程序相似性的常用方法。我们将在后面探讨它的缺点。我们需要让 Ibis 了解 DuckDB 内置的 list_cosine_similarity
函数：
@ibis.udf.scalar.builtin # <1>
def list_cosine_similarity(x, y) -> float: # <1>
"""Compute cosine similarity between two vectors.""" # <1>
1. 使用 Ibis 内置 UDF 定义余弦相似性函数。
💡 如果使用不同的后端，则需要使用其相应的余弦相似性函数或定义自己的函数。
现在，我们可以在文档中搜索类似的文本：
def search_docs(text): # <1>
"""Search documentation for similar text, returning a sorted table""" # <1>
embedding = _embed(text) # <2>
s = (
t.mutate(similarity=list_cosine_similarity(t["embedding"], embedding)) # <3>
.relocate("similarity") # <4>
.order_by(ibis._["similarity"].desc()) # <5>
.cache() # <6>
)
return s
1. 定义一个函数来搜索文档。
2. 嵌入输入文本。
3. 使用嵌入文本和输入文本之间的余弦相似性扩充表格。
4. 将
similarity
列重新定位到表格的前面。5. 按余弦相似性降序对表格进行排序。
6. 缓存结果。
text = "where can I chat with the community about Ibis?"
search_docs(text)
现在我们已经检索到最相似的文档，我们可以在生成响应之前使用该上下文扩充我们的语言模型的输入！在实践中，我们可能希望设置相似性阈值并采用前 N
个结果。将我们的文本分成更小的部分并从这些结果中进行选择也是一个好主意。
让我们再试几个查询：
text = "what do users say about Ibis?"
search_docs(text)
text = "can Ibis complete the one billion row challenge?"
search_docs(text)
text = "teach me about the Polars backend"
search_docs(text)
text = "why does Voltron Data support Ibis?"
search_docs(text)
text = "why should I use Ibis?"
search_docs(text)
text = "what is the Ibis roadmap?"
search_docs(text)
RAG 缺陷
与传统的软件工程和机器学习一样，在围绕语言模型构建系统时，架构并不是通用的。最佳架构取决于系统的应用。使用余弦（或其他）相似性搜索经过数字转换的文本可能并不总是表现最佳。
我们可以通过以下场景来演示这一点：假设我们正在构建一个聊天机器人，它可以根据用户用英语（或语言模型充分理解的任何语言）提出的问题来编写和运行 SQL。我们可能会将过去的查询、生成的 SQL 以及一些评估标准（用户对生成的 SQL 的赞成/反对票）存储在数据库中。然后，我们可以使用上面相同的 RAG 方法，使用用户的过去查询和生成的 SQL 来扩充我们的查询。
以以下英语查询为例：
a = """the first 10 rows of the penguins, ordered by body_mass_g from
lightest to heaviest"""
我们可以计算此查询与其自身的余弦相似性，结果为 1.0 或接近 1.0：
list_cosine_similarity(_embed(a), _embed(a))
现在，让我们翻转顺序。这将导致完全不同的 SQL 查询和结果集：
b = """the first 10 rows of the penguins, ordered by body_mass_g from
heaviest to lightest"""
并计算 a
和 b
的余弦相似性：
list_cosine_similarity(_embed(a), _embed(b))
同样，我们可以构造一个语义等价的查询，该查询将产生与示例 a
相同的 SQL：
c = """of the birds retrieve the initial ten opposite-of-columns,
sorted from biggest to smallest by weight"""
在嵌入空间中，它的余弦相似性要低得多：
list_cosine_similarity(_embed(a), _embed(c))
尽管语义等价，但这两个查询不会被视为相似。你需要小心并理解系统在做什么。
讨论
传统的 RAG 是最佳方法吗？这还不清楚，而且取决于具体情况。
系统架构
通常，基于 RAG 的系统会引入另一个数据库来存储嵌入。希望这篇文章已经证明传统的数据库选项也可以正常工作。虽然我们这里只演示了几行，但 DuckDB 可以轻松扩展到数百万行。
成本和速度
在 OpenAI 上运行此演示的成本不到 0.01 美元。它也运行大约一分钟，但对于很多行来说，这会慢得多。你需要在这里考虑你的架构。一种选择是使用可以在本地运行的开源嵌入模型。用 OpenAI 调用替换它，看看会发生什么！
向量搜索 RAG 的替代方案
很明显，我们应该始终使用最相关的上下文来扩充我们的语言模型调用，即使传统的向量搜索 RAG 不是最佳解决方案。那么问题就变成了如何为给定系统的预期行为检索该上下文。
💡 废除 RAG 作为术语？
“我想废除术语 RAG，而是同意我们应该始终尝试为模型提供适当的上下文，以提供高质量的答案。” - Hamel Husain[6]
向量搜索 RAG 的一些替代方案包括：
• 传统搜索方法
• 使用外部网络搜索引擎
• 使用语言模型进行搜索
最后一点特别有趣——在上面带有 a
和 c
字符串的示例中，嵌入空间中的余弦相似性无法判断查询在语义上是等价的。然而，一个足够好的语言模型肯定可以。正如我们在 语言模型用于 数据[7]中演示的那样，我们甚至可以设置一个分类器并使用 逻辑偏差 技巧[8]来计算具有单个输出标记的相似性。
这种技巧使语言模型更快，但对于实际应用来说，语言模型用于搜索的成本和速度仍然是一个问题。随着语言模型变得更快、运行成本更低，我们可以预期随着时间的推移，这个问题会得到缓解。
对于较大的上下文长度，我们还可以使用对语言模型的初始调用来提取相关的文本部分，然后使用该扩充的上下文进行第二次调用。
后续步骤
自己尝试一下！你只需要一个 OpenAI 帐户和上面的代码。
引用链接
[1]
Using DuckDB + Ibis for RAG: https://ibis-project.org/posts/duckdb-for-rag/[2]
固定大小数组: https://duckdb.org/docs/sql/data_types/array[3]
DuckDB 是 Ibis 的默认后端: https://ibis-project.org/posts/why-duckdb/[4]
MotherDuck: https://motherduck.com[5]
OpenAI 的嵌入模型: https://platform.openai.com/docs/guides/embeddings/embedding-models[6]
Hamel Husain: https://twitter.com/HamelHusain/status/1709740984643596768[7]
语言模型用于 数据: https://ibis-project.org/posts/lms-for-data/[8]
逻辑偏差 技巧: https://www.askmarvin.ai/docs/text/classification/#classifying-text
