Cached generation improvements (+ past_key_values via pipeline function) (#1638)

* add `multi-turn with past_key_values matches without` test

* Add dynamic cache sequence length retrieval test

* sequence length = 0 if empty

* Create DynamicCache update method

* Remove private PreTrainedModel functions & use new methods

* Align with python transformers: load_audio instead of read_audio

* Add past_key_values to possible generation config parameters

* Add test case for past_key_values type

* Type improvements

* Far better error messages!

* Improve typescript tests

ensure good typescript errors

* Apply same typescript improvements to other pipelines

* new typescript tests

* Update preprocess.js

* improve error message

Author: xenova

packages/transformers/docs/plugins/preprocess.js

@@ -14,24 +14,39 @@ function extractBalancedBraces(text, start) {
   return depth === 0 ? { content: text.slice(start + 1, i - 1), endIndex: i } : null;
 }
 
+function stripLeadingGenericParams(expr) {
+  if (expr[0] !== "<") return expr;
+  let depth = 1, i = 1;
+  while (i < expr.length && depth > 0) {
+    if (expr[i] === "<") depth++;
+    else if (expr[i] === ">") depth--;
+    i++;
+  }
+  return depth === 0 ? expr.slice(i) : expr;
+}
+
 function transformType(expr) {
   let result = expr,
     prev = "";
   while (result !== prev) {
     prev = result;
+    result = stripLeadingGenericParams(result); // <T extends X>(...) -> (...)
     result = result
       .replace(/import\([^)]+\)((?:\.\w+)*)/g, (_, p) => p?.slice(1) || "any") // import() -> Type
       .replace(/\w+\[['"][^\]]+['"]\]\s+extends\s+[^}]+/g, "any") // X['p'] extends ... -> any
       .replace(/(\w+)(?:<[^>]+>)?\[[^\]]+\]/g, "$1") // Type[x] or Type<T>[x] -> Type
       .replace(/keyof\s+typeof\s+\w+/g, "string") // keyof typeof X -> string
       .replace(/typeof\s+\w+/g, "Object") // typeof X -> Object
+      .replace(/\binfer\s+\w+/g, "any") // infer K -> any
+      .replace(/\(\s*\w[\w<>, ]*\s+extends\s+\w+\s*\?[^)]*\)/g, "any") // (X extends Y ? A : B) -> any
+      .replace(/(?<!\w)\(\s*(\w+)\s*\)/g, "$1") // (any) -> any (unwrap parens around simple types, not after words like "function")
       .replace(/(\w+)\?\s*:/g, "$1:") // x?: T -> x: T
       .replace(/;\s*([}\n])/g, " $1")
       .replace(/;\s+/g, ", ") // semicolons -> commas
       .replace(/\{\s*\[\w+\s+in\s+[^\]]+\][^}]*\}/g, "Object") // mapped types -> Object
       .replace(/new\s*\([^)]*\)\s*=>\s*\{[^}]*\}/g, "Function") // new () => {...} -> Function
       .replace(/new\s*\([^)]*\)\s*=>\s*\w+/g, "Function") // new () => T -> Function
-      .replace(/\([^()]*\)\s*=>\s*\w[\w<>|[\]]*/g, "Function") // () => T -> Function
+      .replace(/\([^()]*\)\s*=>\s*\w[\w<>|[\], ]*/g, "Function") // () => T -> Function
       .replace(/\{[^{}]*\}\[\]/g, "Array") // {x:T}[] -> Array
       .replace(/(\w+)<[^>]+>\[\]/g, "Array.<$1>") // T<U>[] -> Array.<T>
       .replace(/\([^()]+\)\[\]/g, "Array") // (A|B)[] -> Array

packages/transformers/src/cache_utils.js

@@ -29,6 +29,11 @@ class _DynamicCache {
     get_seq_length() {
         /** @type {Record<string, Tensor>} */
         const self = /** @type {any} */ (this);
+
+        if (Object.keys(self).length === 0) {
+            return 0;
+        }
+
         for (const name in self) {
             if (name.startsWith('past_key_values.')) {
                 return self[name].dims.at(-2);
@@ -37,6 +42,21 @@ class _DynamicCache {
         throw new Error('Unable to determine sequence length from the cache.');
     }
 
+    /**
+     * Update the cache in-place with new entries, disposing replaced GPU tensors.
+     * @param {Record<string, Tensor>} newEntries The new name → Tensor mappings.
+     */
+    update(newEntries) {
+        for (const key in newEntries) {
+            const oldValue = this[key];
+            const newValue = newEntries[key];
+            if (oldValue && oldValue !== newValue && oldValue.location === 'gpu-buffer') {
+                oldValue.dispose();
+            }
+            this[key] = newValue;
+        }
+    }
+
     /**
      * Dispose all contained tensors whose data resides on the GPU.
      * Returns a promise that resolves when all disposals are complete.
@@ -54,7 +74,7 @@ class _DynamicCache {
 }
 
 /**
- * @typedef {_DynamicCache & Record<string, Tensor>} DynamicCache
+ * @typedef {Record<string, Tensor> & _DynamicCache} DynamicCache
  */
 
 export const DynamicCache = /** @type {new (entries?: Record<string, Tensor>) => DynamicCache} */ (

packages/transformers/src/generation/parameters.js

@@ -30,10 +30,13 @@
  * through `streamer.put(token_ids)` and the streamer is responsible for any further processing.
  * @property {number[]} [decoder_input_ids=null] (`number[]`, *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
+ * use these cached states to avoid recomputing them, significantly speeding up sequential generation.
  */
 
 /**
- * @typedef {GenerationFunctionParametersBase & Partial<import('./configuration_utils.js').GenerationConfig> & Record<string, any>} GenerationFunctionParameters
+ * @typedef {GenerationFunctionParametersBase & Partial<import('./configuration_utils.js').GenerationConfig> & {[key: string]: unknown}} GenerationFunctionParameters
  */
 
 export {}; // Ensure this file is treated as a module

packages/transformers/src/models/chatterbox/modeling_chatterbox.js

@@ -168,7 +168,7 @@ export class ChatterboxModel extends ChatterboxPreTrainedModel {
         );
 
         const new_tokens = sequences.slice(null, [
-            params.input_ids.dims[1], // Exclude start of speech token
+            /** @type {Tensor} */ (params.input_ids).dims[1], // Exclude start of speech token
             -1, // Exclude end of speech token
         ]);