Add return_offset_mapping support to PreTrainedTokenizer

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<string>} 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)) }
+                : {}),
         };
     }
 

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