From fa94faa1d8ae046105d39894b44ff78f1cb0a3bc Mon Sep 17 00:00:00 2001 From: Anika Jain Date: Mon, 8 Jun 2026 05:01:24 +0530 Subject: [PATCH 1/4] added return_offset_mapping support to PreTrainedTokenizer --- .../transformers/src/tokenization_utils.js | 88 ++++++++++++-- .../transformers/tests/tokenizers.test.js | 107 ++++++++++++++++++ 2 files changed, 187 insertions(+), 8 deletions(-) diff --git a/packages/transformers/src/tokenization_utils.js b/packages/transformers/src/tokenization_utils.js index 29de6f186..a8e47787e 100644 --- a/packages/transformers/src/tokenization_utils.js +++ b/packages/transformers/src/tokenization_utils.js @@ -126,6 +126,63 @@ 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 +241,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 +255,7 @@ 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_offset_mapping=false] Whether to return character offset mappings. */ /** @@ -359,7 +418,7 @@ 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_offset_mapping = false } = options; let { truncation = null, max_length = null } = options; const return_tensor = /** @type {TReturnTensor} */ (options.return_tensor ?? true); // Different to HF @@ -380,10 +439,10 @@ 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_offset_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_offset_mapping })); } } else { if (text === null || text === undefined) { @@ -397,7 +456,7 @@ 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_offset_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 +503,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 +545,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 +594,12 @@ 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_offset_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_offset_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 +608,7 @@ export class PreTrainedTokenizer input_ids: ids, attention_mask, ...(token_type_ids ? { token_type_ids } : {}), + ...(return_offset_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..05d9f25f2 100644 --- a/packages/transformers/tests/tokenizers.test.js +++ b/packages/transformers/tests/tokenizers.test.js @@ -726,3 +726,110 @@ 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_offset_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_offset_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_offset_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_offset_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_offset_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_offset_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] + ]); + }); +}); From d9db4ec691a7c76959e84a34b03f083ee63daf2d Mon Sep 17 00:00:00 2001 From: Anika Jain Date: Mon, 8 Jun 2026 05:18:24 +0530 Subject: [PATCH 2/4] Add JSDoc limitations note to return_offset_mapping option --- packages/transformers/src/tokenization_utils.js | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/packages/transformers/src/tokenization_utils.js b/packages/transformers/src/tokenization_utils.js index a8e47787e..b38692866 100644 --- a/packages/transformers/src/tokenization_utils.js +++ b/packages/transformers/src/tokenization_utils.js @@ -255,7 +255,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_offset_mapping=false] Whether to return character offset mappings. + * @property {boolean} [return_offset_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. */ /** From 24ffdfef602f89afd9a0686a9779afae5137f427 Mon Sep 17 00:00:00 2001 From: Anika Jain Date: Mon, 8 Jun 2026 05:23:56 +0530 Subject: [PATCH 3/4] auto-fix formatting --- .../transformers/src/tokenization_utils.js | 67 +++++++++++++------ .../transformers/tests/tokenizers.test.js | 55 +++++++++++---- 2 files changed, 88 insertions(+), 34 deletions(-) diff --git a/packages/transformers/src/tokenization_utils.js b/packages/transformers/src/tokenization_utils.js index b38692866..94298c755 100644 --- a/packages/transformers/src/tokenization_utils.js +++ b/packages/transformers/src/tokenization_utils.js @@ -134,25 +134,25 @@ function truncateHelper(item, length) { * @returns {[number, number][]} Array of [start, end] character offsets. */ function computeOffsetMapping(text, tokens, special_tokens_set) { - const offsets = /** @type {[number, number][]} */([]) + 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]); + 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('##')){ + 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')){ + } 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) + actual_text = token.slice(1); preceded_by_space = true; - }else if(token.startsWith('\u2581')){ - // SentencePiece: "▁word" means word boundary (space in original) + } else if (token.startsWith('\u2581')) { + // SentencePiece: "▁word" means word boundary (space in original) actual_text = token.slice(1); preceded_by_space = true; } @@ -176,13 +176,10 @@ function computeOffsetMapping(text, tokens, special_tokens_set) { 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. @@ -422,7 +419,13 @@ export class PreTrainedTokenizer text, options = {}, ) { - const { text_pair = null, add_special_tokens = true, padding = false, return_token_type_ids = null,return_offset_mapping = false } = options; + const { + text_pair = null, + add_special_tokens = true, + padding = false, + return_token_type_ids = null, + return_offset_mapping = false, + } = options; let { truncation = null, max_length = null } = options; const return_tensor = /** @type {TReturnTensor} */ (options.return_tensor ?? true); // Different to HF @@ -443,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, return_offset_mapping }), + this._encode_plus(t, { + text_pair: text_pair[i], + add_special_tokens, + return_token_type_ids, + return_offset_mapping, + }), ); } else { - encodedTokens = text.map((x) => this._encode_plus(x, { add_special_tokens, return_token_type_ids, return_offset_mapping })); + encodedTokens = text.map((x) => + this._encode_plus(x, { add_special_tokens, return_token_type_ids, return_offset_mapping }), + ); } } else { if (text === null || text === undefined) { @@ -460,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, return_offset_mapping })]; + encodedTokens = [ + this._encode_plus(text, { + text_pair, + add_special_tokens, + return_token_type_ids, + return_offset_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. @@ -509,7 +526,7 @@ export class PreTrainedTokenizer max_length, (key) => { if (key === 'input_ids') return this.pad_token_id; - if (key === 'offset_mapping') return [0,0]; + if (key === 'offset_mapping') return [0, 0]; return 0; }, this.padding_side, @@ -602,8 +619,16 @@ export class PreTrainedTokenizer * @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,return_offset_mapping = false } = {}) { - const { ids,tokens, 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_offset_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, @@ -612,7 +637,9 @@ export class PreTrainedTokenizer input_ids: ids, attention_mask, ...(token_type_ids ? { token_type_ids } : {}), - ...(return_offset_mapping ? {offset_mapping: computeOffsetMapping(text,tokens,new Set(this.all_special_tokens))}:{}), + ...(return_offset_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 05d9f25f2..3f48b04dc 100644 --- a/packages/transformers/tests/tokenizers.test.js +++ b/packages/transformers/tests/tokenizers.test.js @@ -747,10 +747,10 @@ describe("Offset mapping", () => { return_offset_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 + [0, 0], // [CLS] — special token, no character span + [0, 5], // hello → "Hello" + [6, 11], // world → "world" + [0, 0], // [SEP] — special token ]); }); @@ -763,10 +763,10 @@ describe("Offset mapping", () => { return_offset_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] + [0, 0], // [CLS] + [0, 5], // token → "token" + [5, 12], // ##ization → "ization" (continues from where "token" ended) + [0, 0], // [SEP] ]); }); @@ -778,8 +778,17 @@ describe("Offset mapping", () => { return_offset_mapping: true, }); expect(offset_mapping).toEqual([ - [[0, 0], [0, 1], [0, 0]], - [[0, 0], [0, 1], [2, 3], [0, 0]], + [ + [0, 0], + [0, 1], + [0, 0], + ], + [ + [0, 0], + [0, 1], + [2, 3], + [0, 0], + ], ]); }); @@ -792,8 +801,18 @@ describe("Offset mapping", () => { return_offset_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]], + [ + [0, 0], + [0, 1], + [0, 0], + [0, 0], + ], // last [0,0] is the [PAD] token + [ + [0, 0], + [0, 1], + [2, 3], + [0, 0], + ], ]); }); @@ -828,8 +847,16 @@ describe("Offset mapping", () => { // 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] + [ + [0, 0], + [0, 1], + [0, 0], + ], // "a" → [CLS] a [SEP] + [ + [0, 0], + [0, 1], + [0, 0], + ], // "a" → [CLS] a [SEP] ]); }); }); From 15146b9597f33c8f7b6a126bc7827d0c8be65573 Mon Sep 17 00:00:00 2001 From: Anika Jain Date: Mon, 8 Jun 2026 23:43:55 +0530 Subject: [PATCH 4/4] spell_error --- packages/transformers/src/tokenization_utils.js | 16 ++++++++-------- packages/transformers/tests/tokenizers.test.js | 12 ++++++------ 2 files changed, 14 insertions(+), 14 deletions(-) diff --git a/packages/transformers/src/tokenization_utils.js b/packages/transformers/src/tokenization_utils.js index 94298c755..15549d508 100644 --- a/packages/transformers/src/tokenization_utils.js +++ b/packages/transformers/src/tokenization_utils.js @@ -252,7 +252,7 @@ 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_offset_mapping=false] Whether to return character-level offset mappings for each token. + * @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, @@ -424,7 +424,7 @@ export class PreTrainedTokenizer add_special_tokens = true, padding = false, return_token_type_ids = null, - return_offset_mapping = false, + return_offsets_mapping = false, } = options; let { truncation = null, max_length = null } = options; const return_tensor = /** @type {TReturnTensor} */ (options.return_tensor ?? true); // Different to HF @@ -450,12 +450,12 @@ export class PreTrainedTokenizer text_pair: text_pair[i], add_special_tokens, return_token_type_ids, - return_offset_mapping, + return_offsets_mapping, }), ); } else { encodedTokens = text.map((x) => - this._encode_plus(x, { add_special_tokens, return_token_type_ids, return_offset_mapping }), + this._encode_plus(x, { add_special_tokens, return_token_type_ids, return_offsets_mapping }), ); } } else { @@ -475,7 +475,7 @@ export class PreTrainedTokenizer text_pair, add_special_tokens, return_token_type_ids, - return_offset_mapping, + return_offsets_mapping, }), ]; } @@ -615,7 +615,7 @@ 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. - * @param {boolean} [options.return_offset_mapping=false] Whether to return offset_mapping + * @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 */ @@ -625,7 +625,7 @@ export class PreTrainedTokenizer text_pair = null, add_special_tokens = true, return_token_type_ids = null, - return_offset_mapping = false, + return_offsets_mapping = false, } = {}, ) { const { ids, tokens, attention_mask, token_type_ids } = this._tokenizer.encode(text, { @@ -637,7 +637,7 @@ export class PreTrainedTokenizer input_ids: ids, attention_mask, ...(token_type_ids ? { token_type_ids } : {}), - ...(return_offset_mapping + ...(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 3f48b04dc..b2cb10bae 100644 --- a/packages/transformers/tests/tokenizers.test.js +++ b/packages/transformers/tests/tokenizers.test.js @@ -744,7 +744,7 @@ describe("Offset mapping", () => { // BERT lowercases, but offsets point to positions in the original text. const { offset_mapping } = tokenizer("Hello world", { return_tensor: false, - return_offset_mapping: true, + return_offsets_mapping: true, }); expect(offset_mapping).toEqual([ [0, 0], // [CLS] — special token, no character span @@ -760,7 +760,7 @@ describe("Offset mapping", () => { // 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_offset_mapping: true, + return_offsets_mapping: true, }); expect(offset_mapping).toEqual([ [0, 0], // [CLS] @@ -775,7 +775,7 @@ describe("Offset mapping", () => { // "b c" → [CLS] b c [SEP] const { offset_mapping } = tokenizer(["a", "b c"], { return_tensor: false, - return_offset_mapping: true, + return_offsets_mapping: true, }); expect(offset_mapping).toEqual([ [ @@ -798,7 +798,7 @@ describe("Offset mapping", () => { const { offset_mapping } = tokenizer(["a", "b c"], { return_tensor: false, padding: true, - return_offset_mapping: true, + return_offsets_mapping: true, }); expect(offset_mapping).toEqual([ [ @@ -823,7 +823,7 @@ describe("Offset mapping", () => { truncation: true, max_length: 1, add_special_tokens: false, - return_offset_mapping: true, + return_offsets_mapping: true, }); expect(offset_mapping).toEqual([[0, 1]]); }); @@ -835,7 +835,7 @@ describe("Offset mapping", () => { const output = tokenizer(["a", "a"], { padding: true, truncation: true, - return_offset_mapping: true, + return_offsets_mapping: true, }); // Confirm the other fields are Tensors (they have a tolist() method).