EmbeddingGemma in Go, tokenizer and all

On the home page of my personal site live two small AI robots, teblo and teblito. Ask them something about me and they answer, mostly with real facts, sometimes hallucinating, and teblito, the smaller one, likes to correct teblo or play tricks on him.

Their lines come from a language model in Paris. But before any of that, something smaller has to understand the question. Is there a prepared answer for it, such as how to reach me? Is someone trying to trick the robots?

Which parts of my projects' documentation should go to the LLM, and which questions should be offered next? And if the language model is not available, which of 96 prepared answers fits best?

Exact matches and the links between answers are plain code and data. Everything that needs to understand what a question means is the job of an embedding model, and it runs right there, inside a Go service on a server with two cores and 1.9 GiB of memory.

This is the story of that part.

Words were not enough

My first plan was keyword search: MiniSearch over a list of example questions. It ranked the right answer first in 11 to 40 percent of cases. It also found answers for nonsense, just as confidently. That was the end of the simplest plan.

Embeddings work differently. A model turns each question into a vector: a list of numbers. Questions with the same meaning end up close together, even when they use different words. "What do you do for a living?" and "what's your job?" share no word that matters, and still the model puts them at 0.92 on a scale up to 1. "How can I reach you?" gets 0.57 against the job question.

The remaining question then was: Which model?

Five candidates, one test

There are now 1,775 example questions for the prepared answers. To choose the embedding model, I used two collections of them, each written by a different language model. In each collection, I held back every third question, matched it against the rest and checked whether the nearest match led to the right answer. The candidates: MiniSearch from before, potion-base-8M (a static Model2Vec model), all-MiniLM-L6-v2, multilingual-e5-small and EmbeddingGemma.

In English, the small all-MiniLM-L6-v2 was as good as EmbeddingGemma at finding the right answer first, and far faster. Had the test ended there, MiniLM would have won. It did not end there. (This benchmark ran in JavaScript and averaged the model's raw output; more on that below.)

One caveat belongs here and not in a footnote: these questions are synthetic. Language models wrote them, not visitors. How well this works for real people I will only know once enough real questions have come in.

The German test

This site is in English, but I live and work in Germany, and some visitors will ask in German. One day I want to translate the most important parts of this site into German too, but I did not want to keep a second set of examples. So I added a second test: 33 German questions, matched against English examples only, and counted as correct when the nearest example belonged to the right topic, such as work, projects or the robots themselves.

32 of 33 for EmbeddingGemma, 11 to 28 for the others. That decided it. Matching the topic is a broader test than matching the exact answer; a later run against both collections merged, with 32 of the questions (I lost one somewhere along the way), found the exact answer for 30. Either way, a German question finds its English neighbours without anyone writing German examples.

The price

EmbeddingGemma is not small.

The q4 export of EmbeddingGemma takes about 200 MB on disk, compared with 23 MB for MiniLM in q8. In the JavaScript benchmark it took 43 milliseconds per question against MiniLM's 2.5, averaged over a hundred runs of one question. 200 MB to tell hi from hire me. A lot for a simple question box, but I'm perfectly fine with that: the language model behind it takes about a second and a half anyway, and the server, which only serves static sites and a few small Go services, has the memory to spare.

Into Go

As said, everything else on that server is small and plain, and I wanted to keep it that way: one Go binary beside the ONNX Runtime library, one systemd unit, no Python. The model itself runs in ONNX Runtime, through onnxruntime_go. That part is easy: open a session, hand it two tensors, read one back. The real work is around it: turning a question into tokens before, and reading the right output after.

Opening the session gives the model one thread and names the output we want (error checks left out):

options.SetIntraOpNumThreads(1)
options.SetInterOpNumThreads(1)
session, err := ort.NewDynamicAdvancedSession(filepath.Join(directory, "onnx", "model_q4.onnx"),
	[]string{"input_ids", "attention_mask"}, []string{"sentence_embedding"}, options)

Why sentence_embedding and not the other output? More on that below.

The tokenizer

A model does not read text, it reads token IDs. EmbeddingGemma ships its tokenizer as a 20 MB tokenizer.json with 262,144 tokens and 514,906 merge rules. I did not want a dependency for it, so I wrote my own. The format can describe many kinds of tokenizers; mine handles only the kind this one file describes. That keeps it short, five steps:

1. Replace every space with ▁. That is all the normalisation this tokenizer does: no lowercasing, no Unicode tricks.
2. Split the text into single characters.
3. Find the neighbouring pair with the best merge rank, merge every occurrence of it, and repeat until no neighbouring pair has a merge rule.
4. Look up each piece. A piece the vocabulary does not know falls back to its UTF-8 bytes, as tokens like <0xC3>.
5. Put <bos> in front and <eos> at the end.

In Go, the heart of it:

text = strings.ReplaceAll(text, " ", "▁")
var symbols []string
for _, r := range text {
	symbols = append(symbols, string(r))
}
for len(symbols) > 1 {
	best, at := int(^uint(0)>>1), -1
	for i := 0; i+1 < len(symbols); i++ {
		if rank, ok := t.ranks[symbols[i]+"\x00"+symbols[i+1]]; ok && rank < best {
			best, at = rank, i
		}
	}
	if at < 0 {
		break
	}
	// merge every occurrence of the pair at `at`, then go round again
}

It is quadratic, and it does not matter: a question has at most 120 characters, and the longer documentation passages are embedded once, at build time.

Here is what it makes of two questions, as typed. (The service itself lowercases a question and strips a trailing question mark before it gets here.)

"Programmiersprachen" breaks down into ▁Program, miers, pr and achen. There is no token for the whole word, and still the model knows what it means, as we will see further down.

Before trusting the tokenizer, I ran all 974 test questions through the reference implementation and through mine: 974 of 974 identical token sequences!

The tests that stayed with the service check the merges, the two-byte fallback for é and the rejection of invalid UTF-8. With the real vocabulary, they also check that four questions decode back to their exact text. The one I like best checks that a missing byte token is an error. In the first prototype a missing entry quietly became token 0, because that is what a Go map returns for a key it does not have. Token 0 is <pad>: where a character without a token of its own stood, the model would have read padding.

It is the kind of bug I like least: nothing crashes, the answers just get a little worse, and nobody knows why.

// the prototype: a missing byte token quietly becomes token 0
ids = append(ids, vocab[fmt.Sprintf("<0x%02X>", c)])

// the service: a missing byte token is an error
for _, value := range []byte(symbol) {
	token := fmt.Sprintf("<0x%02X>", value)
	id, ok := t.vocabulary[token]
	if !ok {
		return nil, fmt.Errorf("tokenizer: missing byte fallback %s", token)
	}
	ids = append(ids, id)
}

With the real vocabulary this never happens: all 256 byte tokens are there. But a wrong file should stop the service the first time it matters, not hand the model a quiet zero.

One more detail is easy to miss: the prefix. EmbeddingGemma was trained with task prompts, and for matching a question against questions both sides get task: sentence similarity | query:. In "who is michael?" that is 7 of 13 tokens. Documentation passages get a prefix of their own. (Questions keep the sentence-similarity prefix when they are matched against documentation passages too. I have not measured whether a retrieval prefix would find better passages.)

const embeddingPrefix = "task: sentence similarity | query: "

func (m *Model) Embed(ctx context.Context, question string) ([]float32, error) {
	return m.embed(ctx, embeddingPrefix+question)
}

func (m *Model) EmbedPassage(ctx context.Context, heading, passage string) ([]float32, error) {
	return m.embed(ctx, "title: "+heading+" | text: "+passage)
}

The output that matters

The ONNX export of EmbeddingGemma has two outputs, last_hidden_state and sentence_embedding. Many pipelines take the first and average it, one vector per token into one vector per question. My first benchmark did exactly that, and it is the benchmark that chose the model.

But the model does not stop at the average. It was trained with three more steps after it: two dense layers, which reshape the vector with weights learned in training, and a normalisation to length 1. last_hidden_state comes out before all three, sentence_embedding after them.

In Go, reading the output takes a type assertion and one more normalisation. The output also has to be freed explicitly: it lives in memory managed by ONNX Runtime, where Go's garbage collector cannot see it. The model has normalised already, but the second pass costs nothing and turns a broken result, all zeros or NaN, into an error (the shape check left out):

output := []ort.Value{nil}
defer func() {
	if output[0] != nil {
		output[0].Destroy()
	}
}()
if err := m.session.Run([]ort.Value{input, attention}, output); err != nil {
	return nil, err
}
tensor, ok := output[0].(*ort.Tensor[float32])
if !ok {
	return nil, errors.New("sentence_embedding is not a float32 tensor")
}
return normalizeVector(tensor.GetData())

The Go version, with sentence_embedding, ranked the right answer first in 50 to 69 percent of cases, against 49 to 67 percent in the JavaScript benchmark. This was not a controlled comparison: the runtime and the tokenizer changed too, so I cannot attribute the difference to the output alone.

What a number means

Every vector is a unit vector, so comparing two of them takes a dot product. There is no vector database behind it, just a loop over the 1,775 examples. (If they ever reach a million, I will think again.) Each answer keeps the score of its best example:

best := make(map[string]float32)
for _, example := range index {
	score := Dot(vector, example.Vector)
	if previous, exists := best[example.ID]; !exists || score > previous {
		best[example.ID] = score
	}
}

For unit vectors the dot product is the cosine of their angle, a number between −1 and 1. It is tempting to read it as a percentage. It is not one.

German and English versions of the same question score 0.88 to 0.96. Pairs of questions with different meanings still score 0.45 to 0.59, and the figure shows only the part of the scale from 0.4 up. So a threshold cannot be guessed. When the language model is not available, a fallback answers, asks back or shrugs. To set its thresholds, I tried a grid of them on held-back questions and counted right answers, wrong ones and questions back. At the threshold I chose, about 48 percent got the right answer directly and 3 percent a wrong one; most of the rest got a question back with the right answer among its choices. Synthetic questions, again: this calibrates the fallback, it does not predict what visitors will see.

Later, with both collections merged into one set of 95 answers and 1,811 questions, I ran three rounds, each time holding back a different third and matching it against the rest. The right answer came first in 79 to 82 percent of cases and among the first three in 92 to 93 percent. Today's 96 answers have not been measured again.

Running it

• One session, shared. ONNX Runtime allows concurrent runs on a session, so every request uses the same one. Each run gets one thread; the service admits two at a time, so two questions at once can still fill both cores.
• Speed. About 25 ms for "who is michael?" and 40 ms for a question of 52 characters, the median of 60 runs, on one core of my laptop (i7-12700H). On the server itself I have not measured it.
• Memory. The prototype peaked at about 300 MB of resident memory while it embedded 1,811 questions, again on the laptop.
• One model, two jobs. The same model also embeds passages of my projects' documentation at build time, with a different prefix. That is how teblo can say something true about hadley or Lima. It deserves its own article.

If you try this yourself

• Measure with your own questions, in the languages your visitors use. Leaderboards could not tell me how German questions fare against English examples.
• Check your tokenizer against the reference on every test question, and make a missing entry an error.
• Use the model as it was trained: with its prefix, and with sentence_embedding if it has one.

Ask teblo something on the home page, if you like in German too, and see whether teblito agrees. Real questions are the test set I don't have yet. (What is kept, and for how long, is on the privacy page.)