从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对象的实例。然后,您可以将其传递给llm.Collection构造函数,同时传递集合的唯一字符串名称和您将与该集合一起使用的嵌入模型的ID:

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)

如果集合已经存在于数据库中,您可以省略modelmodel_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
)

集合类参考#

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

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

  • name - 集合的字符串名称(在数据库中唯一)

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

  • model() - 返回基于model_idEmbeddingModel实例

  • 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])
)