From f3e45923837e9a5a668ac4ec6822c3c7c3a9b05e Mon Sep 17 00:00:00 2001
From: "Tomwang.eth" <73059070+tomwangcn@users.noreply.github.com>
Date: Sat, 7 Nov 2026 19:34:12 +0000
Subject: [PATCH] fix: prepend BOS token for `ugm` vocabulary models that
 request it

`resolveBeginningTokenToPrepend` returned `null` for `ugm` vocabulary
models unconditionally, so UGM models with `add_bos_token = true` in
their GGUF metadata (such as bge-m3) never got the BOS token prepended
to embedding, ranking and completion inputs. This diverged from
`llama.cpp`, which prepends BOS for UGM when `add_bos` is set, and
broke CLS-pooled embeddings, where the first content token was pooled
instead of `<s>`.

Respect `tokens.shouldPrependBosToken` for `ugm` models to match the
`llama.cpp` behavior. Models that don't request BOS (such as T5) are
unaffected.

Closes: #683
---
 src/utils/tokenizerUtils.ts                  |  4 +-
 test/standalone/utils/tokenizerUtils.test.ts | 99 ++++++++++++++++++++
 2 files changed, 102 insertions(+), 1 deletion(-)
 create mode 100644 test/standalone/utils/tokenizerUtils.test.ts

diff --git a/src/utils/tokenizerUtils.ts b/src/utils/tokenizerUtils.ts
index ce459966..187d2fe5 100644
--- a/src/utils/tokenizerUtils.ts
+++ b/src/utils/tokenizerUtils.ts
@@ -14,7 +14,9 @@ export function resolveBeginningTokenToPrepend(vocabularyType: LlamaVocabularyTy
 
 
     if (vocabularyType === LlamaVocabularyType.ugm)
-        return null;
+        return tokens.shouldPrependBosToken
+            ? tokens.bos
+            : null;
 
     if (tokens.shouldPrependBosToken)
         return tokens.bos;
diff --git a/test/standalone/utils/tokenizerUtils.test.ts b/test/standalone/utils/tokenizerUtils.test.ts
new file mode 100644
index 00000000..50d61cad
--- /dev/null
+++ b/test/standalone/utils/tokenizerUtils.test.ts
@@ -0,0 +1,99 @@
+import {describe, expect, test} from "vitest";
+import {LlamaVocabularyType} from "../../../src/bindings/types.js";
+import {resolveBeginningTokenToPrepend, resolveEndTokenToAppend} from "../../../src/utils/tokenizerUtils.js";
+import type {LlamaModelTokens} from "../../../src/evaluator/LlamaModel/LlamaModel.js";
+import type {Token} from "../../../src/types.js";
+
+function createTokens({
+    bos = 1 as Token,
+    eos = 2 as Token,
+    sep = 3 as Token,
+    shouldPrependBosToken = false,
+    shouldAppendEosToken = false
+}: {
+    bos?: Token | null,
+    eos?: Token | null,
+    sep?: Token | null,
+    shouldPrependBosToken?: boolean,
+    shouldAppendEosToken?: boolean
+} = {}) {
+    return {
+        bos,
+        eos,
+        sep,
+        shouldPrependBosToken,
+        shouldAppendEosToken
+    } as LlamaModelTokens;
+}
+
+describe("utils", () => {
+    describe("tokenizerUtils", () => {
+        describe("resolveBeginningTokenToPrepend", () => {
+            test("prepends BOS for ugm models that request it (e.g. bge-m3)", () => {
+                // XLM-RoBERTa-based models (like bge-m3) are converted with a ugm vocabulary
+                // and set `tokenizer.ggml.add_bos_token = true` in the GGUF metadata.
+                // llama.cpp prepends the BOS token for ugm when `add_bos` is set
+                // (`llama_vocab::impl::tokenize` in `llama-vocab.cpp`).
+                const tokens = createTokens({shouldPrependBosToken: true});
+                expect(resolveBeginningTokenToPrepend(LlamaVocabularyType.ugm, tokens)).to.eql(tokens.bos);
+            });
+
+            test("prepends nothing for ugm models that do not request BOS (e.g. T5)", () => {
+                const tokens = createTokens({shouldPrependBosToken: false});
+                expect(resolveBeginningTokenToPrepend(LlamaVocabularyType.ugm, tokens)).to.eql(null);
+            });
+
+            test("prepends BOS for spm models that request it", () => {
+                const tokens = createTokens({shouldPrependBosToken: true});
+                expect(resolveBeginningTokenToPrepend(LlamaVocabularyType.spm, tokens)).to.eql(tokens.bos);
+            });
+
+            test("prepends BOS for bpe models that request it", () => {
+                const tokens = createTokens({shouldPrependBosToken: true});
+                expect(resolveBeginningTokenToPrepend(LlamaVocabularyType.bpe, tokens)).to.eql(tokens.bos);
+            });
+
+            test("prepends nothing for spm models that do not request BOS", () => {
+                const tokens = createTokens({shouldPrependBosToken: false});
+                expect(resolveBeginningTokenToPrepend(LlamaVocabularyType.spm, tokens)).to.eql(null);
+            });
+
+            test("prepends BOS for wpm models", () => {
+                const tokens = createTokens({shouldPrependBosToken: false});
+                expect(resolveBeginningTokenToPrepend(LlamaVocabularyType.wpm, tokens)).to.eql(tokens.bos);
+            });
+
+            test("prepends nothing for rwkv models", () => {
+                const tokens = createTokens({shouldPrependBosToken: true});
+                expect(resolveBeginningTokenToPrepend(LlamaVocabularyType.rwkv, tokens)).to.eql(null);
+            });
+        });
+
+        describe("resolveEndTokenToAppend", () => {
+            test("appends EOS for ugm models", () => {
+                const tokens = createTokens();
+                expect(resolveEndTokenToAppend(LlamaVocabularyType.ugm, tokens)).to.eql(tokens.eos);
+            });
+
+            test("appends SEP for wpm models", () => {
+                const tokens = createTokens();
+                expect(resolveEndTokenToAppend(LlamaVocabularyType.wpm, tokens)).to.eql(tokens.sep);
+            });
+
+            test("appends EOS for spm models that request it", () => {
+                const tokens = createTokens({shouldAppendEosToken: true});
+                expect(resolveEndTokenToAppend(LlamaVocabularyType.spm, tokens)).to.eql(tokens.eos);
+            });
+
+            test("appends nothing for spm models that do not request EOS", () => {
+                const tokens = createTokens({shouldAppendEosToken: false});
+                expect(resolveEndTokenToAppend(LlamaVocabularyType.spm, tokens)).to.eql(null);
+            });
+
+            test("appends nothing for rwkv models", () => {
+                const tokens = createTokens({shouldAppendEosToken: true});
+                expect(resolveEndTokenToAppend(LlamaVocabularyType.rwkv, tokens)).to.eql(null);
+            });
+        });
+    });
+});