编写插件以添加新的嵌入模型

阅读插件教程，了解如何开发和打包插件的详细信息。

本页展示了一个示例插件，用于实现和注册新的嵌入模型。

嵌入模型插件有两个组成部分

1. register_embedding_models() 钩子的实现，它接受一个 register 回调函数，并调用该函数将新模型注册到 LLM 插件系统中。

2. 一个扩展 llm.EmbeddingModel 抽象基类的类。

   此类中唯一必需的方法是 embed_batch(texts)，它接受一个字符串的可迭代对象，并返回一个浮点数列表的迭代器。

以下示例使用 sentence-transformers 软件包提供对 MiniLM-L6 嵌入模型的访问。

import llm
from sentence_transformers import SentenceTransformer


@llm.hookimpl
def register_embedding_models(register):
    model_id = "sentence-transformers/all-MiniLM-L6-v2"
    register(SentenceTransformerModel(model_id, model_id), aliases=("all-MiniLM-L6-v2",))


class SentenceTransformerModel(llm.EmbeddingModel):
    def __init__(self, model_id, model_name):
        self.model_id = model_id
        self.model_name = model_name
        self._model = None

    def embed_batch(self, texts):
        if self._model is None:
            self._model = SentenceTransformer(self.model_name)
        results = self._model.encode(texts)
        return (list(map(float, result)) for result in results)

安装后，此插件提供的模型可以使用 llm embed 命令，如下所示

cat file.txt | llm embed -m sentence-transformers/all-MiniLM-L6-v2

或者通过其注册的别名，如下所示

cat file.txt | llm embed -m all-MiniLM-L6-v2

llm-sentence-transformers 是一个提供嵌入模型的完整插件示例。

使用 llm-embed-jina 通过 CLI 执行 Jina 嵌入 介绍了添加对 Jina 嵌入模型支持的类似过程。

嵌入二进制内容

如果您的模型可以嵌入二进制内容，请使用 supports_binary 属性进行指示

class ClipEmbeddingModel(llm.EmbeddingModel):
    model_id = "clip"
    supports_binary = True
    supports_text= True

supports_text 默认为 True，因此在此处不是必需的。如果您的模型仅支持二进制数据，您可以将其设置为 False。

如果您的模型接受二进制，则您的 .embed_batch() 模型可能会被调用，参数为 Python 字节串列表。如果模型同时接受这两种类型的输入，这些字节串可能会与常规字符串混合。

llm-clip 是一个既可以嵌入二进制内容又可以嵌入文本内容的模型示例。

LLM_RAISE_ERRORS

在开发插件时，请求抛出错误而不是捕获和记录错误可能很有用，这样您就可以从 Python 调试器中访问它们。

设置 LLM_RAISE_ERRORS 环境变量以启用此行为，然后像这样运行 llm

LLM_RAISE_ERRORS=1 python -i -m llm ...

-i 选项表示如果发生错误，Python 将进入交互式 shell。然后您可以使用以下方法在最近的错误处打开调试器

import pdb; pdb.pm()