From 142c63d3c8522682c49ea085fea20ff3ecbe3b01 Mon Sep 17 00:00:00 2001 From: Anika Jain Date: Thu, 11 Jun 2026 06:32:21 +0530 Subject: [PATCH 1/2] Add offset tracking to Encoding --- jest.config.mjs | 2 +- src/core/PreTokenizer.ts | 10 +- src/core/Tokenizer.ts | 100 ++++++--- src/core/TokenizerModel.ts | 6 +- src/core/preTokenizer/BertPreTokenizer.ts | 6 +- src/core/preTokenizer/ByteLevel.ts | 27 ++- src/core/preTokenizer/Digits.ts | 4 +- src/core/preTokenizer/FixedLength.ts | 6 +- src/core/preTokenizer/Metaspace.ts | 4 +- src/core/preTokenizer/Punctuation.ts | 4 +- src/core/preTokenizer/Replace.ts | 7 +- src/core/preTokenizer/Sequence.ts | 17 +- src/core/preTokenizer/Split.ts | 17 +- src/core/preTokenizer/Whitespace.ts | 4 +- src/core/preTokenizer/WhitespaceSplit.ts | 2 +- src/core/tokenizerModelImplementations/BPE.ts | 44 +++- .../tokenizerModelImplementations/Legacy.ts | 2 +- .../tokenizerModelImplementations/Unigram.ts | 13 +- .../WordPiece.ts | 15 +- src/static/types.ts | 1 + src/utils/core.ts | 99 +++++++-- .../data-structures/DictionarySplitter.ts | 10 +- src/utils/data-structures/TokenLattice.ts | 8 + tests/bundle.test.ts | 3 +- tests/edge-cases/offsets.test.ts | 209 ++++++++++++++++++ tests/models/bert/offsets.test.ts | 37 ++++ tests/models/gpt2/gpt2.test.ts | 50 +++++ tests/models/roberta/offsets.test.ts | 143 ++++++++++++ tests/models/t5/offsets.test.ts | 124 +++++++++++ tests/tokenizers.test.ts | 2 +- 30 files changed, 859 insertions(+), 117 deletions(-) create mode 100644 tests/edge-cases/offsets.test.ts create mode 100644 tests/models/bert/offsets.test.ts create mode 100644 tests/models/gpt2/gpt2.test.ts create mode 100644 tests/models/roberta/offsets.test.ts create mode 100644 tests/models/t5/offsets.test.ts diff --git a/jest.config.mjs b/jest.config.mjs index ed46c08..29a25d6 100644 --- a/jest.config.mjs +++ b/jest.config.mjs @@ -169,7 +169,7 @@ export default { { useESM: true, tsconfig: { - module: "es2022", + module: "esnext", target: "es2022", }, }, diff --git a/src/core/PreTokenizer.ts b/src/core/PreTokenizer.ts index c87ec79..da4ca99 100644 --- a/src/core/PreTokenizer.ts +++ b/src/core/PreTokenizer.ts @@ -8,7 +8,7 @@ import type { PreTokenizeTextOptions } from "@static/tokenizer"; */ abstract class PreTokenizer extends Callable< [string | string[], any?], - string[] + Array<[string, [number, number]]> > { /** * Method that should be implemented by subclasses to define the specific pre-tokenization logic. @@ -20,7 +20,7 @@ abstract class PreTokenizer extends Callable< abstract pre_tokenize_text( text: string, options?: PreTokenizeTextOptions, - ): string[]; + ): Array<[string, [number,number]]>; /** * Tokenizes the given text into pre-tokens. @@ -31,11 +31,11 @@ abstract class PreTokenizer extends Callable< pre_tokenize( text: string | string[], options?: PreTokenizeTextOptions, - ): string[] { + ): Array<[string, [number,number]]> { return ( Array.isArray(text) ? text.map((x) => this.pre_tokenize_text(x, options)) - : this.pre_tokenize_text(text, options) + : [this.pre_tokenize_text(text, options)] ).flat(); } @@ -45,7 +45,7 @@ abstract class PreTokenizer extends Callable< * @param options Additional options for the pre-tokenization logic. * @returns An array of pre-tokens. */ - _call(text: string | string[], options?: any): string[] { + _call(text: string | string[], options?: any): Array<[string, [number,number]]> { return this.pre_tokenize(text, options); } } diff --git a/src/core/Tokenizer.ts b/src/core/Tokenizer.ts index 3dfa396..94a378c 100644 --- a/src/core/Tokenizer.ts +++ b/src/core/Tokenizer.ts @@ -1,6 +1,7 @@ import DictionarySplitter from "@utils/data-structures/DictionarySplitter"; import AddedToken from "./AddedToken"; import { + build_alignment_map, clean_up_tokenization, is_integral_number, lowercase_and_remove_accents, @@ -164,7 +165,7 @@ class Tokenizer { return_token_type_ids = null, }: EncodeOptions = {}, ): Encoding { - const { tokens, token_type_ids } = this.tokenize_helper(text, { + const { tokens, token_type_ids, offsets } = this.tokenize_helper(text, { text_pair, add_special_tokens, }); @@ -179,6 +180,7 @@ class Tokenizer { ids: input_ids, tokens, attention_mask: new Array(input_ids.length).fill(1), + offsets, }; if (return_token_type_ids && token_type_ids) { @@ -244,7 +246,7 @@ class Tokenizer { return this.tokenize_helper(text, { text_pair, add_special_tokens }).tokens; } - private encode_text(text: string | null): string[] | null { + private encode_text(text: string | null): Array<[string, [number, number]]> | null { if (text === null) { return null; } @@ -257,26 +259,32 @@ class Tokenizer { // 2. Normalize, then split by normalized added tokens (normalized: true) const sections = this.splitter_unnormalized.split(text); - sections.forEach((section, i) => { - const added_token = this.added_tokens_map.get(section); + sections.forEach(([section_text], i) => { + const added_token = this.added_tokens_map.get(section_text); if (added_token) { if (added_token.lstrip && i > 0) { - sections[i - 1] = sections[i - 1].trimEnd(); + const [s, start] = sections[i - 1]; + sections[i - 1] = [s.trimEnd(), start]; } if (added_token.rstrip && i < sections.length - 1) { - sections[i + 1] = sections[i + 1].trimStart(); + const [s, start] = sections[i + 1]; + const trimmed = s.trimStart(); + sections[i + 1] = [trimmed, start + s.length - trimmed.length]; } } }); - return sections.flatMap((processed_text, section_index) => { - if (processed_text.length === 0) { + return sections.flatMap(([section_text, section_offset], section_index) => { + if (section_text.length === 0) { return []; } - if (this.added_tokens_map.has(processed_text)) { - return [processed_text]; + if (this.added_tokens_map.has(section_text)) { + return [[section_text, [section_offset, section_offset + section_text.length]]]; } + const original_section = section_text; + let processed_text = section_text; + if (this.remove_space === true) { processed_text = processed_text.trim().split(/\s+/).join(" "); } @@ -292,36 +300,56 @@ class Tokenizer { return []; } + // Build alignment: alignment[i] = position in original_section of processed_text[i] + const alignment = build_alignment_map(original_section, processed_text); + // Phase 2: Split by normalized tokens on the normalized text const subsections = this.splitter_normalized.split(processed_text); - subsections.forEach((subsection, j) => { - const added_token = this.added_tokens_map.get(subsection); + subsections.forEach(([sub_text], j) => { + const added_token = this.added_tokens_map.get(sub_text); if (added_token) { if (added_token.lstrip && j > 0) { - subsections[j - 1] = subsections[j - 1].trimEnd(); + const [s, start] = subsections[j - 1]; + subsections[j - 1] = [s.trimEnd(), start]; } if (added_token.rstrip && j < subsections.length - 1) { - subsections[j + 1] = subsections[j + 1].trimStart(); + const [s, start] = subsections[j + 1]; + const trimmed = s.trimStart(); + subsections[j + 1] = [trimmed, start + s.length - trimmed.length]; } } }); - return subsections.flatMap((subsection) => { + // Converts a processed_text span [pt_s, pt_e) to an absolute original-text span. + const to_orig = (pt_s: number, pt_e: number): [number, number] => { + const sec_start = pt_s < alignment.length ? alignment[pt_s] : original_section.length; + const sec_end = pt_e < alignment.length ? alignment[pt_e] : original_section.length; + return [section_offset + sec_start, section_offset + sec_end]; + }; + + return subsections.flatMap(([subsection, sub_offset]) => { if (subsection.length === 0) { return []; } if (this.added_tokens_map.has(subsection)) { - return [subsection]; + return [[subsection, to_orig(sub_offset, sub_offset + subsection.length)]]; } - const section_tokens = + // Pre-tokenizer returns spans relative to subsection; shift to processed_text coords. + const word_pairs: Array<[string, [number, number]]> = this.pre_tokenizer !== null - ? this.pre_tokenizer(subsection, { - section_index, - }) - : [subsection]; - return this.model(section_tokens); + ? this.pre_tokenizer(subsection, { section_index }) + : [[subsection, [0, subsection.length]]]; + + const pt_word_pairs: Array<[string, [number, number]]> = word_pairs.map( + ([w, [ws, we]]) => [w, [sub_offset + ws, sub_offset + we]], + ); + + // Model produces processed_text-relative subword spans; map to original-text coords. + return this.model(pt_word_pairs).map( + ([t, [pt_s, pt_e]]) => [t, to_orig(pt_s, pt_e)] as [string, [number, number]], + ); }); }); } @@ -329,13 +357,29 @@ class Tokenizer { private tokenize_helper( text: string, { text_pair = null, add_special_tokens = true }: TokenizeOptions, - ): { tokens: Array; token_type_ids?: Array } { - const tokens1 = this.encode_text(text); - const tokens2 = this.encode_text(text_pair || null); + ): { tokens: string[]; token_type_ids?: number[]; offsets: Array<[number, number]> } { + const pairs1 = this.encode_text(text); + const pairs2 = this.encode_text(text_pair || null); + + const strings1 = pairs1?.map(([t]) => t) ?? null; + const strings2 = pairs2?.map(([t]) => t) ?? null; + + const { tokens, token_type_ids } = this.post_processor + ? this.post_processor(strings1, strings2, add_special_tokens) + : { tokens: merge_arrays(strings1 ?? [], strings2 ?? []) }; + + // Align output tokens with their spans. Special tokens added by the post-processor + // (not present in the original pairs) receive a [0, 0] placeholder span. + const all_pairs = [...(pairs1 ?? []), ...(pairs2 ?? [])]; + let pair_i = 0; + const offsets: Array<[number, number]> = tokens.map((t) => { + if (pair_i < all_pairs.length && all_pairs[pair_i][0] === t) { + return all_pairs[pair_i++][1]; + } + return [0, 0]; + }); - return this.post_processor - ? this.post_processor(tokens1, tokens2, add_special_tokens) - : { tokens: merge_arrays(tokens1 ?? [], tokens2 ?? []) }; + return { tokens, token_type_ids, offsets }; } /** diff --git a/src/core/TokenizerModel.ts b/src/core/TokenizerModel.ts index e3bd9c4..a066be1 100644 --- a/src/core/TokenizerModel.ts +++ b/src/core/TokenizerModel.ts @@ -6,7 +6,7 @@ import type { TokenizerModelConfig } from "@static/tokenizer"; /** * Abstract base class for tokenizer models. */ -abstract class TokenizerModel extends Callable<[string[]], string[]> { +abstract class TokenizerModel extends Callable<[Array<[string, [number, number]]>], Array<[string, [number, number]]>> { config: TokenizerModelConfig; vocab: string[]; /** A mapping of tokens to ids. */ @@ -37,7 +37,7 @@ abstract class TokenizerModel extends Callable<[string[]], string[]> { * @param tokens The tokens to encode. * @returns The encoded tokens. */ - _call(tokens: string[]): string[] { + _call(tokens: Array<[string, [number,number]]>): Array<[string, [number,number]]> { let result = this.encode(tokens); if (this.fuse_unk) { result = fuse_unk(result, this.tokens_to_ids, this.unk_token_id); @@ -50,7 +50,7 @@ abstract class TokenizerModel extends Callable<[string[]], string[]> { * @param tokens The tokens to encode. * @returns The encoded tokens. */ - abstract encode(tokens: string[]): string[]; + abstract encode(tokens: Array<[string, [number,number]]>): Array<[string, [number, number]]>; } export default TokenizerModel; diff --git a/src/core/preTokenizer/BertPreTokenizer.ts b/src/core/preTokenizer/BertPreTokenizer.ts index e7ae867..d43e5eb 100644 --- a/src/core/preTokenizer/BertPreTokenizer.ts +++ b/src/core/preTokenizer/BertPreTokenizer.ts @@ -30,8 +30,10 @@ class BertPreTokenizer extends PreTokenizer { * @param options Additional options for the pre-tokenization logic. * @returns An array of tokens. */ - pre_tokenize_text(text: string, options?: any): string[] { - return text.trim().match(this.pattern) || []; + + pre_tokenize_text(text: string, options?: any): Array<[string, [number, number]]> { + const trimOffset = text.length - text.trimStart().length; + return [...text.trim().matchAll(this.pattern)].map(m => [m[0],[trimOffset + m.index!, trimOffset + m.index! + m[0].length]]); } } diff --git a/src/core/preTokenizer/ByteLevel.ts b/src/core/preTokenizer/ByteLevel.ts index 0e866f1..9cced8e 100644 --- a/src/core/preTokenizer/ByteLevel.ts +++ b/src/core/preTokenizer/ByteLevel.ts @@ -54,22 +54,31 @@ class ByteLevel extends PreTokenizer { * @param options Additional options for the pre-tokenization logic. * @returns An array of tokens. */ - pre_tokenize_text(text: string, options?: any): string[] { - // Add a leading space if the option is enabled - if (this.add_prefix_space && !text.startsWith(" ")) { + pre_tokenize_text(text: string, options?: any): Array<[string, [number, number]]> { + // Track whether we insert a synthetic space so we can correct span positions + const prefixInserted = this.add_prefix_space && !text.startsWith(" "); + if (prefixInserted) { text = " " + text; } - // Split on whitespace and punctuation - const tokens = this.use_regex ? text.match(this.pattern) || [] : [text]; + // Capture raw token strings with their positions in the (possibly prefixed) text + const rawTokens: Array<[string, number]> = this.use_regex + ? [...text.matchAll(this.pattern)].map((m) => [m[0], m.index!]) + : [[text, 0]]; + + // Offset converts positions in the prefixed text back to the original input + const offset = prefixInserted ? -1 : 0; // Maps all our bytes to unicode strings, avoiding control tokens of the BPE (spaces in our case) - return tokens.map((token) => - Array.from( + return rawTokens.map(([token, index]) => { + const start = Math.max(0, index + offset); + const end = Math.max(0, index + token.length + offset); + const encoded = Array.from( this.text_encoder.encode(token), (byte) => this.byte_encoder[byte], - ).join(""), - ); + ).join(""); + return [encoded, [start, end]]; + }); } } diff --git a/src/core/preTokenizer/Digits.ts b/src/core/preTokenizer/Digits.ts index 886f860..3525482 100644 --- a/src/core/preTokenizer/Digits.ts +++ b/src/core/preTokenizer/Digits.ts @@ -26,8 +26,8 @@ class Digits extends PreTokenizer { * @param text The text to tokenize. * @returns An array of tokens. */ - pre_tokenize_text(text: string): string[] { - return text.match(this.pattern) || []; + pre_tokenize_text(text: string): Array<[string, [number, number]]> { + return [...text.matchAll(this.pattern)].map(m => [m[0], [m.index!, m.index! + m[0].length]]); } } diff --git a/src/core/preTokenizer/FixedLength.ts b/src/core/preTokenizer/FixedLength.ts index 6427666..0ecef80 100644 --- a/src/core/preTokenizer/FixedLength.ts +++ b/src/core/preTokenizer/FixedLength.ts @@ -22,10 +22,10 @@ class FixedLength extends PreTokenizer { * @param text The text to be pre-tokenized. * @returns An array of tokens produced by splitting the input text into fixed-length tokens. */ - pre_tokenize_text(text: string): string[] { - const tokens = []; + pre_tokenize_text(text: string): Array<[string, [number, number]]> { + const tokens: Array<[string, [number, number]]> = []; for (let i = 0; i < text.length; i += this._length) { - tokens.push(text.slice(i, i + this._length)); + tokens.push([text.slice(i, i + this._length), [i, i + this._length]]); } return tokens; } diff --git a/src/core/preTokenizer/Metaspace.ts b/src/core/preTokenizer/Metaspace.ts index 4bc8fb6..b741f2a 100644 --- a/src/core/preTokenizer/Metaspace.ts +++ b/src/core/preTokenizer/Metaspace.ts @@ -36,7 +36,7 @@ class Metaspace extends PreTokenizer { * @param options The options for the pre-tokenization. * @returns A new list of pre-tokenized tokens. */ - pre_tokenize_text(text: string, options?: PreTokenizeTextOptions): string[] { + pre_tokenize_text(text: string, options?: PreTokenizeTextOptions): Array<[string, [number,number]]> { const { section_index = undefined } = options ?? {}; let normalized = text.replaceAll(" ", this.str_rep); @@ -52,7 +52,7 @@ class Metaspace extends PreTokenizer { ) { normalized = this.str_rep + normalized; } - return [normalized]; + return [[normalized,[0,text.length]]]; } } diff --git a/src/core/preTokenizer/Punctuation.ts b/src/core/preTokenizer/Punctuation.ts index edc1bf2..34f315c 100644 --- a/src/core/preTokenizer/Punctuation.ts +++ b/src/core/preTokenizer/Punctuation.ts @@ -27,8 +27,8 @@ class Punctuation extends PreTokenizer { * @param text The text to tokenize. * @returns An array of tokens. */ - pre_tokenize_text(text: string): string[] { - return text.match(this.pattern) || []; + pre_tokenize_text(text: string): Array<[string, [number, number]]> { + return [...text.matchAll(this.pattern)].map(m => [m[0], [m.index!, m.index! + m[0].length]]); } } diff --git a/src/core/preTokenizer/Replace.ts b/src/core/preTokenizer/Replace.ts index c862655..22bc927 100644 --- a/src/core/preTokenizer/Replace.ts +++ b/src/core/preTokenizer/Replace.ts @@ -24,11 +24,12 @@ class Replace extends PreTokenizer { * @param text The text to be pre-tokenized. * @returns An array of tokens produced by replacing certain characters. */ - pre_tokenize_text(text: string): string[] { + pre_tokenize_text(text: string): Array<[string, [number, number]]> { + const span: [number, number] = [0, text.length]; if (this.pattern === null) { - return [text]; + return [[text, span]]; } - return [text.replaceAll(this.pattern, this.config.content ?? "")]; + return [[text.replaceAll(this.pattern, this.config.content ?? ""), span]]; } } diff --git a/src/core/preTokenizer/Sequence.ts b/src/core/preTokenizer/Sequence.ts index 6e79b7c..f0fef64 100644 --- a/src/core/preTokenizer/Sequence.ts +++ b/src/core/preTokenizer/Sequence.ts @@ -26,15 +26,18 @@ class Sequence extends PreTokenizer { * @param options Additional options for the pre-tokenization logic. * @returns The pre-tokenized text. */ - pre_tokenize_text(text: string, options?: any): string[] { - // Use reduce to apply each tokenizer to the text + pre_tokenize_text(text: string, options?: any): Array<[string, [number, number]]> { return this.tokenizers.reduce( - (pre_tokenized_text, tokenizer) => { - return tokenizer - ? tokenizer.pre_tokenize(pre_tokenized_text, options) - : pre_tokenized_text; + (pairs, tokenizer) => { + if (!tokenizer) return pairs; + return pairs.flatMap(([word, [start]]) => + tokenizer.pre_tokenize(word, options).map(([subWord, [subStart, subEnd]]) => [ + subWord, + [subStart + start, subEnd + start] as [number, number], + ]), + ); }, - [text] as string[], + [[text, [0, text.length]]] as Array<[string, [number, number]]>, ); } } diff --git a/src/core/preTokenizer/Split.ts b/src/core/preTokenizer/Split.ts index 07732bd..a742eb8 100644 --- a/src/core/preTokenizer/Split.ts +++ b/src/core/preTokenizer/Split.ts @@ -29,15 +29,26 @@ class Split extends PreTokenizer { * @param text The text to tokenize. * @returns An array of tokens. */ - pre_tokenize_text(text: string): string[] { + pre_tokenize_text(text: string): Array<[string, [number, number]]> { if (this.pattern === null) { return []; } if (this.config.invert) { - return text.match(this.pattern) || []; + return [...text.matchAll(this.pattern)].map((m) => [m[0], [m.index!, m.index! + m[0].length]]); } else if (this.config.behavior?.toLowerCase() === "removed") { - return text.split(this.pattern).filter((x) => x); + const result: Array<[string, [number, number]]> = []; + let prev = 0; + for (const match of text.matchAll(this.pattern)) { + if (prev < match.index!) { + result.push([text.slice(prev, match.index!), [prev, match.index!]]); + } + prev = match.index! + match[0].length; + } + if (prev < text.length) { + result.push([text.slice(prev), [prev, text.length]]); + } + return result; } else { return regex_split(text, this.pattern); } diff --git a/src/core/preTokenizer/Whitespace.ts b/src/core/preTokenizer/Whitespace.ts index 6997b07..796a623 100644 --- a/src/core/preTokenizer/Whitespace.ts +++ b/src/core/preTokenizer/Whitespace.ts @@ -10,8 +10,8 @@ class Whitespace extends PreTokenizer { * @param options Additional options for the pre-tokenization logic. * @returns An array of tokens produced by splitting the input text on whitespace. */ - pre_tokenize_text(text: string, options?: any): string[] { - return text.match(/\w+|[^\w\s]+/g) || []; + pre_tokenize_text(text: string, options?: any): Array<[string, [number, number]]> { + return [...text.matchAll(/\w+|[^\w\s]+/g)].map(m => [m[0], [m.index!, m.index! + m[0].length]]); } } diff --git a/src/core/preTokenizer/WhitespaceSplit.ts b/src/core/preTokenizer/WhitespaceSplit.ts index df93e6b..bb85eeb 100644 --- a/src/core/preTokenizer/WhitespaceSplit.ts +++ b/src/core/preTokenizer/WhitespaceSplit.ts @@ -10,7 +10,7 @@ class WhitespaceSplit extends PreTokenizer { * @param text The text to be pre-tokenized. * @returns An array of tokens produced by splitting the input text on whitespace. */ - pre_tokenize_text(text: string): string[] { + pre_tokenize_text(text: string): Array<[string, [number, number]]> { return whitespace_split(text); } } diff --git a/src/core/tokenizerModelImplementations/BPE.ts b/src/core/tokenizerModelImplementations/BPE.ts index 66b6abd..564c859 100644 --- a/src/core/tokenizerModelImplementations/BPE.ts +++ b/src/core/tokenizerModelImplementations/BPE.ts @@ -249,33 +249,57 @@ class BPE extends TokenizerModel { * @param tokens The input sequence of tokens to encode. * @returns The resulting subword tokens after applying the BPE algorithm to the input sequence of tokens. */ - encode(tokens: string[]): string[] { - const output_tokens: string[] = []; + encode(tokens: Array<[string, [number, number]]>): Array<[string, [number, number]]> { + const output_tokens: Array<[string, [number, number]]> = []; - for (const token of tokens) { + for (const [token, [word_start, word_end]] of tokens) { if (this.ignore_merges && this.tokens_to_ids.has(token)) { - output_tokens.push(token); + output_tokens.push([token, [word_start, word_end]]); continue; } + const bpe_token_list = this.bpe(token); - for (const t of bpe_token_list) { + // Walk left-to-right through bpe_token_list to recover character positions. + // Each BPE entry is a contiguous slice of the original word (plus decorating + // suffixes that have no corresponding characters), so a cursor gives us spans. + let pos = 0; + for (let i = 0; i < bpe_token_list.length; i++) { + const t = bpe_token_list[i]; + const is_last = i === bpe_token_list.length - 1; + + // Strip decorating suffixes to find the characters this token actually covers. + let core = t; + if (!is_last && this.continuing_subword_suffix) { + core = core.slice(0, core.length - this.continuing_subword_suffix.length); + } + if (is_last && this.end_of_word_suffix) { + core = core.slice(0, core.length - this.end_of_word_suffix.length); + } + const chars_covered = core.length; + const span: [number, number] = [word_start + pos, word_start + pos + chars_covered]; + pos += chars_covered; + if (this.tokens_to_ids.has(t)) { - output_tokens.push(t); + output_tokens.push([t, span]); } else if (this.byte_fallback) { - const byte_tokens = Array.from(this.text_encoder!.encode(t)).map( + const byte_tokens = Array.from(this.text_encoder!.encode(core)).map( (x) => `<0x${x.toString(16).toUpperCase().padStart(2, "0")}>`, ); if (byte_tokens.every((x) => this.tokens_to_ids.has(x))) { // Ensure the byte tokens are actually in the vocabulary, otherwise // we fall back to the unknown token. For more information, see // https://github.com/huggingface/transformers/issues/28096. - output_tokens.push(...byte_tokens); + // All byte tokens share the same span — they collectively represent + // the characters at that position; individual bytes can't be subdivided. + for (const bt of byte_tokens) { + output_tokens.push([bt, span]); + } } else if (this.unk_token != null) { - output_tokens.push(this.unk_token); + output_tokens.push([this.unk_token, span]); } } else if (this.unk_token != null) { - output_tokens.push(this.unk_token); + output_tokens.push([this.unk_token, span]); } } } diff --git a/src/core/tokenizerModelImplementations/Legacy.ts b/src/core/tokenizerModelImplementations/Legacy.ts index 65cddf5..6bc7acb 100644 --- a/src/core/tokenizerModelImplementations/Legacy.ts +++ b/src/core/tokenizerModelImplementations/Legacy.ts @@ -58,7 +58,7 @@ class Legacy extends TokenizerModel { } } - encode(tokens: string[]): string[] { + encode(tokens: Array<[string, [number, number]]>): Array<[string, [number, number]]> { return tokens; } } diff --git a/src/core/tokenizerModelImplementations/Unigram.ts b/src/core/tokenizerModelImplementations/Unigram.ts index 6c6e70f..8f2ab82 100644 --- a/src/core/tokenizerModelImplementations/Unigram.ts +++ b/src/core/tokenizerModelImplementations/Unigram.ts @@ -111,11 +111,14 @@ class Unigram extends TokenizerModel { * @param tokens The tokens to encode. * @returns An array of encoded tokens. */ - encode(tokens: string[]): string[] { - const to_return: string[] = []; - for (const token of tokens) { - const tokenized = this.tokenize(token); - to_return.push(...tokenized); + encode(tokens: Array<[string, [number, number]]>): Array<[string, [number, number]]> { + const to_return: Array<[string, [number, number]]> = []; + for (const [token, [word_start]] of tokens) { + const lattice = new TokenLattice(token, this.bos_token_id, this.eos_token_id); + this.populate_nodes(lattice); + for (const [subWord, [subStart, subEnd]] of lattice.token_spans()) { + to_return.push([subWord, [word_start + subStart, word_start + subEnd]]); + } } return to_return; } diff --git a/src/core/tokenizerModelImplementations/WordPiece.ts b/src/core/tokenizerModelImplementations/WordPiece.ts index f59aabb..48aa961 100644 --- a/src/core/tokenizerModelImplementations/WordPiece.ts +++ b/src/core/tokenizerModelImplementations/WordPiece.ts @@ -32,22 +32,23 @@ class WordPieceTokenizer extends TokenizerModel { * @param tokens The tokens to encode. * @returns An array of encoded tokens. */ - encode(tokens: string[]): string[] { - const output_tokens: string[] = []; - for (const token of tokens) { + encode(tokens: Array<[string, [number, number]]>): Array<[string, [number, number]]> { + const output_tokens: Array<[string, [number, number]]> = []; + for (const [token, [word_start, word_end]] of tokens) { const chars = [...token]; if (chars.length > this.max_input_chars_per_word) { - output_tokens.push(this.unk_token!); + output_tokens.push([this.unk_token!, [word_start, word_end]]); continue; } let is_unknown = false; let start = 0; - const sub_tokens: string[] = []; + const sub_tokens: Array<[string, [number, number]]> = []; while (start < chars.length) { let end = chars.length; let current_substring: string | null = null; + const sub_start = start; while (start < end) { let substr = chars.slice(start, end).join(""); @@ -65,11 +66,11 @@ class WordPieceTokenizer extends TokenizerModel { is_unknown = true; break; } - sub_tokens.push(current_substring); + sub_tokens.push([current_substring, [word_start + sub_start, word_start + end]]); start = end; } if (is_unknown) { - output_tokens.push(this.unk_token!); + output_tokens.push([this.unk_token!, [word_start, word_end]]); } else { output_tokens.push(...sub_tokens); } diff --git a/src/static/types.ts b/src/static/types.ts index 197273a..4379243 100644 --- a/src/static/types.ts +++ b/src/static/types.ts @@ -1,6 +1,7 @@ export interface Encoding { ids: number[]; tokens: string[]; + offsets: Array<[number,number]>; attention_mask: number[]; token_type_ids?: number[]; } diff --git a/src/utils/core.ts b/src/utils/core.ts index 455392d..9dbce11 100644 --- a/src/utils/core.ts +++ b/src/utils/core.ts @@ -106,31 +106,101 @@ export const escape_reg_exp = (string: string): string => * Helper function to fuse consecutive unknown tokens. */ export const fuse_unk = ( - arr: Array, + arr: Array<[string, [number, number]]>, tokens_to_ids: Map, unk_token_id: number, -) => { - const fused = []; +): Array<[string, [number, number]]> => { + const fused: Array<[string, [number, number]]> = []; let i = 0; while (i < arr.length) { fused.push(arr[i]); - const token_id = tokens_to_ids.get(arr[i]) ?? unk_token_id; + const token_id = tokens_to_ids.get(arr[i][0]) ?? unk_token_id; if (token_id !== unk_token_id) { ++i; continue; } while ( ++i < arr.length && - (tokens_to_ids.get(arr[i]) ?? unk_token_id) === unk_token_id + (tokens_to_ids.get(arr[i][0]) ?? unk_token_id) === unk_token_id ) { - if (tokens_to_ids.get(fused.at(-1)) !== unk_token_id) { - fused[fused.length - 1] += arr[i]; + if (tokens_to_ids.get(fused.at(-1)![0]) !== unk_token_id) { + const last = fused[fused.length - 1]; + fused[fused.length - 1] = [last[0] + arr[i][0], [last[1][0], arr[i][1][1]]]; } } } return fused; }; +/** + * Builds an alignment map from normalized-string positions back to original-string positions. + * map[i] gives the index in `original` that normalized character i came from. + * Handles 1-to-1 (lowercase), 1-to-many (NFD expansion), many-to-1 (NFC compression), + * and many-to-0 (deletion, e.g. accent stripping) — all monotone transformations. + */ +export function build_alignment_map(original: string, normalized: string): number[] { + const map = new Array(normalized.length).fill(0); + let orig_i = 0; + let norm_i = 0; + + while (norm_i < normalized.length) { + if (orig_i >= original.length) { + map[norm_i++] = original.length; + continue; + } + + const oc = original[orig_i]; + const nc = normalized[norm_i]; + + // Case 1: direct match or lowercase match (1-to-1) + if (oc === nc || oc.toLowerCase() === nc) { + map[norm_i++] = orig_i++; + continue; + } + + let matched = false; + + // Case 2: expansion — oc decomposes into multiple norm chars (NFD/NFKD) + for (const form of ["NFD", "NFKD"] as const) { + for (const candidate of [oc.normalize(form), oc.toLowerCase().normalize(form)]) { + if (candidate.length > 1 && normalized.slice(norm_i, norm_i + candidate.length) === candidate) { + for (let k = 0; k < candidate.length; k++) map[norm_i + k] = orig_i; + norm_i += candidate.length; + orig_i++; + matched = true; + break; + } + } + if (matched) break; + } + if (matched) continue; + + // Case 3: contraction — multiple orig chars compose into fewer norm chars (NFC/NFKC) + for (let look = 2; look <= 4 && !matched; look++) { + if (orig_i + look > original.length) break; + const slice = original.slice(orig_i, orig_i + look); + for (const form of ["NFC", "NFKC"] as const) { + for (const candidate of [slice.normalize(form), slice.normalize(form).toLowerCase()]) { + if (normalized.slice(norm_i, norm_i + candidate.length) === candidate) { + for (let k = 0; k < candidate.length; k++) map[norm_i + k] = orig_i; + norm_i += candidate.length; + orig_i += look; + matched = true; + break; + } + } + if (matched) break; + } + } + if (matched) continue; + + // Case 4: deletion — oc was removed by normalization (e.g. combining mark stripped) + orig_i++; + } + + return map; +} + export const is_chinese_char = (cp: number): boolean => (cp >= 0x4e00 && cp <= 0x9fff) || (cp >= 0x3400 && cp <= 0x4dbf) || @@ -181,21 +251,21 @@ export const object_to_map = (obj: Object): Map => * @param regex The regex to split on. * @returns The split string. */ -export const regex_split = (text: string, regex: RegExp): string[] => { - const result: string[] = []; +export const regex_split = (text: string, regex: RegExp): Array<[string, [number, number]]> => { + const result: Array<[string, [number, number]]> = []; let prev = 0; for (const match of text.matchAll(regex)) { const full_match = match[0]; if (prev < match.index!) { - result.push(text.slice(prev, match.index)); + result.push([text.slice(prev, match.index!), [prev, match.index!]]); } if (full_match.length > 0) { - result.push(full_match); + result.push([full_match, [match.index!, match.index! + full_match.length]]); } prev = match.index! + full_match.length; } if (prev < text.length) { - result.push(text.slice(prev)); + result.push([text.slice(prev), [prev, text.length]]); } return result; }; @@ -226,5 +296,6 @@ export const validate_object = ( * @param {string} text The text to split. * @returns {string[]} The split string. */ -export const whitespace_split = (text: string): Array => - text.match(/\S+/g) || []; +export const whitespace_split = (text: string): Array<[string, [number,number]]> => + [...text.matchAll(/\S+/g)].map(m => [m[0], [m.index!, m.index! + m[0].length]]); + diff --git a/src/utils/data-structures/DictionarySplitter.ts b/src/utils/data-structures/DictionarySplitter.ts index b57205a..23afdcc 100644 --- a/src/utils/data-structures/DictionarySplitter.ts +++ b/src/utils/data-structures/DictionarySplitter.ts @@ -44,8 +44,8 @@ class DictionarySplitter { * @param text The input text to split. * @returns An array of tokens. */ - split(text: string): string[] { - const result: string[] = []; + split(text: string): Array<[string, number]> { + const result: Array<[string, number]> = []; const n = text.length; let start = 0; let i = 0; @@ -65,9 +65,9 @@ class DictionarySplitter { if (match) { if (i > start) { - result.push(text.slice(start, i)); + result.push([text.slice(start, i),start]); } - result.push(match); + result.push([match,i]); i += match.length; start = i; } else { @@ -75,7 +75,7 @@ class DictionarySplitter { } } if (start < n) { - result.push(text.slice(start)); + result.push([text.slice(start),start]); } return result; } diff --git a/src/utils/data-structures/TokenLattice.ts b/src/utils/data-structures/TokenLattice.ts index 6fb96a4..532cdb9 100644 --- a/src/utils/data-structures/TokenLattice.ts +++ b/src/utils/data-structures/TokenLattice.ts @@ -182,6 +182,14 @@ class TokenLattice { return nodes.map((x) => this.piece(x)); } + /** + * @returns The most likely sequence of tokens with their character spans. + */ + token_spans(): Array<[string, [number, number]]> { + const nodes = this.viterbi(); + return nodes.map((x) => [this.piece(x), [x.pos, x.pos + x.length]]); + } + /** * @returns The most likely sequence of token ids. */ diff --git a/tests/bundle.test.ts b/tests/bundle.test.ts index e39b62b..06e040a 100644 --- a/tests/bundle.test.ts +++ b/tests/bundle.test.ts @@ -25,7 +25,8 @@ const TARGET_OUTPUT = `[ '▁Hello', '▁World' ] { ids: [ 1, 15043, 2787 ], tokens: [ '', '▁Hello', '▁World' ], - attention_mask: [ 1, 1, 1 ] + attention_mask: [ 1, 1, 1 ], + offsets: [ [ 0, 0 ], [ 11, 11 ], [ 11, 11 ] ] } Hello World `; diff --git a/tests/edge-cases/offsets.test.ts b/tests/edge-cases/offsets.test.ts new file mode 100644 index 0000000..a5d614b --- /dev/null +++ b/tests/edge-cases/offsets.test.ts @@ -0,0 +1,209 @@ +import fetchConfigById from "../utils/fetchConfigById"; +import { Tokenizer } from "../../src"; + +// ─── Why each case is dangerous for offset tracking ──────────────────────────── +// +// 1. EMPTY STRING +// encode_text("") → DictionarySplitter.split("") → [] (n=0, while-loop never runs). +// all_pairs is empty, so every token in the post-processor output mismatches the +// pair walk and receives [0, 0]. The risk: callers that iterate over offsets and +// subtract start from end to get a character count would get 0, which is correct, +// but callers that assume at least one real span exists would read a stale pair_i. +// +// 2. WHITESPACE-ONLY STRING +// encode_text(" ") → DictionarySplitter gives one non-empty section " ". +// BertNormalizer produces " " (clean_text maps whitespace to ASCII space). +// BertPreTokenizer: text.trim() = "", matchAll on "" → []. Model receives []. +// processed_text is non-empty (length=3 > 0) so the early-exit at line 299 doesn't +// fire — the code reaches the model and just produces nothing. all_pairs = []. +// Risk: a naive check `if (!pairs1.length) early-return` would be wrong; the code +// must allow the post-processor to still run (and produce its sentinels). +// +// 3. ACCENTED INPUT (NFC) — "cafe\u0301" stripped to "cafe" +// Two Unicode forms of the same word produce different alignment maps and +// therefore different end-offsets for the same logical token: +// +// NFC "café" (4 chars, precomposed): +// BertNormalizer: NFD(é)="é", strip Mn → "e". Result: "cafe" (4 chars). +// build_alignment_map("café", "cafe"): +// 'c','a','f' match via Case 1. Then orig='é', norm='e': +// Case 1 fails ('é' !== 'e', toLowerCase same). +// Case 2: NFD('é')="é" (2 chars) not found in normalized "cafe" at [3..5]. +// Case 3: orig_i+2=5 > length=4, break. +// Case 4 DELETION: orig_i++ (skips 'é'), norm_i stays. +// Next iter: orig_i=4 >= length=4 → map[3] = 4. +// alignment=[0,1,2,4]. to_orig(0,4): sec_end=alignment[4] OOB → original.length=4 → [0,4]. +// +// NFD "cafe\u0301" (5 chars, decomposed — must use \u escape, V8 normalizes literals): +// BertNormalizer: NFD("cafe\u0301")="cafe\u0301", strip Mn → "cafe" (4 chars). +// build_alignment_map("cafe\u0301", "cafe"): +// 'c','a','f': Case 1. Then orig='e', norm='e': Case 1 MATCHES → map[3]=3. +// Loop exits at norm_i=4; combining accent at orig[4] is NEVER consumed. +// alignment=[0,1,2,3]. to_orig(0,4): sec_end=alignment[4] OOB → original.length=5 → [0,5]. +// +// Same semantic text, same token "cafe", but NFC→[0,4] and NFD→[0,5]. +// +// 4. TEXT PAIR — B-sequence offsets are independent; [SEP] between gets [0, 0] +// BertProcessing pair: [CLS] A [SEP] B [SEP] (one separator, unlike RoBERTa's two). +// all_pairs = pairs1 ++ pairs2. tokenize_helper's offset walk: +// [CLS] mismatches all_pairs[0] → [0, 0]. +// A tokens match consecutively. +// [SEP] mismatches all_pairs[pair_i] which is the first B token → [0, 0]. +// B tokens match. Their offsets come from encode_text(text_pair) independently: +// positions are relative to text_pair, not to text. +// Final [SEP] → pair_i exhausted → [0, 0]. +// +// 5. SPECIAL TOKEN EMBEDDED IN TEXT — "hello [SEP] world" with BERT +// DictionarySplitter (unnormalized pass) recognises "[SEP]" as an added token. +// sections = [("hello ", 0), ("[SEP]", 6), (" world", 11)]. +// encode_text returns: +// ("hello", [0, 5]) — from the "hello " section +// ("[SEP]", [6, 11]) — added_tokens_map hit, span = [offset, offset+len] +// ("world", [12, 17]) — from the " world" section, section_offset=11 +// all_pairs has all three entries. Post-processor: [CLS] hello [SEP] world [SEP]. +// Offset walk (sequential string-match): +// [CLS] → mismatch all_pairs[0]="hello" → [0, 0] +// hello → match → [0, 5] +// [SEP] → all_pairs[1][0]="[SEP]" MATCHES → [6, 11] ← real offset +// world → match → [12, 17] +// [SEP] → pair_i=3 = all_pairs.length → [0, 0] ← sentinel +// The in-text [SEP] is distinguishable from the post-processor [SEP] only by +// sequential position in the token stream, not by any flag on the token string. + +describe("Offset edge cases (BERT)", () => { + let uncased: Tokenizer; + + beforeAll(async () => { + const { tokenizerJson, tokenizerConfig } = + await fetchConfigById("Xenova/bert-base-uncased"); + uncased = new Tokenizer(tokenizerJson, tokenizerConfig); + }); + + // ── 1. Empty string ─────────────────────────────────────────────────────────── + + test("empty string — only sentinels, both [0, 0]", () => { + // DictionarySplitter.split("") → [] → all_pairs = []. Every post-processor + // token mismatches the pair walk → all offsets are the [0, 0] sentinel. + const { tokens, offsets } = uncased.encode(""); + expect(tokens).toEqual(["[CLS]", "[SEP]"]); + expect(offsets).toEqual([[0, 0], [0, 0]]); + }); + + // ── 2. Whitespace-only string ───────────────────────────────────────────────── + + test("whitespace-only string — BertPreTokenizer produces nothing, sentinels only", () => { + // encode_text(" ") reaches the model with an empty token list because + // BertPreTokenizer's text.trim() = "" — but processed_text.length=3 is non-zero + // so the early-exit guard doesn't fire. all_pairs = []. + const { tokens, offsets } = uncased.encode(" "); + expect(tokens).toEqual(["[CLS]", "[SEP]"]); + expect(offsets).toEqual([[0, 0], [0, 0]]); + }); + + // ── 3. Accented input ───────────────────────────────────────────────────────── + + test("NFC accent — Case 4 deletion fires; whole-token offset correctly spans original", () => { + // "café" (4 chars, precomposed e-acute). BertNormalizer strips accent → "cafe". + // build_alignment_map: 'é' fails Cases 1-3 → Case 4 deletion (orig_i++). + // Next iter: orig_i=4 >= length=4 → map[3] = original.length = 4. + // alignment = [0, 1, 2, 4]. + // to_orig(0, 4): alignment[4] OOB → sec_end = original_section.length = 4 → [0, 4]. + // The full token correctly spans the original word even though 'e-acute' couldn't + // be individually aligned — only the whole-word span is reliable when Case 4 fires. + const { tokens, offsets } = uncased.encode("café"); // NFC precomposed + expect(tokens).toEqual(["[CLS]", "cafe", "[SEP]"]); + expect(offsets).toEqual([ + [0, 0], + [0, 4], // whole "cafe" maps back to original [0,4] = "café" + [0, 0], + ]); + }); + + test("NFD accent — combining mark un-consumed; end overflows to original.length=5", () => { + // "cafe\u0301" (5 chars: c,a,f,e, combining-acute). Must use \u escape — V8 + // normalizes source-file NFD string literals to NFC at parse time. + // + // BertNormalizer: strip Mn removes ́ → "cafe" (4 chars). + // build_alignment_map("cafe\u0301", "cafe"): + // 'c','a','f': Case 1, orig_i=3, norm_i=3. + // orig[3]='e', norm[3]='e': Case 1 MATCHES → map[3]=3, orig_i=4, norm_i=4. + // Loop exits. Combining accent at orig[4] is NEVER consumed. + // alignment = [0, 1, 2, 3]. + // to_orig(0, 4): alignment[4] OOB → sec_end = original_section.length = 5 → [0, 5]. + // NFC and NFD of the same word produce different end-offsets: [0,4] vs [0,5]. + const nfd = "cafe\u0301"; // explicit NFD: 5 JS chars + const { tokens, offsets } = uncased.encode(nfd); + expect(tokens).toEqual(["[CLS]", "cafe", "[SEP]"]); + expect(offsets).toEqual([ + [0, 0], + [0, 5], // end=5: spans 'e' (orig[3]) AND the unconsumed combining accent (orig[4]) + [0, 0], + ]); + }); + + // ── 4. Text pair ────────────────────────────────────────────────────────────── + + test("text pair — B spans are independent of A; single [SEP] separator gets [0, 0]", () => { + // Structure: [CLS] A [SEP] B [SEP] + // The [SEP] between A and B is injected by BertProcessing — not in all_pairs. + // When the offset walk hits it, all_pairs[pair_i] = first-B-token != "[SEP]" → [0, 0]. + // B-sequence offsets come from encode_text("world") independently: + // "world" → [0, 5] relative to text_pair, not concatenated after text. + const enc = uncased.encode("hello", { + text_pair: "world", + return_token_type_ids: true, + }); + expect(enc.tokens).toEqual(["[CLS]", "hello", "[SEP]", "world", "[SEP]"]); + expect(enc.offsets).toEqual([ + [0, 0], // [CLS] + [0, 5], // hello — from encode_text("hello") + [0, 0], // [SEP] separator — post-processor sentinel + [0, 5], // world — from encode_text("world"), independent of A's positions + [0, 0], // [SEP] — post-processor sentinel + ]); + // A-side tokens → type_id 0; B-side tokens → type_id 1. + expect(enc.token_type_ids).toEqual([0, 0, 0, 1, 1]); + }); + + // ── 5. Special token embedded in text ──────────────────────────────────────── + + test("[SEP] embedded in text gets its real character span", () => { + // "hello [SEP] world" (length 17). + // DictionarySplitter (unnormalized) splits on "[SEP]": + // sections = [("hello ", 0), ("[SEP]", 6), (" world", 11)] + // encode_text emits: + // ("hello", [0, 5]) — from "hello " section + // ("[SEP]", [6, 11]) — added_tokens_map hit; span = [offset, offset+len] + // ("world", [12, 17]) — section_offset=11, "world" at position 1 in " world" + // Post-processor: [CLS] hello [SEP] world [SEP] + // Offset walk (sequential string-match, not by token type): + // [CLS] → mismatch all_pairs[0]="hello" → [0, 0] + // hello → match → [0, 5] + // [SEP] → all_pairs[1][0]="[SEP]" MATCHES → [6, 11] ← real span from text + // world → match → [12, 17] + // [SEP] → pair_i exhausted → [0, 0] ← post-processor sentinel + const { tokens, offsets } = uncased.encode("hello [SEP] world"); + expect(tokens).toEqual(["[CLS]", "hello", "[SEP]", "world", "[SEP]"]); + expect(offsets).toEqual([ + [0, 0], // [CLS] + [0, 5], // hello + [6, 11], // [SEP] from text — real span, not sentinel + [12, 17], // world + [0, 0], // [SEP] from post-processor — sentinel + ]); + }); + + test("[SEP] embedded in text with add_special_tokens:false — all spans are real", () => { + // Without special tokens, the post-processor does not wrap, so the output is + // exactly the three pairs from encode_text — all in all_pairs, all matching. + const { tokens, offsets } = uncased.encode("hello [SEP] world", { + add_special_tokens: false, + }); + expect(tokens).toEqual(["hello", "[SEP]", "world"]); + expect(offsets).toEqual([ + [0, 5], + [6, 11], + [12, 17], + ]); + }); +}); diff --git a/tests/models/bert/offsets.test.ts b/tests/models/bert/offsets.test.ts new file mode 100644 index 0000000..3fa3559 --- /dev/null +++ b/tests/models/bert/offsets.test.ts @@ -0,0 +1,37 @@ +// tests/models/bert/offsets.test.ts +import fetchConfigById from "../../utils/fetchConfigById"; +import { Tokenizer } from "../../../src"; + +describe("BERT offset mapping", () => { + let tokenizer: Tokenizer; + + beforeAll(async () => { + const { tokenizerJson, tokenizerConfig } = + await fetchConfigById("Xenova/bert-base-uncased"); + tokenizer = new Tokenizer(tokenizerJson, tokenizerConfig); + }); + + test("Hello World — special tokens get [0,0], words get character spans", () => { + const { tokens, offsets } = tokenizer.encode("Hello World"); + expect(tokens).toEqual(["[CLS]", "hello", "world", "[SEP]"]); + expect(offsets).toEqual([[0, 0], [0, 5], [6, 11], [0, 0]]); + }); + + test("text pair — B-sequence offsets are independent of A", () => { + const { tokens, offsets } = tokenizer.encode("hello", { + text_pair: "world", + }); + // [CLS] hello [SEP] world [SEP] + expect(offsets).toEqual([[0, 0], [0, 5], [0, 0], [0, 5], [0, 0]]); + }); + + test("subword split preserves character-level spans", async () => { + // bert-base-cased: "Héllo" → H + ##é + ##llo + const { tokenizerJson: cJson, tokenizerConfig: cCfg } = + await fetchConfigById("Xenova/bert-base-cased"); + const cased = new Tokenizer(cJson, cCfg); + const { offsets } = cased.encode("Héllo"); + // [CLS] H ##é ##llo [SEP] + expect(offsets).toEqual([[0,0], [0,1], [1,2], [2,5], [0,0]]); + }); +}); diff --git a/tests/models/gpt2/gpt2.test.ts b/tests/models/gpt2/gpt2.test.ts new file mode 100644 index 0000000..34c5171 --- /dev/null +++ b/tests/models/gpt2/gpt2.test.ts @@ -0,0 +1,50 @@ +import fetchConfigById from "../../utils/fetchConfigById"; +import { Tokenizer } from "../../../src"; + +describe("GPT-2 offset mapping (BPE + ByteLevel)", () => { + let tokenizer: Tokenizer; + + beforeAll(async () => { + const { tokenizerJson, tokenizerConfig } = + await fetchConfigById("Xenova/gpt2"); + tokenizer = new Tokenizer(tokenizerJson, tokenizerConfig); + }); + + test("space is absorbed into the following token's span, not between spans", () => { + // The GPT-2 regex uses ' ?\p{L}+', so the space at index 5 is consumed by + // the second match — ĠWorld covers [5, 11], not [6, 11]. + const { tokens, offsets } = tokenizer.encode("Hello World"); + expect(tokens).toEqual(["Hello", "ĠWorld"]); + expect(offsets).toEqual([ + [0, 5], // "Hello" + [5, 11], // "ĠWorld" — span includes the space at index 5 + ]); + }); + + test("no special tokens means no [0,0] sentinel offsets", () => { + // Unlike BERT, GPT-2 has no CLS/SEP; every offset is a real character span. + const { tokens, offsets } = tokenizer.encode("Hello World"); + expect(offsets).toHaveLength(tokens.length); + expect(offsets.every(([s, e]) => s !== 0 || e !== 0)).toBe(true); + }); + + test("BPE sub-split: each piece gets its own sub-span within the chunk", () => { + // "trailing space " splits into three pre-tokenizer chunks: + // "trailing" at [0, 8] → BPE: ["tra"(3 chars), "iling"(5 chars)] + // " space" at [8, 14] → BPE: ["Ġspace"] (single token) + // " " at [14, 17] → BPE: ["Ġ", "Ġ", "Ġ"] (one space each) + // + // BPE.encode walks the sub-token list with a char cursor, giving each piece + // [word_start + pos, word_start + pos + piece_length]. + const { tokens, offsets } = tokenizer.encode("trailing space "); + expect(tokens).toEqual(["tra", "iling", "Ġspace", "Ġ", "Ġ", "Ġ"]); + expect(offsets).toEqual([ + [0, 3], // "tra" + [3, 8], // "iling" + [8, 14], // "Ġspace" + [14, 15], // first trailing space + [15, 16], // second trailing space + [16, 17], // third trailing space + ]); + }); +}); diff --git a/tests/models/roberta/offsets.test.ts b/tests/models/roberta/offsets.test.ts new file mode 100644 index 0000000..6e49304 --- /dev/null +++ b/tests/models/roberta/offsets.test.ts @@ -0,0 +1,143 @@ +import fetchConfigById from "../../utils/fetchConfigById"; +import { Tokenizer } from "../../../src"; + +// RoBERTa pipeline: (no normalizer) → ByteLevel → BPE → RobertaProcessing +// +// How RoBERTa differs from GPT-2 in offset terms: +// +// 1. SPECIAL TOKENS — [0, 0] sentinels +// RobertaProcessing wraps every output with and . These tokens are +// injected by the post-processor *after* encode_text runs, so they are never +// present in `all_pairs`. tokenize_helper assigns offsets by walking `all_pairs` +// in order and matching on token string; any token that doesn't match the next +// pair entry falls through to [0, 0]. always mismatches (it's first, before +// any content token), and always mismatches (content is exhausted). +// +// 2. TEXT PAIRS — double separator, all sentinels +// RobertaProcessing follows the convention: +// A B +// The two tokens between A and B are both injected by the post-processor. +// When tokenize_helper scans that region, all_pairs[pair_i] is already pointing +// at the first B token ("good"), so neither matches → both get [0, 0]. +// The closing after B also gets [0, 0] (pair_i has advanced past all pairs). +// +// 3. SPACE ABSORPTION — identical to GPT-2 +// ByteLevel uses the same GPT-2 regex (' ?\p{L}+' etc.), so the space before +// the second word is consumed by that match and included in its span start. +// Note: ByteLevel.trim_offsets is set to true in the config but is currently +// marked @todo in the implementation — offsets are NOT trimmed, so the space +// remains absorbed into the following token's span start. +// +// 4. TOKEN TYPE IDs (text pairs only) +// RobertaProcessing fills type_id=0 for all A-side tokens (including and +// the closing of A), and type_id=1 for the middle , all B tokens, and +// the final . + +describe("RoBERTa offset mapping (BPE + ByteLevel + RobertaProcessing)", () => { + let tokenizer: Tokenizer; + + beforeAll(async () => { + const { tokenizerJson, tokenizerConfig } = + await fetchConfigById("Xenova/all-distilroberta-v1"); + tokenizer = new Tokenizer(tokenizerJson, tokenizerConfig); + }); + + test(" and get [0,0]; content tokens keep ByteLevel spans", () => { + // ByteLevel regex on "Hello world": + // match "Hello" at index 0 → span [0, 5] + // match " world" at index 5 → span [5, 11] (space absorbed into start) + // BPE keeps both as single tokens (no sub-split for these words). + // RobertaProcessing: Hello Ġworld + // → not in all_pairs → [0, 0] + // Hello → all_pairs[0] match → [0, 5] + // Ġworld → all_pairs[1] match → [5, 11] + // → pair_i past end → [0, 0] + const { tokens, offsets } = tokenizer.encode("Hello world"); + expect(tokens).toEqual(["", "Hello", "Ġworld", ""]); + expect(offsets).toEqual([ + [0, 0], // — post-processor sentinel + [0, 5], // Hello + [5, 11], // Ġworld — span starts at 5, absorbing the space + [0, 0], // — post-processor sentinel + ]); + }); + + test("add_special_tokens:false strips both sentinels and their [0,0] offsets", () => { + // Without special tokens the post-processor is still called but wrapping is skipped, + // so the output is identical to raw ByteLevel+BPE — same as GPT-2 for this input. + const { tokens, offsets } = tokenizer.encode("Hello world", { + add_special_tokens: false, + }); + expect(tokens).toEqual(["Hello", "Ġworld"]); + expect(offsets).toEqual([ + [0, 5], + [5, 11], + ]); + }); + + test("BPE sub-split: each piece gets its own sub-span, sentinels wrap", () => { + // BPE splits "tokenization" into "token" (5 chars) and "ization" (7 chars). + // BPE.encode walks a cursor: pos=0 → "token" [0,5]; pos=5 → "ization" [5,12]. + // word_start=0 (ByteLevel match starts at 0 for no-prefix-space input). + // RobertaProcessing wraps: token ization + const { tokens, offsets } = tokenizer.encode("tokenization"); + expect(tokens).toEqual(["", "token", "ization", ""]); + expect(offsets).toEqual([ + [0, 0], + [0, 5], // "token" + [5, 12], // "ization" + [0, 0], + ]); + }); + + test("text pair — double separator; both separator tokens get [0,0]", () => { + // Single sequence: Hello Ġworld + // Pair sequence: Hello Ġworld good Ġmorning + // + // all_pairs = [Hello,[0,5]], [Ġworld,[5,11]], [good,[0,4]], [Ġmorning,[4,12]] + // + // tokenize_helper walk: + // → mismatch all_pairs[0]="Hello" → [0, 0] + // Hello → match all_pairs[0] → [0, 5], pair_i=1 + // Ġworld → match all_pairs[1] → [5, 11], pair_i=2 + // → mismatch all_pairs[2]="good" → [0, 0] ← end-of-A separator + // → mismatch all_pairs[2]="good" → [0, 0] ← start-of-B separator + // good → match all_pairs[2] → [0, 4], pair_i=3 + // Ġmorning → match all_pairs[3] → [4, 12], pair_i=4 + // → pair_i == all_pairs.length → [0, 0] ← end-of-B sentinel + const enc = tokenizer.encode("Hello world", { + text_pair: "good morning", + return_token_type_ids: true, + }); + expect(enc.tokens).toEqual([ + "", "Hello", "Ġworld", "", + "", "good", "Ġmorning", "", + ]); + expect(enc.offsets).toEqual([ + [0, 0], // + [0, 5], // Hello + [5, 11], // Ġworld + [0, 0], // end-of-A + [0, 0], // start-of-B ← this is the extra RoBERTa separator + [0, 4], // good + [4, 12], // Ġmorning + [0, 0], // end-of-B + ]); + // A-side (including and first ) → type_id 0 + // B-side (including both middle and final ) → type_id 1 + expect(enc.token_type_ids).toEqual([0, 0, 0, 0, 1, 1, 1, 1]); + }); + + test("pair offset independence — B-sequence spans are relative to their own text, not A", () => { + // Sanity check: "good morning" encoded alone vs as pair-B produces the same offsets. + // Offsets reference the *original string passed for that sequence*, not a concatenated buffer. + const alone = tokenizer.encode("good morning", { add_special_tokens: false }); + const pair = tokenizer.encode("Hello world", { + text_pair: "good morning", + add_special_tokens: false, + }); + // pair = Hello Ġworld good Ġmorning (no sentinels) + const pairBOffsets = pair.offsets.slice(2); // drop Hello, Ġworld + expect(alone.offsets).toEqual(pairBOffsets); + }); +}); diff --git a/tests/models/t5/offsets.test.ts b/tests/models/t5/offsets.test.ts new file mode 100644 index 0000000..932350c --- /dev/null +++ b/tests/models/t5/offsets.test.ts @@ -0,0 +1,124 @@ +import fetchConfigById from "../../utils/fetchConfigById"; +import { Tokenizer } from "../../../src"; + +// T5 pipeline: Precompiled normalizer → Sequence[WhitespaceSplit, Metaspace] → Unigram → TemplateProcessing (adds ) +// +// Key behaviours that shape every offset: +// +// 1. WhitespaceSplit splits on whitespace *before* Metaspace runs, so no space +// character ever enters the Metaspace replacement logic for T5. Instead, +// Metaspace prepends ▁ (U+2581) to every resulting chunk unconditionally +// (prepend_scheme defaults to "always"). +// +// 2. The Unigram lattice is built over the ▁-prefixed string (e.g. "▁Hello", +// 6 chars). The model returns node.pos and node.pos+node.length spans within +// *that* string, then shifts them by word_start (the pre-token's start in +// processed_text). This means the full-token span in processed_text is +// [word_start, word_start + len("▁Hello")] = [0, 6], which reaches *one +// character past* the end of the original word — landing on the following +// space (or the end of text). to_orig then maps that processed-text index +// through the alignment into the original character span. +// +// 3. The Precompiled (SentencePiece) normalizer may recompose decomposed Unicode. +// build_alignment_map tracks these multi-char → single-char contractions so +// that even NFD input produces offsets anchored in the original raw string. +// +// 4. Special tokens added by the post-processor () are not present in the +// pre-model pair list and receive the [0, 0] sentinel offset. + +describe("T5 offset mapping (Unigram + Metaspace + Precompiled normalizer)", () => { + let tokenizer: Tokenizer; + + beforeAll(async () => { + const { tokenizerJson, tokenizerConfig } = + await fetchConfigById("google-t5/t5-small"); + tokenizer = new Tokenizer(tokenizerJson, tokenizerConfig); + }); + + test("Hello world — ▁ expansion pushes first token's end into the following space", () => { + // WhitespaceSplit gives ["Hello",[0,5]] and ["world",[6,11]]. + // Metaspace prepends ▁ to each → "▁Hello" (6 chars) and "▁world" (6 chars). + // Unigram keeps both as single tokens: lattice positions [0,6] within each. + // "▁Hello": word_start=0, processed span [0,6]. + // to_orig(0,6) → alignment[0]=0, alignment[6]=6 → [0,6]. + // Original[0..6) = "Hello " — the ▁'s extra char absorbs the whitespace. + // "▁world": word_start=6, processed span [6,12]. + // to_orig(6,12) → alignment[6]=6, 12 >= len("Hello world")=11 → end=11 → [6,11]. + // "" is injected by the post-processor and gets [0,0]. + const { tokens, offsets } = tokenizer.encode("Hello world"); + expect(tokens).toEqual(["▁Hello", "▁world", ""]); + expect(offsets).toEqual([ + [0, 6], // "▁Hello" — end 6 reaches into the space, not 5 + [6, 11], // "▁world" + [0, 0], // — post-processor sentinel + ]); + }); + + test("subword split — each Unigram lattice node maps independently to the original span", () => { + // "tokenization" → WhitespaceSplit gives one chunk at [0,12]. + // Metaspace: "▁tokenization" (13 chars), word_start=0. + // Unigram splits at best Viterbi path: "▁token" (pos=0,len=6) + "ization" (pos=6,len=7). + // "▁token": processed span [0+0, 0+6] = [0,6] → to_orig → [0,6] + // "ization": processed span [0+6, 0+13] = [6,13] → to_orig → alignment[6]=6, + // 13 >= len("tokenization")=12 → end=12 → [6,12] + const { tokens, offsets } = tokenizer.encode("tokenization"); + expect(tokens).toEqual(["▁token", "ization", ""]); + expect(offsets).toEqual([ + [0, 6], // "▁token" + [6, 12], // "ization" + [0, 0], + ]); + }); + + test("café NFC — alignment is identity; accent codepoint inside a single char", () => { + // NFC "café" (U+00E9 for é, length 4 in JS). + // Normalizer leaves NFC input unchanged. WhitespaceSplit gives ["café",[0,4]]. + // Metaspace: "▁café" (5 chars), word_start=0. + // Unigram keeps it as one token: processed span [0,5]. + // to_orig(0,5) → alignment[0]=0, 5 >= len("café")=4 → end=4 → [0,4]. + const nfc = "café"; // precomposed é + const { tokens, offsets } = tokenizer.encode(nfc); + expect(tokens).toEqual(["▁café", ""]); + expect(offsets).toEqual([ + [0, 4], // single NFC codepoint, 4 chars total + [0, 0], + ]); + }); + + test("café NFD — alignment map contracts decomposed e+combining-accent back to original span", () => { + // NFD "café" (U+0065 + U+0301, length 5 in JS). + // The Precompiled normalizer recomposes to NFC "café" (length 4). + // build_alignment_map Case 3 (contraction): "é" → "é". + // alignment = [0, 1, 2, 3, 3] (positions 3 and 4 in NFD both map to orig-3). + // Unigram token "▁café" spans processed [0,5]. + // to_orig(0,5) → alignment[0]=0; 5 >= len(normalized)=4 → end = len(original)=5 → [0,5]. + // End=5 spans both the 'e' (orig-3) and the combining accent (orig-4). + const nfd = "café"; // decomposed é + const { tokens, offsets } = tokenizer.encode(nfd); + expect(tokens).toEqual(["▁café", ""]); + expect(offsets).toEqual([ + [0, 5], // spans all 5 original chars including the combining mark + [0, 0], + ]); + }); + + test("leading-space input — ▁ not prepended to first word (already starts with ▁ equivalent); trailing subword gets zero-width span", () => { + // " already spaced" — WhitespaceSplit gives ["already",[1,8]], ["spaced",[9,15]]. + // Metaspace prepends ▁ to each (neither starts with ▁ already). + // Unigram on "▁already" (8 chars, word_start=1): single token, processed span [1,9]. + // to_orig(1,9) → [1,9] = "already " (absorbs the space at orig-8). + // Unigram on "▁spaced" (7 chars, word_start=9): splits "▁space"(6) + "d"(1). + // "▁space": processed [9,15] → to_orig(9,15) → alignment[9]=9, end=15 → [9,15]. + // "d": processed [15,16] → to_orig(15,16) → both 15 and 16 >= len=15 → [15,15]. + // "d" gets a zero-width [15,15] because its lattice position overshoots the + // original text length after the ▁ expansion offset. + const { tokens, offsets } = tokenizer.encode(" already spaced"); + expect(tokens).toEqual(["▁already", "▁space", "d", ""]); + expect(offsets).toEqual([ + [1, 9], // "already " in original (absorbs the trailing space) + [9, 15], // "spaced" in original + [15, 15], // "d" — zero-width, ▁ expansion pushed its start past original's end + [0, 0], + ]); + }); +}); diff --git a/tests/tokenizers.test.ts b/tests/tokenizers.test.ts index 9f65bbf..f5b045a 100644 --- a/tests/tokenizers.test.ts +++ b/tests/tokenizers.test.ts @@ -290,7 +290,7 @@ describe("Tokenizer methods", () => { const text = `${bos_token}abc ${added_token} ${normalized_special_token} ${unnormalized_special_token} xyz${eos_token}`; const encoded = tokenizer.encode(text); - expect(encoded).toEqual({ + expect(encoded).toMatchObject({ ids: [1, 4, 8, 7, 10, 11, 4, 12, 4, 4, 0, 0, 0, 2], tokens: ["", " ", "ab", "c", " ", " ", " ", "", " ", " ", "", "", "", ""], attention_mask: [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], From cad8d96b941b305c5caedfe2085cb07ae79c7959 Mon Sep 17 00:00:00 2001 From: Anika Jain Date: Fri, 12 Jun 2026 17:47:12 +0530 Subject: [PATCH 2/2] fix for CI failures --- tests/utils/fetchConfigById.ts | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/tests/utils/fetchConfigById.ts b/tests/utils/fetchConfigById.ts index 937c8eb..2c7a818 100644 --- a/tests/utils/fetchConfigById.ts +++ b/tests/utils/fetchConfigById.ts @@ -27,8 +27,15 @@ const fetchConfigById = async ( const remoteUrl = `https://huggingface.co/${modelId}/resolve/main/tokenizer.json`; const remoteUrlConfig = `https://huggingface.co/${modelId}/resolve/main/tokenizer_config.json`; + const headers: Record = { "User-Agent": "tokenizers.js-tests/1.0" }; + if (process.env.HF_TOKEN) headers["Authorization"] = `Bearer ${process.env.HF_TOKEN}`; + const loadJson = async (url: string) => { - const response = await fetch(url); + const response = await fetch(url, { headers }); + if (!response.ok) { + const text = await response.text(); + throw new Error(`Failed to fetch ${url} (HTTP ${response.status}): ${text.slice(0, 300)}`); + } return await response.json(); };