go-encoderfile

Go package for working with Mozilla's encoderfile tool.

Important

This package is NOT feature complete meaning that it currently only focuses on embeddings-specific functionality of an encoderfile instance.

Motivation

Honestly, the package exists mostly to provide functionality to pool the per-token embeddings returned by an encoderfile instance in to a single set of embedding for the input phrase. If I had more time I might have dug in to the encoderfile source and provided a patch to do the same thing there but I don't (have the time) so this exists instead.

Example

package main

import (
	"context"
	"encoding/json"
	"flag"
	"log"
	"os"

	"github.com/sfomuseum/go-encoderfile/client"
	"github.com/sfomuseum/go-encoderfile/embeddings"
)

func main() {

	var client_uri string
	var normalize bool
	var pool bool

	flag.StringVar(&client_uri, "client-uri", "http://localhost:8080", "A registered sfomuseum/go-encoderfile/client.Client URI.")
	flag.BoolVar(&normalize, "normalize", true, "Boolean flag to indicate the encoder server endpoint should normalize embeddings.")
	flag.BoolVar(&pool, "pool", true, "Boolean flag to indicate the per-token embeddings for the input should be pooled in to a single embedding.")

	flag.Parse()

	input := flag.Args()

	ctx := context.Background()

	cl, _ := client.NewClient(ctx, client_uri)

	var output any

	rsp, _ := cl.Embeddings(ctx, input, normalize)

	output = rsp

	if pool {

		pooled, _ := embeddings.PoolOutputResults(rsp)
		output = pooled
	}

	enc := json.NewEncoder(os.Stdout)
	enc.Encode(output)
}

Building encoderfile

$> cargo build --bin encoderfile --release

$> mkdir work
$> python -m venv work
$> cd work
$> ./bin/activate

$> ./bin/pip install "optimum[onnxruntime]"
$> ./bin/pip install sentence-transformers

Note: The docs say to run pip install "optimum[exporters]" but that no longer works with newer versions of optimum.

$> ./bin/optimum-cli export onnx \
	--model sentence-transformers/all-MiniLM-L6-v2 \
	--task feature-extraction \
	--library transformers \
	./models/minilm

$> cat minilm.embedding.cfg
encoderfile:
  name: minilm
  model_type: embedding
  output_path: ./build/minilm.encoderfile
  path:
    model_config_path: ./models/minilm/config.json
    model_weights_path: ./models/minilm/model.onnx
    tokenizer_path: ./models/minilm/tokenizer.json

$> ../target/release/encoderfile build -f minilm.embedding.cfg
$> chmod 755 build/minilm.encoderfile

$> ./build/minilm.encoderfile serve --http-hostname localhost --disable-grpc