datasette-faiss

为指定的 Datasette 表维护一个 FAISS 索引

有关此项目的背景信息，请参阅语义搜索答案：使用 GPT3 + OpenAI 嵌入对文档进行问答。

安装

在 Datasette 所在的同一环境中安装此插件。

datasette install datasette-faiss

用法

此插件在启动时为指定的表创建内存中的 FAISS 索引，使用 IndexFlatL2 FAISS 索引类型。

如果表在服务器启动后被修改，索引尚不会反映这些更改。

配置

要建立索引的表必须包含 id 和 embedding 列。embedding 列必须是一个 blob，其中包含已使用以下 Python 函数编码的浮点数数组形式的嵌入向量

def encode(vector):
    return struct.pack("f" * len(vector), *vector)

您可以像这样从该包导入该函数

from datasette_faiss import encode

您可以通过将此内容添加到 metadata.json 来指定哪些表应为其创建索引

{
    "plugins": {
        "datasette-faiss": {
            "tables": [
                ["blog", "embeddings"]
            ]
        }
    }
}

每个表都是一个数组，列出了数据库名称和表名称。

如果您使用 metadata.yml，配置应如下所示

plugins:
  datasette-faiss:
    tables:
    - ["blog", "embeddings"]

SQL 函数

此插件在 Datasette 中提供了四个新的 SQL 函数

faiss_search(database, table, embedding, k)

返回在指定的数据库和表中找到的、距离给定 embedding 最近的 k 个邻居。例如

select faiss_search('blog', 'embeddings', (select embedding from embeddings where id = 3), 5)

这将返回一个 JSON 数组，包含 blog 数据库中 embeddings 表中距离指定嵌入向量最近的五个记录的 ID。返回的值如下所示

["1", "1249", "1011", "5", "10"]

您可以使用 SQLite 的 json_each() 函数将其转换为类似表的序列，以便进行连接（JOIN）操作。

以下是一个执行此操作的示例查询

with related as (
  select value from json_each(
    faiss_search(
      'blog',
      'embeddings',
      (select embedding from embeddings limit 1),
      5
    )
  )
)
select * from blog_entry, related
where id = value

faiss_search_with_scores(database, table, embedding, k)

接受与上述函数相同的参数，但返回值是一个 JSON 对数组，每对包含一个 ID 和一个分数 - 示例如下

[
    ["1", 0.0],
    ["1249", 0.21042244136333466],
    ["1011", 0.29391372203826904],
    ["5", 0.29505783319473267],
    ["10", 0.31554925441741943]
]

faiss_encode(json_vector)

给定一个浮点数的 JSON 数组，返回可用于其他函数的二进制嵌入向量 blob

select faiss_encode('[2.4, 4.1, 1.8]')
-- Returns a 12 byte blob
select hex(faiss_encode('[2.4, 4.1, 1.8]'))
-- Returns 9A991940333383406666E63F

faiss_decode(vector_blob)

与 faiss_encode() 功能相反。

select faiss_decode(X'9A991940333383406666E63F')

返回结果

[2.4000000953674316, 4.099999904632568, 1.7999999523162842]

请注意，浮点运算的结果在往返转换时可能与预期值不完全一致。

faiss_agg(id, embedding, compare_embedding, k)

此聚合函数可用于查找表中每个唯一的 id 值对应的、距离 compare_embedding 最近的 k 个邻居。例如

select faiss_agg(
    id, embedding, (select embedding from embeddings where id = 3), 5
) from embeddings

与 faiss_search() 函数不同，此函数不依赖于插件首次运行时创建的每表索引。相反，每次运行聚合函数时都会构建一个索引。

这意味着它只适用于较小的数据集 - 一旦超过 10,000 个左右，此函数的性能可能会变得非常昂贵，难以接受。

该函数返回一个 JSON 数组，包含距离得分最近的 k 行的 ID，如下所示

[1324, 344, 5562, 553, 2534]

您可以使用 json_each() 函数将其转换为类似表的序列，以便进行连接（JOIN）操作。

尝试一个 fais_agg() 示例查询.

faiss_agg_with_scores(id, embedding, compare_embedding, k)

这类似于 faiss_agg() 聚合函数，但它返回一个对列表，每对包含一个 ID 和相应的分数 - 如果 k 是 2，看起来像这样

[[2412, 0.25], [1245, 24.25]]

尝试一个 fais_agg_with_scores() 示例查询.

开发

要在本地设置此插件，首先检出代码。然后创建一个新的虚拟环境

cd datasette-faiss
python3 -m venv venv
source venv/bin/activate

现在安装依赖和测试依赖

pip install -e '.[test]'

运行测试

pytest

datasette-faiss by simonw

README 源代码