Using a different model instead of "gte-small" - Self Hosted

[kallebysantos]

Hi there,
Since Supabase releases the 'AI inference in supabase edge functions' the inference of embedding became more simpler.
But for some use cases the default embed model gte-small is not so usefull.
For example if you are using a non-english language it can result wrong semantic search or bad score results.

Since that I'd investigate the source code of the edge-functions image and realize that Supabase team has hard coded the gte-small. So if you want to change it to a custom model is there some modifications to do.
The simpler way is by overwriting the default model inside de docker image. Since de api lib just look for a folder called gte-small we can change the model.onnx for whatever we want. In my specific case I choose the paraphrase-multilingual-MiniLM-L12-v2 because this model works very well with Portuguese.

But for it work, we need to get an ONNX runtime converted model, so I we use the Xenova version of it.

So lets get it work!

First of all create a custom Dockerfile

# Dockerfile.functions

FROM supabase/edge-runtime:v1.41.2 AS build

# Change here with the HF model name. It must be ONNX runtime compatible
ENV MODEL_NAME=Xenova/paraphrase-multilingual-MiniLM-L12-v2

# Setting up
RUN apt update && apt install curl -y

# Replacing 'gte-small' by $MODEL_NAME
WORKDIR /etc/sb_ai/models/gte-small

# Downloading the custom model
RUN curl -L -o model.onnx https://huggingface.co/${MODEL_NAME}/resolve/main/onnx/model_quantized.onnx?download=true
RUN curl -L -o tokenizer.json https://huggingface.co/${MODEL_NAME}/resolve/main/tokenizer.json?download=true

FROM supabase/edge-runtime:v1.41.2 AS run

COPY --from=build /etc/sb_ai/models/gte-small /etc/sb_ai/models/gte-small

Then modify your compose to build and use our custom image

services:
  # others supabase services ...
  supabase-functions:
      container_name: supabase-edge-functions
      
      # Here you need to specify that you want do docker build a image with your custom model
      image: supabase/edge-runtime:v1.41.2-custom-model
      build:
        context: .
        dockerfile: Dockerfile.functions
      restart: unless-stopped
      # The rest of config doesn't matter, you can keep it as default or your custom

Finally we can call our custom model

Here you need to keep the gte-small reference for api compatibility, because it was hard coded

/// <reference types="https://esm.sh/@supabase/functions-js/src/edge-runtime.d.ts" />
// gte-small has been overriden by paraphrase-multilingual-MiniLM-L12-v2

// Here you need to keep the 'gte-small' for API compatility. 
// But since we overwrite the model inside container it will use our custom model
const embedder = new Supabase.ai.Session("gte-small");

Deno.serve(async (req) => {
  // ...

  // @ts-ignore Embedder generates unknown type
  const outputEmbed: number[] = await embedder.run("Olá Mundo!", {
    mean_pool: true,
    normalize: true,
  });

  // Rest of your code ...
}

I now that is a little tricky but should work until Supabase release a built-in support for custom models, as well some GPU support should be nice too

[laktek]

Nice that drop-in model worked. Let me look into easier ways to support custom embedding models.

[kallebysantos]

Hi guys, just to let you know. I'd fork the edge-runtime and I'm currently working on this feature.
You can follow it here.
Currently is already possible to install and use different kinds of embedding/feature-extraction models. So my next step is to implement a Transformers like API to handle more kinds of pipelines.