Improve tokenizer types based on input function parameters (#1641)

* Add tokenizer types tests

* update tests

* Improve types

* add back docs

* Update preprocess.js

* fix build

Author: xenova

packages/transformers/docs/plugins/preprocess.js

@@ -25,12 +25,47 @@ function stripLeadingGenericParams(expr) {
   return depth === 0 ? expr.slice(i) : expr;
 }
 
+function hasTopLevelConditional(expr) {
+  let angle = 0,
+    paren = 0,
+    brace = 0,
+    bracket = 0;
+
+  for (let i = 0; i < expr.length; ++i) {
+    const ch = expr[i];
+    if (ch === "<") angle++;
+    else if (ch === ">") angle = Math.max(0, angle - 1);
+    else if (ch === "(") paren++;
+    else if (ch === ")") paren = Math.max(0, paren - 1);
+    else if (ch === "{") brace++;
+    else if (ch === "}") brace = Math.max(0, brace - 1);
+    else if (ch === "[") bracket++;
+    else if (ch === "]") bracket = Math.max(0, bracket - 1);
+    else if (ch === "?" && angle === 0 && paren === 0 && brace === 0 && bracket === 0) {
+      return true;
+    }
+  }
+  return false;
+}
+
 function transformType(expr) {
   let result = expr,
     prev = "";
   while (result !== prev) {
     prev = result;
     result = stripLeadingGenericParams(result); // <T extends X>(...) -> (...)
+    if (/^Promise<function\s*\(/.test(result)) {
+      result = "Promise.<Function>";
+      continue;
+    }
+    if (/^function\s*\(/.test(result)) {
+      result = "Function";
+      continue;
+    }
+    if (hasTopLevelConditional(result)) {
+      result = "any";
+      continue;
+    }
     result = result
       .replace(/import\([^)]+\)((?:\.\w+)*)/g, (_, p) => p?.slice(1) || "any") // import() -> Type
       .replace(/\w+\[['"][^\]]+['"]\]\s+extends\s+[^}]+/g, "any") // X['p'] extends ... -> any
@@ -53,7 +88,6 @@ function transformType(expr) {
       .replace(/(\w+)\[\]/g, "Array") // T[] -> Array (simple)
       .replace(/\[\w+\]/g, "Array") // [T] single-element tuple -> Array
       .replace(/\[[^\[\]]*,[^\[\]]*\]/g, "Array") // tuples with commas -> Array
-      .replace(/\w+\s+extends\s+[^?]+\?\s*[^:]+\s*:\s*[^,}>)]+/g, "any") // conditionals -> any
       .replace(/\bnew\s+([A-Z]\w*)\b/g, "$1") // new Type -> Type
       .replace(/,?\s*\[\s*\w+\s*:\s*\w+\s*\]\s*:\s*\w+/g, "") // [key: string]: any -> (removed)
       .replace(/\(\s*(\w+)\s*&\s*\{\s*\}\s*\)/g, "$1") // (string & {}) -> string

packages/transformers/src/generation/parameters.js

@@ -28,7 +28,7 @@
  * @property {import('./streamers.js').BaseStreamer} [streamer=null] (`BaseStreamer`, *optional*):
  * Streamer object that will be used to stream the generated sequences. Generated tokens are passed
  * through `streamer.put(token_ids)` and the streamer is responsible for any further processing.
- * @property {number[]} [decoder_input_ids=null] (`number[]`, *optional*):
+ * @property {number[]|import('../utils/tensor.js').Tensor} [decoder_input_ids=null] (`number[]` or `Tensor`, *optional*):
  * If the model is an encoder-decoder model, this argument is used to pass the `decoder_input_ids`.
  * @property {import('../cache_utils.js').DynamicCache | null} [past_key_values=null] (`DynamicCache`, *optional*):
  * A cache object that stores previously computed key/value states. When provided, the model will

packages/transformers/src/models/whisper/modeling_whisper.js

@@ -2,6 +2,7 @@ import { cat, mean, Tensor, stack, std_mean } from '../../utils/tensor.js';
 import { PreTrainedModel } from '../modeling_utils.js';
 import { WhisperGenerationConfig } from './generation_whisper.js';
 import { whisper_language_to_code } from './common_whisper.js';
+import { prepareTensorForDecode } from '../../tokenization_utils.js';
 import {
     LogitsProcessorList,
     SuppressTokensAtBeginLogitsProcessor,
@@ -116,7 +117,9 @@ export class WhisperForConditionalGeneration extends WhisperPreTrainedModel {
     }) {
         generation_config = this._prepare_generation_config(generation_config, kwargs);
 
-        const init_tokens = kwargs.decoder_input_ids ?? this._retrieve_init_tokens(generation_config);
+        const init_tokens = kwargs.decoder_input_ids instanceof Tensor
+            ? prepareTensorForDecode(kwargs.decoder_input_ids)
+            : kwargs.decoder_input_ids ?? this._retrieve_init_tokens(generation_config);
 
         if (generation_config.return_timestamps) {
             logits_processor ??= new LogitsProcessorList();

packages/transformers/src/pipelines/text-generation.js

@@ -124,11 +124,15 @@ export class TextGenerationPipeline
             // If the input is a chat, we need to apply the chat template
             inputs = /** @type {string[]} */ (
                 /** @type {Chat[]} */ (texts).map((x) =>
-                    this.tokenizer.apply_chat_template(x, {
-                        tokenize: false,
-                        add_generation_prompt: true,
-                        ...tokenizer_kwargs,
-                    }),
+                    /** @type {string} */ (
+                        /** @type {unknown} */ (
+                            this.tokenizer.apply_chat_template(x, {
+                                tokenize: false,
+                                add_generation_prompt: true,
+                                ...tokenizer_kwargs,
+                            })
+                        )
+                    ),
                 )
             );
             // Chat template handles these already

packages/transformers/src/tokenization_utils.js

@@ -167,7 +167,68 @@ function getSpecialTokens(tokenizer) {
     return special;
 }
 
-export class PreTrainedTokenizer extends Callable {
+/**
+ * @template {string|string[]} TText
+ * @typedef {TText extends string ? number[] : number[][]} BatchEncodingArrayItem
+ */
+
+/**
+ * @template {string|string[]} TText
+ * @template {boolean} [TReturnTensor=true]
+ * @typedef {TReturnTensor extends true ? Tensor : BatchEncodingArrayItem<TText>} BatchEncodingItem
+ */
+
+/**
+ * @template TItem
+ * @typedef {Object} BatchEncoding
+ * @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.
+ */
+
+/**
+ * @template {string|string[]} TText
+ * @template {boolean} [TReturnTensor=true]
+ * @typedef {Object} TokenizerCallOptions
+ * @property {TText extends string ? string|null : string[]|null} [text_pair=null] Optional second sequence to be encoded. If set, must be the same type as text.
+ * @property {boolean|'max_length'} [padding=false] Whether to pad the input sequences.
+ * @property {boolean} [add_special_tokens=true] Whether or not to add the special tokens associated with the corresponding model.
+ * @property {boolean|null} [truncation=null] Whether to truncate the input sequences.
+ * @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.
+ */
+
+/**
+ * @typedef {<TText extends string | string[], TReturnTensor extends boolean = true>(text: TText, options?: TokenizerCallOptions<TText, TReturnTensor>) => BatchEncoding<BatchEncodingItem<TText, TReturnTensor>>} PreTrainedTokenizerCallback
+ */
+
+/**
+ * @template {boolean} [TTokenize=true]
+ * @template {boolean} [TReturnTensor=true]
+ * @template {boolean} [TReturnDict=true]
+ * @typedef {Object} ApplyChatTemplateOptions
+ * @property {string|null} [chat_template=null] A Jinja template to use for this conversion.
+ * @property {Object[]|null} [tools=null] A list of tools (callable functions) that will be accessible to the model.
+ * @property {Record<string, string>[]|null} [documents=null] Documents that will be accessible to the model.
+ * @property {boolean} [add_generation_prompt=false] Whether to end the prompt with the token(s) that indicate the start of an assistant message.
+ * @property {TTokenize} [tokenize=true] Whether to tokenize the output. If false, the output will be a string.
+ * @property {boolean} [padding=false] Whether to pad sequences to the maximum length. Has no effect if tokenize is false.
+ * @property {boolean} [truncation=false] Whether to truncate sequences to the maximum length. Has no effect if tokenize is false.
+ * @property {number|null} [max_length=null] Maximum length (in tokens) to use for padding or truncation. Has no effect if tokenize is false.
+ * @property {TReturnTensor} [return_tensor=true] Whether to return the output as a Tensor or an Array. Has no effect if tokenize is false.
+ * @property {TReturnDict} [return_dict=true] Whether to return a dictionary with named outputs. Has no effect if tokenize is false.
+ * @property {Object} [tokenizer_kwargs={}] Additional options to pass to the tokenizer.
+ */
+
+/**
+ * @template {boolean} [TTokenize=true]
+ * @template {boolean} [TReturnTensor=true]
+ * @template {boolean} [TReturnDict=true]
+ * @typedef {TTokenize extends false ? string : TReturnDict extends false ? BatchEncodingItem<string, TReturnTensor> : BatchEncoding<BatchEncodingItem<string, TReturnTensor>>} ApplyChatTemplateReturn
+ */
+
+export class PreTrainedTokenizer extends /** @type {new (tokenizerJSON: Object, tokenizerConfig: Object) => PreTrainedTokenizerCallback} */ (Callable) {
     return_token_type_ids = false;
 
     padding_side = 'right';
@@ -281,43 +342,23 @@ export class PreTrainedTokenizer extends Callable {
         }
     }
 
-    /**
-     * @typedef {number[]|number[][]|Tensor} BatchEncodingItem
-     *
-     * @typedef {Object} BatchEncoding Holds the output of the tokenizer's call function.
-     * @property {BatchEncodingItem} input_ids List of token ids to be fed to a model.
-     * @property {BatchEncodingItem} attention_mask List of indices specifying which tokens should be attended to by the model.
-     * @property {BatchEncodingItem} [token_type_ids] List of token type ids to be fed to a model.
-     */
-
     /**
      * Encode/tokenize the given text(s).
-     * @param {string|string[]} text The text to tokenize.
-     * @param {Object} options An optional object containing the following properties:
-     * @param {string|string[]} [options.text_pair=null] Optional second sequence to be encoded. If set, must be the same type as text.
-     * @param {boolean|'max_length'} [options.padding=false] Whether to pad the input sequences.
-     * @param {boolean} [options.add_special_tokens=true] Whether or not to add the special tokens associated with the corresponding model.
-     * @param {boolean|null} [options.truncation=null] Whether to truncate the input sequences.
-     * @param {number|null} [options.max_length=null] Maximum length of the returned list and optionally padding length.
-     * @param {boolean} [options.return_tensor=true] Whether to return the results as Tensors or arrays.
-     * @param {boolean|null} [options.return_token_type_ids=null] Whether to return the token type ids.
-     * @returns {BatchEncoding} Object to be passed to the model.
+     * @template {string|string[]} TText
+     * @template {boolean} [TReturnTensor=true]
+     * @param {TText} text The text to tokenize.
+     * @param {TokenizerCallOptions<TText, TReturnTensor>} [options] Additional tokenization options.
+     * @returns {BatchEncoding<BatchEncodingItem<TText, TReturnTensor>>} Object to be passed to the model.
      */
     _call(
         // Required positional arguments
         text,
-
-        // Optional keyword arguments
-        {
-            text_pair = null,
-            add_special_tokens = true,
-            padding = false,
-            truncation = null,
-            max_length = null,
-            return_tensor = true, // Different to HF
-            return_token_type_ids = null,
-        } = {},
+        options = {},
     ) {
+        const { text_pair = null, add_special_tokens = true, padding = false, return_token_type_ids = null } = options;
+        let { truncation = null, max_length = null } = options;
+        const return_tensor = /** @type {TReturnTensor} */ (options.return_tensor ?? true); // Different to HF
+
         const isBatched = Array.isArray(text);
 
         let encodedTokens;
@@ -457,7 +498,7 @@ export class PreTrainedTokenizer extends Callable {
             }
         }
 
-        return /** @type {BatchEncoding} */ (result);
+        return /** @type {BatchEncoding<BatchEncodingItem<TText, TReturnTensor>>} */ (result);
     }
 
     /**
@@ -662,7 +703,10 @@ export class PreTrainedTokenizer extends Callable {
      *
      * @param {Message[]} conversation A list of message objects with `"role"` and `"content"` keys,
      * representing the chat history so far.
-     * @param {Object} options An optional object containing the following properties:
+     * @template {boolean} [TTokenize=true]
+     * @template {boolean} [TReturnTensor=true]
+     * @template {boolean} [TReturnDict=true]
+     * @param {Object} [options] An optional object containing the following properties:
      * @param {string|null} [options.chat_template=null] A Jinja template to use for this conversion. If
      * this is not passed, the model's chat template will be used instead.
      * @param {Object[]} [options.tools=null]
@@ -681,33 +725,34 @@ export class PreTrainedTokenizer extends Callable {
      * the start of an assistant message. This is useful when you want to generate a response from the model.
      * Note that this argument will be passed to the chat template, and so it must be supported in the
      * template for this argument to have any effect.
-     * @param {boolean} [options.tokenize=true] Whether to tokenize the output. If false, the output will be a string.
+     * @param {TTokenize} [options.tokenize=true] Whether to tokenize the output. If false, the output will be a string.
      * @param {boolean} [options.padding=false] Whether to pad sequences to the maximum length. Has no effect if tokenize is false.
      * @param {boolean} [options.truncation=false] Whether to truncate sequences to the maximum length. Has no effect if tokenize is false.
      * @param {number|null} [options.max_length=null] Maximum length (in tokens) to use for padding or truncation. Has no effect if tokenize is false.
      * If not specified, the tokenizer's `max_length` attribute will be used as a default.
-     * @param {boolean} [options.return_tensor=true] Whether to return the output as a Tensor or an Array. Has no effect if tokenize is false.
-     * @param {boolean} [options.return_dict=true] Whether to return a dictionary with named outputs. Has no effect if tokenize is false.
+     * @param {TReturnTensor} [options.return_tensor=true] Whether to return the output as a Tensor or an Array. Has no effect if tokenize is false.
+     * @param {TReturnDict} [options.return_dict=true] Whether to return a dictionary with named outputs. Has no effect if tokenize is false.
      * @param {Object} [options.tokenizer_kwargs={}] Additional options to pass to the tokenizer.
-     * @returns {string | Tensor | number[]| number[][]|BatchEncoding} The tokenized output.
+     * @returns {ApplyChatTemplateReturn<TTokenize, TReturnTensor, TReturnDict>} The tokenized output.
      */
     apply_chat_template(
         conversation,
-        {
+        options = /** @type {ApplyChatTemplateOptions<TTokenize, TReturnTensor, TReturnDict>} */ ({}),
+    ) {
+        let {
             tools = null,
             documents = null,
             chat_template = null,
             add_generation_prompt = false,
-            tokenize = true,
+            tokenize = /** @type {TTokenize} */ (true),
             padding = false,
             truncation = false,
             max_length = null,
-            return_tensor = true,
-            return_dict = true,
+            return_tensor = /** @type {TReturnTensor} */ (true),
+            return_dict = /** @type {TReturnDict} */ (true),
             tokenizer_kwargs = {},
             ...kwargs
-        } = {},
-    ) {
+        } = options;
         chat_template = this.get_chat_template({ chat_template, tools });
 
         if (typeof chat_template !== 'string') {
@@ -748,10 +793,10 @@ export class PreTrainedTokenizer extends Callable {
                 return_tensor,
                 ...tokenizer_kwargs,
             });
-            return return_dict ? out : out.input_ids;
+            return /** @type {ApplyChatTemplateReturn<TTokenize, TReturnTensor, TReturnDict>} */ (return_dict ? out : out.input_ids);
         }
 
-        return rendered;
+        return /** @type {ApplyChatTemplateReturn<TTokenize, TReturnTensor, TReturnDict>} */ (rendered);
     }
 }
 

packages/transformers/tests/types.test.js

@@ -40,7 +40,7 @@ function getDiagnosticsFromSource(source) {
 
 describe("TypeScript compilation succeeds", () => {
   const DIR = "tests/types/";
-  const FILES = ["pipelines.ts", "cache.ts"];
+  const FILES = ["pipelines.ts", "cache.ts", "tokenizers.ts"];
   for (const file of FILES) {
     it(`compiles ${file} without errors`, () => {
       const diagnostics = getDiagnostics(`${DIR}${file}`);