Writing plugins to add new embedding models#
Read the plugin tutorial for details on how to develop and package a plugin.
This page shows an example plugin that implements and registers a new embedding model.
There are two components to an embedding model plugin:
An implementation of the
register_embedding_models()
hook, which takes aregister
callback function and calls it to register the new model with the LLM plugin system.A class that extends the
llm.EmbeddingModel
abstract base class.The only required method on this class is
embed_batch(texts)
, which takes an iterable of strings and returns an iterator over lists of floating point numbers.
The following example uses the sentence-transformers package to provide access to the MiniLM-L6 embedding model.
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)
Once installed, the model provided by this plugin can be used with the llm embed command like this:
cat file.txt | llm embed -m sentence-transformers/all-MiniLM-L6-v2
Or via its registered alias like this:
cat file.txt | llm embed -m all-MiniLM-L6-v2
llm-sentence-transformers is a complete example of a plugin that provides an embedding model.
Execute Jina embeddings with a CLI using llm-embed-jina talks through a similar process to add support for the Jina embeddings models.
Embedding binary content#
If your model can embed binary content, use the supports_binary
property to indicate that:
class ClipEmbeddingModel(llm.EmbeddingModel):
model_id = "clip"
supports_binary = True
supports_text= True
supports_text
defaults to True
and so is not necessary here. You can set it to False
if your model only supports binary data.
If your model accepts binary, your .embed_batch()
model may be called with a list of Python bytestrings. These may be mixed with regular strings if the model accepts both types of input.
llm-clip is an example of a model that can embed both binary and text content.