从 Python 使用嵌入#

您可以使用其模型 ID 或别名加载嵌入模型，如下所示

import llm

embedding_model = llm.get_embedding_model("3-small")

要嵌入字符串，并返回一个浮点数的 Python 列表，请使用 .embed() 方法

vector = embedding_model.embed("my happy hound")

如果嵌入模型可以处理二进制输入，您可以使用字节字符串调用 .embed()。您可以检查 supports_binary 属性以查看是否支持此功能

if embedding_model.supports_binary:
    vector = embedding_model.embed(open("my-image.jpg", "rb").read())

embedding_model.supports_text 属性指示模型是否支持文本输入。

许多嵌入模型在一次性嵌入多个字符串或二进制字符串时效率更高。要一次性嵌入多个字符串，请使用 .embed_multi() 方法

vectors = list(embedding_model.embed_multi(["my happy hound", "my dissatisfied cat"]))

这返回一个生成器，为每个字符串生成一个嵌入向量。

嵌入是分批计算的。默认情况下，除非底层嵌入模型定义了其首选的批处理大小，否则所有项将在单个批处理中处理。您可以通过 batch_size=N 传递自定义批处理大小，例如

vectors = list(embedding_model.embed_multi(lines_from_file, batch_size=20))

使用集合#

llm.Collection 类可用于在 Python 代码中处理嵌入的集合。

集合是一组命名的嵌入向量，每个向量及其 ID 存储在 SQLite 数据库表中。

要以这种方式处理嵌入，您需要一个 sqlite-utils Database 对象的实例。然后，您可以将其与集合的唯一字符串名称以及将用于该集合的嵌入模型的 ID 一起传递给 llm.Collection 构造函数

import sqlite_utils
import llm

# This collection will use an in-memory database that will be
# discarded when the Python process exits
collection = llm.Collection("entries", model_id="3-small")

# Or you can persist the database to disk like this:
db = sqlite_utils.Database("my-embeddings.db")
collection = llm.Collection("entries", db, model_id="3-small")

# You can pass a model directly using model= instead of model_id=
embedding_model = llm.get_embedding_model("3-small")
collection = llm.Collection("entries", db, model=embedding_model)

如果集合已存在于数据库中，您可以省略 model 或 model_id 参数 - 模型 ID 将从 collections 表中读取。

要嵌入单个字符串并将其存储在集合中，请使用 embed() 方法

collection.embed("hound", "my happy hound")

这会将字符串“my happy hound”的嵌入存储在 entries 集合中，键为 hound。

添加 store=True 以将文本内容本身与嵌入向量一起存储在数据库表中。

要将附加元数据附加到项，请传递一个 JSON 兼容的字典作为 metadata= 参数

collection.embed("hound", "my happy hound", metadata={"name": "Hound"}, store=True)

这些附加元数据将以 JSON 格式存储在嵌入数据库表的 metadata 列中。

批量存储嵌入#

collection.embed_multi() 方法可用于一次性存储多个项的嵌入。这对于某些嵌入模型可能更高效。

collection.embed_multi(
    [
        ("hound", "my happy hound"),
        ("cat", "my dissatisfied cat"),
    ],
    # Add this to store the strings in the content column:
    store=True,
)

要包含与每个项一起存储的元数据，请调用 embed_multi_with_metadata()

collection.embed_multi_with_metadata(
    [
        ("hound", "my happy hound", {"name": "Hound"}),
        ("cat", "my dissatisfied cat", {"name": "Cat"}),
    ],
    # This can also take the store=True argument:
    store=True,
)

batch_size= 参数默认为 100，除非嵌入模型本身定义了更小的批处理大小，否则将使用此值。如果您在嵌入大型集合时遇到内存问题，可以调整此值

collection.embed_multi(
    (
        (i, line)
        for i, line in enumerate(lines_in_file)
    ),
    batch_size=10
)

Collection 类参考#

集合实例具有以下属性和方法

  • id - 集合在数据库中的整数 ID

  • name - 集合的字符串名称（在数据库中是唯一的）

  • model_id - 用于此集合的嵌入模型的字符串 ID

  • model() - 基于 model_id 返回 EmbeddingModel 实例

  • count() - 返回集合中的项的整数数量

  • embed(id: str, text: str, metadata: dict=None, store: bool=False) - 嵌入给定字符串并将其存储在给定 ID 下的集合中。可以选择包含元数据（存储为 JSON）并将文本内容本身存储在数据库表中。

  • embed_multi(entries: Iterable, store: bool=False, batch_size: int=100) - 见上文

  • embed_multi_with_metadata(entries: Iterable, store: bool=False, batch_size: int=100) - 见上文

  • similar(query: str, number: int=10) - 返回与给定查询字符串的嵌入最相似的条目列表

  • similar_by_id(id: str, number: int=10) - 返回与给定 ID 的项的嵌入最相似的条目列表

  • similar_by_vector(vector: List[float], number: int=10, skip_id: str=None) - 返回与给定嵌入向量最相似的条目列表，可以选择跳过给定 ID 的条目

  • delete() - 从数据库中删除集合及其嵌入

还有一个类方法 Collection.exists(db, name)，它返回一个布尔值，可用于确定集合是否存在于数据库中

if Collection.exists(db, "entries"):
    print("The entries collection exists")

检索相似项#

填充嵌入集合后，您可以使用 similar() 方法检索与给定字符串最相似的条目。

此方法使用暴力搜索方法，计算与每个文档的距离分数。这对于小型集合来说是可以的，但无法扩展到大型集合。请参阅 issue 216 以了解通过插件提供的向量索引添加更具可伸缩性方法的计划。

for entry in collection.similar("hound"):
    print(entry.id, entry.score)

该字符串将首先使用集合的模型进行嵌入。

返回的 entry 对象是具有以下属性的对象

  • id - 项的字符串 ID

  • score - 项与查询字符串之间的浮点相似度分数

  • content - 项的字符串文本内容，如果已存储 - 否则为 None

  • metadata - 项的字典（来自 JSON）元数据，如果已存储 - 否则为 None

这默认返回 10 个最相似的项。您可以通过传递不同的 number= 参数来更改此设置

for entry in collection.similar("hound", number=5):
    print(entry.id, entry.score)

方法 similar_by_id() 接受集合中另一个项的 ID，并根据已为其存储的嵌入返回与该项最相似的项

for entry in collection.similar_by_id("cat"):
    print(entry.id, entry.score)

该项本身将从结果中排除。

SQL 架构#

以下是嵌入数据库使用的 SQL 架构

CREATE TABLE [collections] (
   [id] INTEGER PRIMARY KEY,
   [name] TEXT,
   [model] TEXT
)
CREATE TABLE "embeddings" (
   [collection_id] INTEGER REFERENCES [collections]([id]),
   [id] TEXT,
   [embedding] BLOB,
   [content] TEXT,
   [content_blob] BLOB,
   [content_hash] BLOB,
   [metadata] TEXT,
   [updated] INTEGER,
   PRIMARY KEY ([collection_id], [id])
)