diff --git a/.gitignore b/.gitignore index 21c721c4a..10344ec19 100644 --- a/.gitignore +++ b/.gitignore @@ -6,6 +6,7 @@ __pycache__ node_modules deno.lock package-lock.json +*.local.* # Do not track build artifacts/generated files packages/*/dist diff --git a/packages/transformers/src/generation/logits_process.js b/packages/transformers/src/generation/logits_process.js index 647a30806..7bb9545eb 100644 --- a/packages/transformers/src/generation/logits_process.js +++ b/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(); } diff --git a/packages/transformers/src/models/modeling_utils.js b/packages/transformers/src/models/modeling_utils.js index ec9f17487..f5caec12c 100644 --- a/packages/transformers/src/models/modeling_utils.js +++ b/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); } diff --git a/packages/transformers/tests/utils/generation.test.js b/packages/transformers/tests/utils/generation.test.js index 231985571..db991affc 100644 --- a/packages/transformers/tests/utils/generation.test.js +++ b/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); diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index bea34792e..3ed937a04 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -427,89 +427,105 @@ packages: resolution: {integrity: sha512-excjX8DfsIcJ10x1Kzr4RcWe1edC9PquDRRPx3YVCvQv+U5p7Yin2s32ftzikXojb1PIFc/9Mt28/y+iRklkrw==} cpu: [arm64] os: [linux] + libc: [glibc] '@img/sharp-libvips-linux-arm@1.2.4': resolution: {integrity: sha512-bFI7xcKFELdiNCVov8e44Ia4u2byA+l3XtsAj+Q8tfCwO6BQ8iDojYdvoPMqsKDkuoOo+X6HZA0s0q11ANMQ8A==} cpu: [arm] os: [linux] + libc: [glibc] '@img/sharp-libvips-linux-ppc64@1.2.4': resolution: {integrity: sha512-FMuvGijLDYG6lW+b/UvyilUWu5Ayu+3r2d1S8notiGCIyYU/76eig1UfMmkZ7vwgOrzKzlQbFSuQfgm7GYUPpA==} cpu: [ppc64] os: [linux] + libc: [glibc] '@img/sharp-libvips-linux-riscv64@1.2.4': resolution: {integrity: sha512-oVDbcR4zUC0ce82teubSm+x6ETixtKZBh/qbREIOcI3cULzDyb18Sr/Wcyx7NRQeQzOiHTNbZFF1UwPS2scyGA==} cpu: [riscv64] os: [linux] + libc: [glibc] '@img/sharp-libvips-linux-s390x@1.2.4': resolution: {integrity: sha512-qmp9VrzgPgMoGZyPvrQHqk02uyjA0/QrTO26Tqk6l4ZV0MPWIW6LTkqOIov+J1yEu7MbFQaDpwdwJKhbJvuRxQ==} cpu: [s390x] os: [linux] + libc: [glibc] '@img/sharp-libvips-linux-x64@1.2.4': resolution: {integrity: sha512-tJxiiLsmHc9Ax1bz3oaOYBURTXGIRDODBqhveVHonrHJ9/+k89qbLl0bcJns+e4t4rvaNBxaEZsFtSfAdquPrw==} cpu: [x64] os: [linux] + libc: [glibc] '@img/sharp-libvips-linuxmusl-arm64@1.2.4': resolution: {integrity: sha512-FVQHuwx1IIuNow9QAbYUzJ+En8KcVm9Lk5+uGUQJHaZmMECZmOlix9HnH7n1TRkXMS0pGxIJokIVB9SuqZGGXw==} cpu: [arm64] os: [linux] + libc: [musl] '@img/sharp-libvips-linuxmusl-x64@1.2.4': resolution: {integrity: sha512-+LpyBk7L44ZIXwz/VYfglaX/okxezESc6UxDSoyo2Ks6Jxc4Y7sGjpgU9s4PMgqgjj1gZCylTieNamqA1MF7Dg==} cpu: [x64] os: [linux] + libc: [musl] '@img/sharp-linux-arm64@0.34.5': resolution: {integrity: sha512-bKQzaJRY/bkPOXyKx5EVup7qkaojECG6NLYswgktOZjaXecSAeCWiZwwiFf3/Y+O1HrauiE3FVsGxFg8c24rZg==} engines: {node: ^18.17.0 || ^20.3.0 || >=21.0.0} cpu: [arm64] os: [linux] + libc: [glibc] '@img/sharp-linux-arm@0.34.5': resolution: {integrity: sha512-9dLqsvwtg1uuXBGZKsxem9595+ujv0sJ6Vi8wcTANSFpwV/GONat5eCkzQo/1O6zRIkh0m/8+5BjrRr7jDUSZw==} engines: {node: ^18.17.0 || ^20.3.0 || >=21.0.0} cpu: [arm] os: [linux] + libc: [glibc] '@img/sharp-linux-ppc64@0.34.5': resolution: {integrity: sha512-7zznwNaqW6YtsfrGGDA6BRkISKAAE1Jo0QdpNYXNMHu2+0dTrPflTLNkpc8l7MUP5M16ZJcUvysVWWrMefZquA==} engines: {node: ^18.17.0 || ^20.3.0 || >=21.0.0} cpu: [ppc64] os: [linux] + libc: [glibc] '@img/sharp-linux-riscv64@0.34.5': resolution: {integrity: sha512-51gJuLPTKa7piYPaVs8GmByo7/U7/7TZOq+cnXJIHZKavIRHAP77e3N2HEl3dgiqdD/w0yUfiJnII77PuDDFdw==} engines: {node: ^18.17.0 || ^20.3.0 || >=21.0.0} cpu: [riscv64] os: [linux] + libc: [glibc] '@img/sharp-linux-s390x@0.34.5': resolution: {integrity: sha512-nQtCk0PdKfho3eC5MrbQoigJ2gd1CgddUMkabUj+rBevs8tZ2cULOx46E7oyX+04WGfABgIwmMC0VqieTiR4jg==} engines: {node: ^18.17.0 || ^20.3.0 || >=21.0.0} cpu: [s390x] os: [linux] + libc: [glibc] '@img/sharp-linux-x64@0.34.5': resolution: {integrity: sha512-MEzd8HPKxVxVenwAa+JRPwEC7QFjoPWuS5NZnBt6B3pu7EG2Ge0id1oLHZpPJdn3OQK+BQDiw9zStiHBTJQQQQ==} engines: {node: ^18.17.0 || ^20.3.0 || >=21.0.0} cpu: [x64] os: [linux] + libc: [glibc] '@img/sharp-linuxmusl-arm64@0.34.5': resolution: {integrity: sha512-fprJR6GtRsMt6Kyfq44IsChVZeGN97gTD331weR1ex1c1rypDEABN6Tm2xa1wE6lYb5DdEnk03NZPqA7Id21yg==} engines: {node: ^18.17.0 || ^20.3.0 || >=21.0.0} cpu: [arm64] os: [linux] + libc: [musl] '@img/sharp-linuxmusl-x64@0.34.5': resolution: {integrity: sha512-Jg8wNT1MUzIvhBFxViqrEhWDGzqymo3sV7z7ZsaWbZNDLXRJZoRGrjulp60YYtV4wfY8VIKcWidjojlLcWrd8Q==} engines: {node: ^18.17.0 || ^20.3.0 || >=21.0.0} cpu: [x64] os: [linux] + libc: [musl] '@img/sharp-wasm32@0.34.5': resolution: {integrity: sha512-OdWTEiVkY2PHwqkbBI8frFxQQFekHaSSkUIJkwzclWZe64O1X4UlUjqqqLaPbUpMOQk6FBu/HtlGXNblIs0huw==} @@ -800,41 +816,49 @@ packages: resolution: {integrity: sha512-34gw7PjDGB9JgePJEmhEqBhWvCiiWCuXsL9hYphDF7crW7UgI05gyBAi6MF58uGcMOiOqSJ2ybEeCvHcq0BCmQ==} cpu: [arm64] os: [linux] + libc: [glibc] '@unrs/resolver-binding-linux-arm64-musl@1.11.1': resolution: {integrity: sha512-RyMIx6Uf53hhOtJDIamSbTskA99sPHS96wxVE/bJtePJJtpdKGXO1wY90oRdXuYOGOTuqjT8ACccMc4K6QmT3w==} cpu: [arm64] os: [linux] + libc: [musl] '@unrs/resolver-binding-linux-ppc64-gnu@1.11.1': resolution: {integrity: sha512-D8Vae74A4/a+mZH0FbOkFJL9DSK2R6TFPC9M+jCWYia/q2einCubX10pecpDiTmkJVUH+y8K3BZClycD8nCShA==} cpu: [ppc64] os: [linux] + libc: [glibc] '@unrs/resolver-binding-linux-riscv64-gnu@1.11.1': resolution: {integrity: sha512-frxL4OrzOWVVsOc96+V3aqTIQl1O2TjgExV4EKgRY09AJ9leZpEg8Ak9phadbuX0BA4k8U5qtvMSQQGGmaJqcQ==} cpu: [riscv64] os: [linux] + libc: [glibc] '@unrs/resolver-binding-linux-riscv64-musl@1.11.1': resolution: {integrity: sha512-mJ5vuDaIZ+l/acv01sHoXfpnyrNKOk/3aDoEdLO/Xtn9HuZlDD6jKxHlkN8ZhWyLJsRBxfv9GYM2utQ1SChKew==} cpu: [riscv64] os: [linux] + libc: [musl] '@unrs/resolver-binding-linux-s390x-gnu@1.11.1': resolution: {integrity: sha512-kELo8ebBVtb9sA7rMe1Cph4QHreByhaZ2QEADd9NzIQsYNQpt9UkM9iqr2lhGr5afh885d/cB5QeTXSbZHTYPg==} cpu: [s390x] os: [linux] + libc: [glibc] '@unrs/resolver-binding-linux-x64-gnu@1.11.1': resolution: {integrity: sha512-C3ZAHugKgovV5YvAMsxhq0gtXuwESUKc5MhEtjBpLoHPLYM+iuwSj3lflFwK3DPm68660rZ7G8BMcwSro7hD5w==} cpu: [x64] os: [linux] + libc: [glibc] '@unrs/resolver-binding-linux-x64-musl@1.11.1': resolution: {integrity: sha512-rV0YSoyhK2nZ4vEswT/QwqzqQXw5I6CjoaYMOX0TqBlWhojUf8P94mvI7nuJTeaCkkds3QE4+zS8Ko+GdXuZtA==} cpu: [x64] os: [linux] + libc: [musl] '@unrs/resolver-binding-wasm32-wasi@1.11.1': resolution: {integrity: sha512-5u4RkfxJm+Ng7IWgkzi3qrFOvLvQYnPBmjmZQ8+szTK/b31fQCnleNl1GgEt7nIsZRIf5PLhPwT0WM+q45x/UQ==}