Binary embeddings

[simonw]

Refs:

• Support for embedding binary files #253

Still needed:

• Decide what to do about --store - should it store the binary content in a new content_blob column? See these notes.
• The model.supports_binary boolean needs to be respected - it should raise errors if you attempt to embed binary content against a text-only model
• Should there be a way to mark a model as ONLY accepting binary content? Probably yes - though the models I have played with so far like CLIP and ImageBind are happy to accept either kind of content.
• Lots of tests, using a mock embedding model that can handle binary content.
• Build a plugin that uses this (I have a draft llm-clip one already).

[simonw]

I can build plugins using:

• https://www.sbert.net/examples/applications/image-search/README.html#usage
• https://github.com/mlfoundations/open_clip
• https://github.com/facebookresearch/ImageBind - looks like I'll need to vendor their code since it isn't on PyPI

[simonw]

Here's my rough CLIP plugin:

import llm
from PIL import Image
from sentence_transformers import SentenceTransformer, util
import io


@llm.hookimpl
def register_embedding_models(register):
    register(ClipEmbeddingModel())


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

    def __init__(self):
        self._model = None
        self._processor = None
        self._tokenizer = None

    def embed_batch(self, items):
        # Embeds a mix of text strings and binary images
        if self._model is None:
            self._model = SentenceTransformer('clip-ViT-B-32')

        to_embed = []

        for item in items:
            if isinstance(item, bytes):
                # If the item is a byte string, treat it as image data and convert to Image object
                to_embed.append(Image.open(io.BytesIO(item)))
            elif isinstance(item, str):
                to_embed.append(item)

        embeddings = self._model.encode(to_embed)
        return [[float(num) for num in embedding] for embedding in embeddings]

Though looking at this, it's really just a sentence transformer that can accept both binary and text. It could go in llm-sentence-transformers.

I could still have llm-clip be a plugin that just depends on this and then registers the right model.

[simonw]

I'm going to try the content_blob column and see how it feels.

[simonw]

If a model only supports binary and does not support text, maybe we can have it treat all input as --binary even if you forget to use that flag?

[simonw]

I'm landing this. Future tests will happen as I write the plugins.