Add Generation Extension Points for Constrained Decoding

.gitignore

@@ -6,6 +6,7 @@ __pycache__
 node_modules
 deno.lock
 package-lock.json
+*.local.*
 
 # Do not track build artifacts/generated files
 packages/*/dist

packages/transformers/src/generation/logits_process.js

@@ -88,6 +88,19 @@ export class LogitsProcessorList extends Callable {
         return toReturn;
     }
 
+    /**
+     * Calls Transformers.js-specific post-sampling hooks on processors that need to update state after token selection.
+     * The hook observes the full batch after the current generation step has been appended.
+     *
+     * @param {number[]} token_ids The sampled token IDs for the current generation step.
+     * @param {bigint[][]} input_ids The input IDs after appending the sampled tokens.
+     */
+    onTokensSampled(token_ids, input_ids) {
+        for (const processor of this.processors) {
+            processor.onTokensSampled?.(token_ids, input_ids);
+        }
+    }
+
     [Symbol.iterator]() {
         return this.processors.values();
     }

packages/transformers/src/models/modeling_utils.js

@@ -1003,13 +1003,19 @@ export class PreTrainedModel extends Callable {
                     // TODO: If branching, use previous beam as a starting point
                     // update generated ids, model inputs, and length for next step
                     scores[batch_idx] += logProb;
-                    all_input_ids[batch_idx].push(bigint);
                     generated_input_ids.push([bigint]);
 
                     // TODO: Support beam search
                     break;
                 }
             }
+            for (let batch_idx = 0; batch_idx < generated_input_ids.length; ++batch_idx) {
+                all_input_ids[batch_idx].push(generated_input_ids[batch_idx][0]);
+            }
+            prepared_logits_processor.onTokensSampled(
+                generated_input_ids.map(([token_id]) => Number(token_id)),
+                all_input_ids,
+            );
             if (streamer) {
                 streamer.put(generated_input_ids);
             }

packages/transformers/tests/utils/generation.test.js

@@ -11,6 +11,9 @@ import {
   // Other
   TextStreamer,
   DynamicCache,
+  LogitsProcessor,
+  LogitsProcessorList,
+  StoppingCriteria,
   random,
   full,
 } from "../../src/transformers.js";
@@ -208,6 +211,73 @@ describe("Generation parameters", () => {
       MAX_TEST_EXECUTION_TIME,
     );
 
+    it(
+      "calls logits processor post-sample hook after full batch step",
+      async () => {
+        class RecordingLogitsProcessor extends LogitsProcessor {
+          snapshots = [];
+
+          _call(input_ids, logits) {
+            return logits;
+          }
+
+          onTokensSampled(token_ids, input_ids) {
+            this.snapshots.push({
+              token_ids,
+              lengths: input_ids.map((ids) => ids.length),
+            });
+          }
+        }
+
+        const processor = new RecordingLogitsProcessor();
+        const logits_processor = new LogitsProcessorList();
+        logits_processor.push(processor);
+
+        const outputs = await generate(model, tokenizer, [DUMMY_TEXT, DUMMY_TEXT], {
+          max_new_tokens: 2,
+          logits_processor,
+        });
+
+        const generated_tokens = outputs.tolist().map((tokens) => tokens.slice(-2).map(Number));
+        expect(processor.snapshots).toEqual([
+          {
+            token_ids: generated_tokens.map((tokens) => tokens[0]),
+            lengths: [3, 3],
+          },
+          {
+            token_ids: generated_tokens.map((tokens) => tokens[1]),
+            lengths: [4, 4],
+          },
+        ]);
+      },
+      MAX_TEST_EXECUTION_TIME,
+    );
+
+    it(
+      "supports custom stopping criteria",
+      async () => {
+        class StopAfterLengthCriteria extends StoppingCriteria {
+          constructor(max_length) {
+            super();
+            this.max_length = max_length;
+          }
+
+          _call(input_ids) {
+            return input_ids.map((ids) => ids.length >= this.max_length);
+          }
+        }
+
+        const outputs = await generate(model, tokenizer, DUMMY_TEXT, {
+          max_new_tokens: 5,
+          stopping_criteria: new StopAfterLengthCriteria(3),
+        });
+
+        // BOS + DUMMY_TEXT + exactly one generated token
+        expect(outputs.dims.at(-1)).toEqual(3);
+      },
+      MAX_TEST_EXECUTION_TIME,
+    );
+
     afterAll(async () => {
       await model?.dispose();
     }, MAX_MODEL_DISPOSE_TIME);