Add offset tracking to Encoding

jest.config.mjs

@@ -173,7 +173,7 @@ export default {
       {
         useESM: true,
         tsconfig: {
-          module: "es2022",
+          module: "esnext",
           target: "es2022",
         },
       },

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

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,50 +300,86 @@ 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]],
+        );
       });
     });
   }
 
   private tokenize_helper(
     text: string,
     { text_pair = null, add_special_tokens = true }: TokenizeOptions,
-  ): { tokens: Array<string>; token_type_ids?: Array<number> } {
-    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 };
   }
 
   /**

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;

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

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

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

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

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

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

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