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); + }); + }); + }); +});