跳到内容

嵌入

embeddings embeddings

嵌入数据库是实现语义搜索的引擎。数据被转换为嵌入向量，其中相似的概念将产生相似的向量。这些向量用于构建大大小小的索引。索引用于查找含义相同、但不一定使用相同关键词的结果。

以下代码片段展示了如何构建和搜索嵌入索引。

from txtai import Embeddings

# Create embeddings model, backed by sentence-transformers & transformers
embeddings = Embeddings(path="sentence-transformers/nli-mpnet-base-v2")

data = [
  "US tops 5 million confirmed virus cases",
  "Canada's last fully intact ice shelf has suddenly collapsed, " +
  "forming a Manhattan-sized iceberg",
  "Beijing mobilises invasion craft along coast as Taiwan tensions escalate",
  "The National Park Service warns against sacrificing slower friends " +
  "in a bear attack",
  "Maine man wins $1M from $25 lottery ticket",
  "Make huge profits without work, earn up to $100,000 a day"
]

# Index the list of text
embeddings.index(data)

print(f"{'Query':20} Best Match")
print("-" * 50)

# Run an embeddings search for each query
for query in ("feel good story", "climate change", "public health story", "war",
              "wildlife", "asia", "lucky", "dishonest junk"):
    # Extract uid of first result
    # search result format: (uid, score)
    uid = embeddings.search(query, 1)[0][0]

    # Print text
    print(f"{query:20} {data[uid]}")

构建

嵌入实例是配置驱动的，基于构造函数中传入的内容。向量被存储起来，还可以选择存储内容。内容存储提供了额外的过滤和数据检索选项。

上述示例通过path参数设置了特定的嵌入向量模型。也可以创建一个没有配置的嵌入实例。

embeddings = Embeddings()

在这种情况下，加载和搜索数据时，使用默认的 transformers 向量模型对数据进行向量化。请参阅模型指南获取当前的模型建议。

索引

创建新的嵌入实例后，下一步是将数据添加到其中。

embeddings.index(rows)

index 方法接受一个可迭代对象，并支持以下每种元素的格式。

  • (id, data, tags) - 默认处理格式
元素 描述
id 唯一记录 id
data 要索引的输入数据，可以是文本、字典或对象
tags 可选的标签字符串，用于在数据索引时进行标记/标注
  • (id, data)

与上述相同，但不包含标签。

  • data

要索引的单个元素。在这种情况下，将自动生成唯一的 id。请注意，对于生成的 id，upsert 和 delete 调用需要进行单独的搜索来获取目标 id。

当 data 字段是字典时，文本通过 text 键传入，二进制对象通过 object 键传入。请注意，必须启用内容存储才能存储元数据，必须启用对象存储才能存储二进制对象数据。如果提供了 id 和 tags 键，它们将被提取出来。

输入的可迭代对象可以是列表或生成器。生成器有助于索引非常大的数据集，因为任何时候只有部分数据在内存中。

有关索引的更多信息，请参见索引指南。

数据索引完成后，就可以进行搜索了。

embeddings.search(query, limit)

search 方法接受两个参数：查询和查询限制。结果格式根据是否存储内容而有所不同。

  • 当不存储内容时，结果是 (id, score) 列表
  • 当存储内容时，结果是 {**query columns} 列表

支持自然语言查询和 SQL 查询。更多信息请参见查询指南。

资源管理

嵌入数据库是上下文管理器。以下代码块在完成后会自动关闭并释放资源。

# Create a new Embeddings database, index data and save
with Embeddings() as embeddings:
  embeddings.index(rows)
  embeddings.save(path)

# Search a saved Embeddings database
with Embeddings().load(path) as embeddings:
  embeddings.search(query)

虽然调用 close 不总是必需的（资源会被垃圾回收），但在不再需要共享资源（如数据库连接）时尽快释放它们是最好的做法。

更多示例

请参阅此链接获取嵌入示例的完整列表。