diff --git a/packages/transformers/src/tokenization_utils.js b/packages/transformers/src/tokenization_utils.js index 29de6f186..15549d508 100644 --- a/packages/transformers/src/tokenization_utils.js +++ b/packages/transformers/src/tokenization_utils.js @@ -126,6 +126,60 @@ function truncateHelper(item, length) { } } +/** + * Computes character offset mapping from token strings. + * @param {string} text The original (unnormalized) text. + * @param {string[]} tokens The list of token strings from encoding. + * @param {Set} special_tokens_set Set of special token strings. + * @returns {[number, number][]} Array of [start, end] character offsets. + */ +function computeOffsetMapping(text, tokens, special_tokens_set) { + const offsets = /** @type {[number, number][]} */ ([]); + const textLower = text.toLowerCase(); + let pos = 0; + for (const token of tokens) { + if (special_tokens_set.has(token)) { + offsets.push([0, 0]); + continue; + } + let actual_text = token; + let preceded_by_space = false; + if (token.startsWith('##')) { + // BERT WordPiece: "##ing" means this subword directly follows the previous (no space) + actual_text = token.slice(2); + } else if (token.startsWith('\u0120')) { + // GPT-2/RoBERTa BPE: "Ġword" means there was a space before "word" in the original + actual_text = token.slice(1); + preceded_by_space = true; + } else if (token.startsWith('\u2581')) { + // SentencePiece: "▁word" means word boundary (space in original) + actual_text = token.slice(1); + preceded_by_space = true; + } + if (actual_text.length === 0) { + offsets.push([pos, pos]); + continue; + } + // If this token was preceded by a space in the original, skip past spaces + if (preceded_by_space) { + while (pos < text.length && (text[pos] === ' ' || text[pos] === '\n' || text[pos] === '\t')) { + pos++; + } + } + // Search forward from pos (case-insensitive, since some tokenizers lowercase) + const tokenLower = actual_text.toLowerCase(); + const idx = textLower.indexOf(tokenLower, pos); + + if (idx === -1) { + offsets.push([0, 0]); // fallback + } else { + offsets.push([idx, idx + actual_text.length]); + pos = idx + actual_text.length; + } + } + return offsets; +} + /** * Returns the value of the first matching key in the tokenizer config object. * @param {Object} config The tokenizer config object. @@ -184,6 +238,7 @@ function getSpecialTokens(tokenizer) { * @property {TItem} input_ids List of token ids to be fed to a model. * @property {TItem} attention_mask List of indices specifying which tokens should be attended to by the model. * @property {TItem} [token_type_ids] List of token type ids to be fed to a model. + * @property {[number, number][][]|[number, number][]} [offset_mapping] Character offsets for each token. */ /** @@ -197,6 +252,11 @@ function getSpecialTokens(tokenizer) { * @property {number|null} [max_length=null] Maximum length of the returned list and optionally padding length. * @property {TReturnTensor} [return_tensor=true] Whether to return the results as Tensors or arrays. * @property {boolean|null} [return_token_type_ids=null] Whether to return the token type ids. + * @property {boolean} [return_offsets_mapping=false] Whether to return character-level offset mappings for each token. + * Each entry is a `[start, end]` pair pointing to the token's span in the original string. + * Special tokens (e.g. `[CLS]`, `[SEP]`, `[PAD]`) always map to `[0, 0]`. + * Note: accuracy may be reduced for byte-level BPE tokenizers (e.g. GPT-2) with non-ASCII input, + * and offsets for `text_pair` tokens are not supported. */ /** @@ -359,7 +419,13 @@ export class PreTrainedTokenizer text, options = {}, ) { - const { text_pair = null, add_special_tokens = true, padding = false, return_token_type_ids = null } = options; + const { + text_pair = null, + add_special_tokens = true, + padding = false, + return_token_type_ids = null, + return_offsets_mapping = false, + } = options; let { truncation = null, max_length = null } = options; const return_tensor = /** @type {TReturnTensor} */ (options.return_tensor ?? true); // Different to HF @@ -380,10 +446,17 @@ export class PreTrainedTokenizer } encodedTokens = text.map((t, i) => - this._encode_plus(t, { text_pair: text_pair[i], add_special_tokens, return_token_type_ids }), + this._encode_plus(t, { + text_pair: text_pair[i], + add_special_tokens, + return_token_type_ids, + return_offsets_mapping, + }), ); } else { - encodedTokens = text.map((x) => this._encode_plus(x, { add_special_tokens, return_token_type_ids })); + encodedTokens = text.map((x) => + this._encode_plus(x, { add_special_tokens, return_token_type_ids, return_offsets_mapping }), + ); } } else { if (text === null || text === undefined) { @@ -397,7 +470,14 @@ export class PreTrainedTokenizer } // For single input, we just wrap in an array, and then unwrap later. - encodedTokens = [this._encode_plus(text, { text_pair, add_special_tokens, return_token_type_ids })]; + encodedTokens = [ + this._encode_plus(text, { + text_pair, + add_special_tokens, + return_token_type_ids, + return_offsets_mapping, + }), + ]; } // At this point, `encodedTokens` is batched, of shape [batch_size, tokens]. // However, array may be jagged. So, we may need pad to max_length. @@ -444,7 +524,11 @@ export class PreTrainedTokenizer padHelper( encodedTokens[i], max_length, - (key) => (key === 'input_ids' ? this.pad_token_id : 0), + (key) => { + if (key === 'input_ids') return this.pad_token_id; + if (key === 'offset_mapping') return [0, 0]; + return 0; + }, this.padding_side, ); } @@ -482,12 +566,19 @@ export class PreTrainedTokenizer const dims = [encodedTokens.length, encodedTokens[0].input_ids.length]; for (const key of Object.keys(encodedTokens[0])) { + if (key === 'offset_mapping') continue; result[key] = new Tensor( 'int64', BigInt64Array.from(encodedTokens.flatMap((x) => x[key]).map(BigInt)), dims, ); } + if ('offset_mapping' in encodedTokens[0]) { + result.offset_mapping = encodedTokens.map((x) => x.offset_mapping); + if (!isBatched) { + result.offset_mapping = result.offset_mapping[0]; + } + } } else { for (const key of Object.keys(encodedTokens[0])) { result[key] = encodedTokens.map((x) => x[key]); @@ -524,11 +615,20 @@ export class PreTrainedTokenizer * @param {string|null} [options.text_pair=null] The optional second text to encode. * @param {boolean} [options.add_special_tokens=true] Whether or not to add the special tokens associated with the corresponding model. * @param {boolean|null} [options.return_token_type_ids=null] Whether to return token_type_ids. - * @returns {{input_ids: number[], attention_mask: number[], token_type_ids?: number[]}} An object containing the encoded text. + * @param {boolean} [options.return_offsets_mapping=false] Whether to return offset_mapping + * @returns {{input_ids: number[], attention_mask: number[], token_type_ids?: number[], offset_mapping?: [number,number][]}} An object containing the encoded text. * @private */ - _encode_plus(text, { text_pair = null, add_special_tokens = true, return_token_type_ids = null } = {}) { - const { ids, attention_mask, token_type_ids } = this._tokenizer.encode(text, { + _encode_plus( + text, + { + text_pair = null, + add_special_tokens = true, + return_token_type_ids = null, + return_offsets_mapping = false, + } = {}, + ) { + const { ids, tokens, attention_mask, token_type_ids } = this._tokenizer.encode(text, { text_pair, add_special_tokens, return_token_type_ids: return_token_type_ids ?? this.return_token_type_ids, @@ -537,6 +637,9 @@ export class PreTrainedTokenizer input_ids: ids, attention_mask, ...(token_type_ids ? { token_type_ids } : {}), + ...(return_offsets_mapping + ? { offset_mapping: computeOffsetMapping(text, tokens, new Set(this.all_special_tokens)) } + : {}), }; } diff --git a/packages/transformers/tests/tokenizers.test.js b/packages/transformers/tests/tokenizers.test.js index 1e7977ba9..b2cb10bae 100644 --- a/packages/transformers/tests/tokenizers.test.js +++ b/packages/transformers/tests/tokenizers.test.js @@ -726,3 +726,137 @@ describe("Chat templates", () => { } }); }); + +describe("Offset mapping", () => { + let tokenizer; + beforeAll(async () => { + tokenizer = await AutoTokenizer.from_pretrained("Xenova/bert-base-uncased"); + }, MAX_TOKENIZER_LOAD_TIME); + + it("does not include offset_mapping in output by default", () => { + const output = tokenizer("Hello world", { return_tensor: false }); + expect(output.offset_mapping).toBeUndefined(); + }); + + it("returns correct offsets for a single string", () => { + // "Hello world": H(0)e(1)l(2)l(3)o(4) (5)w(6)o(7)r(8)l(9)d(10) + // Tokens: [CLS] hello world [SEP] + // BERT lowercases, but offsets point to positions in the original text. + const { offset_mapping } = tokenizer("Hello world", { + return_tensor: false, + return_offsets_mapping: true, + }); + expect(offset_mapping).toEqual([ + [0, 0], // [CLS] — special token, no character span + [0, 5], // hello → "Hello" + [6, 11], // world → "world" + [0, 0], // [SEP] — special token + ]); + }); + + it("returns correct offsets for subword tokens (BERT WordPiece ## prefix)", () => { + // "tokenization" is the canonical WordPiece example from the BERT paper. + // BERT-base-uncased splits it as: token + ##ization + // t(0)o(1)k(2)e(3)n(4)i(5)z(6)a(7)t(8)i(9)o(10)n(11) + const { offset_mapping } = tokenizer("tokenization", { + return_tensor: false, + return_offsets_mapping: true, + }); + expect(offset_mapping).toEqual([ + [0, 0], // [CLS] + [0, 5], // token → "token" + [5, 12], // ##ization → "ization" (continues from where "token" ended) + [0, 0], // [SEP] + ]); + }); + + it("returns correct offsets for batched input (jagged, return_tensor=false)", () => { + // "a" → [CLS] a [SEP] + // "b c" → [CLS] b c [SEP] + const { offset_mapping } = tokenizer(["a", "b c"], { + return_tensor: false, + return_offsets_mapping: true, + }); + expect(offset_mapping).toEqual([ + [ + [0, 0], + [0, 1], + [0, 0], + ], + [ + [0, 0], + [0, 1], + [2, 3], + [0, 0], + ], + ]); + }); + + it("pads offset_mapping with [0, 0] for padding tokens", () => { + // "a" → [CLS] a [SEP] (3 tokens, padded to 4) + // "b c" → [CLS] b c [SEP] (4 tokens, no padding needed) + const { offset_mapping } = tokenizer(["a", "b c"], { + return_tensor: false, + padding: true, + return_offsets_mapping: true, + }); + expect(offset_mapping).toEqual([ + [ + [0, 0], + [0, 1], + [0, 0], + [0, 0], + ], // last [0,0] is the [PAD] token + [ + [0, 0], + [0, 1], + [2, 3], + [0, 0], + ], + ]); + }); + + it("truncates offset_mapping to match truncated token sequence", () => { + // "b c" without special tokens → [b, c], truncated to max_length=1 → [b] + const { offset_mapping } = tokenizer("b c", { + return_tensor: false, + truncation: true, + max_length: 1, + add_special_tokens: false, + return_offsets_mapping: true, + }); + expect(offset_mapping).toEqual([[0, 1]]); + }); + + it("returns offset_mapping as a plain array even when return_tensor=true", () => { + // All other outputs (input_ids, attention_mask, token_type_ids) become Tensors. + // offset_mapping must stay a plain JS array because it contains nested [start, end] pairs + // that cannot be flattened into a 1D BigInt64Array. + const output = tokenizer(["a", "a"], { + padding: true, + truncation: true, + return_offsets_mapping: true, + }); + + // Confirm the other fields are Tensors (they have a tolist() method). + expect(output.input_ids.tolist()).toEqual([ + [101n, 1037n, 102n], + [101n, 1037n, 102n], + ]); + + // offset_mapping is a plain nested array, never a Tensor. + expect(Array.isArray(output.offset_mapping)).toBe(true); + expect(output.offset_mapping).toEqual([ + [ + [0, 0], + [0, 1], + [0, 0], + ], // "a" → [CLS] a [SEP] + [ + [0, 0], + [0, 1], + [0, 0], + ], // "a" → [CLS] a [SEP] + ]); + }); +});