From b179ef029cb176f76e6e64f346026eb5735745f3 Mon Sep 17 00:00:00 2001 From: Nico Martin Date: Thu, 9 Jul 2026 12:40:55 +0200 Subject: [PATCH 1/4] first draft for #1328 using llguidance --- .gitignore | 1 + PLAN.local.md | 202 +++++++++++++++ packages/transformers/package.json | 1 + .../src/generation/grammar/llguidance.js | 54 +++++ .../generation/grammar/tokenizer_bridge.js | 51 ++++ .../src/generation/logits_process.js | 229 ++++++++++++++++++ .../transformers/src/models/modeling_utils.js | 11 + .../src/pipelines/text-generation.js | 57 ++++- .../tests/utils/generation.test.js | 38 +++ .../tests/utils/logits_process.test.js | 124 ++++++++++ pnpm-lock.yaml | 8 + test/response_format.js | 84 +++++++ 12 files changed, 859 insertions(+), 1 deletion(-) create mode 100644 PLAN.local.md create mode 100644 packages/transformers/src/generation/grammar/llguidance.js create mode 100644 packages/transformers/src/generation/grammar/tokenizer_bridge.js create mode 100644 test/response_format.js 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/PLAN.local.md b/PLAN.local.md new file mode 100644 index 000000000..0deb8e5f8 --- /dev/null +++ b/PLAN.local.md @@ -0,0 +1,202 @@ +# Plan: `response_format` (JSON Schema constrained decoding) in Transformers.js + +## Goal + +Add OpenAI-compatible structured output support to Transformers.js, backed by +[llguidance](https://github.com/guidance-ai/llguidance) (Rust, compiled to WASM) for +grammar-constrained decoding. Guarantees schema-valid JSON output instead of +prompt-and-hope + post-hoc validation. + +## Target API + +```js +const pipe = await pipeline( + "text-generation", + "onnx-community/gemma-4-E2B-it-ONNX", + { + device: "webgpu", + dtype: "q4f16", + } +); + +const schema = { + type: "array", + items: { type: "string" }, +}; + +await pipe(messages, { + max_new_tokens: 1024, + response_format: { type: "json_schema", json_schema: schema }, +}); +``` + +- `response_format` is a new option accepted by the `pipeline()` call function (text-generation + pipeline first, generalizable later). +- Mirrors OpenAI's `response_format` shape so it's a drop-in mental model for JS devs. +- Internally resolves to attaching a constrained `LogitsProcessor` to the `generate()` call. +- `type: "json_object"` (schema-less, valid-JSON-only) should be supported too as a cheap + first milestone (no schema compilation needed, just JSON grammar). + +## Why llguidance + +Already covered in prior research — key points to keep in mind while implementing: + +- No pre-computation/startup cost (unlike Outlines), fast enough to compute masks on-the-fly + per step (~50μs on native CPU for a 128k vocab). +- Proven to run in a browser (merged into Chromium for `window.ai`'s `responseConstraint`), + though Chromium's actual wiring code is closed-source — we can't copy it, only use the + open-source llguidance engine itself. +- Tokenizer integration is a **one-time setup cost**, not a per-step cost: llguidance builds a + `TokTrie` from the full vocab (as raw bytes) once; per-step calls are just + `compute_mask()` → apply bitmask to logits → `commit_token()`. No decoding in the hot loop. + +## Architecture overview + +``` +response_format (user input) + │ + ▼ +schema → llguidance grammar compiler (WASM) ─┐ + │ │ (one-time per generate() call) + ▼ │ +tokenizer vocab → TokTrie (WASM) ─────────────┘ (one-time per model load, cached) + │ + ▼ +SchemaConstrainedLogitsProcessor (new LogitsProcessor subclass) + │ each decode step: + │ mask = interpreter.compute_mask() + │ logits[i] = -Infinity where mask bit is 0 + │ (after sampling) interpreter.commit_token(sampled_id) + ▼ +generate() loop (existing, unchanged) → LogitsProcessorList → sampler +``` + +## Phases + +### Phase 0 — Confirm/extend the plug-in point + +- Verify `generate()`'s public `logits_processor` option (appended after built-ins) is + sufficient, or needs a small extension to support processors that need a `commit_token`-style + callback *after* sampling (llguidance requires this — most existing `LogitsProcessor`s only + hook pre-sampling). +- Likely need a small addition to the generation loop: an optional post-sample hook + (e.g. `processor.onTokenSampled?.(tokenId)`), since `LogitsProcessorList` today assumes + stateless-after-`_call` processors. +- File: `src/generation/logits_process.js`, `src/models/modeling_utils.js` (generate loop). + +### Phase 1 — Spike: compile llguidance to WASM, measure feasibility + +- Compile llguidance's `parser` crate to `wasm32-unknown-unknown` (or `wasm32-wasi` if easier, + but prefer unknown-unknown for browser bundle compatibility). Check if a maintained + wasm-bindgen target already exists upstream before building our own. +- Measure and report: + - Bundle size (gzip) of the compiled `.wasm` — this is the #1 go/no-go metric for a + browser-shipped library. + - Per-step `compute_mask()` latency through the WASM boundary + JS marshaling (not just + native Rust numbers) for a realistic vocab size (Gemma/Llama ~256k, ~128k). + - One-time `TokTrie` build latency for a full model vocab. +- Decision gate: if bundle size or per-step latency is unacceptable, fall back to a + JSON-Schema-only hand-rolled JS grammar (Jsonformer-style: deterministic structural tokens, + masking only on leaf values) as a lighter-weight v1. Document the fallback trigger criteria + before starting Phase 2 so this isn't an open-ended detour later. + +### Phase 2 — Tokenizer → llguidance bridge (one-time per model) + +- Build an adapter that extracts, once per model/tokenizer load: + - `tokens: bytes[]` — every vocab entry as **raw UTF-8 bytes**, not the human-readable + decoded form. For byte-level BPE (GPT-2/Llama-style, e.g. `Ġgazed`), this means resolving + the byte-level encoding table to actual bytes, not just using `tokenizer.decode()`. + - `eos_token_id`, `bos_token_id`, `special_token_ids` — already available on Transformers.js + tokenizer objects, just need to be surfaced in the shape llguidance expects. +- Pass this to the WASM module once to construct the `LLTokenizer` / `TokTrie`. +- Cache the resulting handle per model instance (same caching pattern as `ModelRegistry`) — + rebuilding the trie per `generate()` call is wasted work if the model doesn't change. +- File: new `src/generation/grammar/tokenizer_bridge.js` (or similar). + +### Phase 3 — `SchemaConstrainedLogitsProcessor` + +```js +class SchemaConstrainedLogitsProcessor extends LogitsProcessor { + constructor(schema, llguidanceTokenizer) { + super(); + // compile schema -> llguidance grammar (LLInterpreter), once per generate() call + } + + _call(input_ids, logits) { + const { mask } = this.interpreter.compute_mask(); + // apply mask (bitset, vocab_size/32 elements) to logits typed array + return logits; + } + + onTokenSampled(tokenId) { + this.interpreter.commit_token(tokenId); + } +} +``` + +- Needs efficient bitmask application against Transformers.js's logits representation + (likely a `Tensor`/`Float32Array`) — avoid per-element JS loops if possible, use typed-array + ops. +- Handle the "grammar reached a stop state" signal from llguidance to force/allow EOS. +- File: `src/generation/logits_process.js`. + +### Phase 4 — Wire up `response_format` in the pipeline + +- Add `response_format` to the text-generation pipeline's call options + (`src/pipelines.js`, `TextGenerationPipeline`). +- Validation: + - `type: "json_object"` → JSON-only grammar, no schema compilation. + - `type: "json_schema"` → compile `json_schema` field via llguidance. + - Unsupported schema features → throw a clear error before generation starts (mirror Chrome's + `NotSupportedError` behavior — fail fast, not mid-generation). +- Translate into a `SchemaConstrainedLogitsProcessor` instance, append to the + `logits_processor` list passed into `generate()`. +- Decide default behavior on schema-includes-context-window cost: unlike Chrome, we probably + should **not** auto-inject the schema into the prompt by default (grammar constraint alone + should be sufficient and cheaper on context budget) — but expose an opt-in flag if empirically + output quality benefits from also showing the schema in-prompt (worth A/B testing once + working). + +### Phase 5 — Batch generation handling + +- llguidance's `Constraint`/`LLInterpreter` is single-sequence/stateful. +- For Transformers.js's batched `generate()`, need one interpreter instance **per batch item**, + each tracking its own grammar state and getting masked independently. +- Scope: v1 may explicitly restrict `response_format` to `batch_size === 1` and throw otherwise, + documented as a known limitation, with batched support as a fast-follow. + +### Phase 6 — Tests & docs + +- Unit tests: schema compilation, mask correctness on known small vocabs (reuse llguidance's + own JSON Schema Test Suite harness pattern if feasible). +- Integration test: run a small ONNX model end-to-end with a nested schema (object + array + + enum + required fields), assert `JSON.parse()` succeeds and matches schema on every run + (structural guarantee — should never flake). +- Perf test: measure generation-loop overhead with/without constrained decoding on WebGPU vs + WASM backend. +- Docs: new guide page + `response_format` API reference, modeled on OpenAI's docs since the + shape is intentionally familiar. + +## Open risks (carry into implementation, don't resolve in the plan) + +1. **WASM bundle size** — could be the single blocking issue; resolve in Phase 1 before + committing further engineering time. +2. **Post-sample hook** — `LogitsProcessorList`/`generate()` today may not have a clean seam for + `commit_token()`-style post-sampling callbacks; needs a small core API addition. +3. **Byte-level tokenizer edge cases** — getting the raw-bytes extraction wrong will silently + produce incorrect masks (accepting/rejecting wrong tokens) rather than an obvious crash — + needs solid test coverage, not just "it compiled." +4. **Batching** — v1 scoping to single-sequence generation is a real product limitation worth + flagging in the RFC/issue up front, not discovering mid-implementation. +5. **WebGPU/WASM boundary latency** — mask computation happening off the main compute path + (CPU, via WASM) needs to not stall the GPU forward pass; may need to run concurrently + (per llguidance's own recommendation: run `compute_mask()` while logits are being computed). + +## Suggested order of work for a first PR + +1. Phase 1 spike (throwaway branch, just prove bundle size + latency are acceptable). +2. Phase 0 + Phase 2 (core plug-in point + tokenizer bridge) — no user-facing API yet. +3. Phase 3 (`SchemaConstrainedLogitsProcessor`) with a manual/internal test using `generate()` + directly (no pipeline API yet). +4. Phase 4 (`response_format` on the pipeline) — first user-facing milestone. +5. Phase 5 (batching) and Phase 6 (tests/docs) as follow-ups, potentially separate PRs. \ No newline at end of file diff --git a/packages/transformers/package.json b/packages/transformers/package.json index 8ea694b11..86f7819d7 100644 --- a/packages/transformers/package.json +++ b/packages/transformers/package.json @@ -57,6 +57,7 @@ "dependencies": { "@huggingface/jinja": "^0.5.6", "@huggingface/tokenizers": "^0.1.3", + "llguidance": "0.1.4", "onnxruntime-node": "1.24.3", "onnxruntime-web": "1.26.0-dev.20260416-b7804b056c", "sharp": "^0.34.5" diff --git a/packages/transformers/src/generation/grammar/llguidance.js b/packages/transformers/src/generation/grammar/llguidance.js new file mode 100644 index 000000000..fe5a1f23d --- /dev/null +++ b/packages/transformers/src/generation/grammar/llguidance.js @@ -0,0 +1,54 @@ +/** + * @module generation/grammar/llguidance + */ + +let cachedRuntimePromise = null; +const normalizedRuntimeCache = new WeakMap(); + +function normalizeRuntime(runtime, options = {}) { + const cached = normalizedRuntimeCache.get(runtime); + if (cached) { + return cached; + } + + const createTokenizer = runtime.createTokenizer ?? runtime.create_tokenizer; + const createInterpreter = runtime.createInterpreter ?? runtime.create_interpreter; + + if (typeof createTokenizer !== 'function' || typeof createInterpreter !== 'function') { + throw new Error( + 'Invalid llguidance runtime: expected createTokenizer/createInterpreter functions exposed by the WASM module.', + ); + } + + const normalized = { ...runtime, ...options, createTokenizer, createInterpreter }; + normalizedRuntimeCache.set(runtime, normalized); + return normalized; +} + +async function loadBundledRuntime() { + const { loadBundledLLGuidance } = await import('llguidance'); + return normalizeRuntime(await loadBundledLLGuidance(), { acceptsTokenizerObjects: true }); +} + +/** + * Loads the llguidance WASM runtime. + * + * Consumers can provide a compatible runtime on + * `globalThis.__transformers_llguidance`; otherwise, the bundled `llguidance` + * package runtime is loaded lazily. + * + * @param {Object|null} runtime Optional runtime, primarily useful for tests or custom integrations. + * @returns {Promise} + */ +export async function loadLLGuidanceRuntime(runtime = null) { + if (runtime) { + return normalizeRuntime(runtime); + } + + if (globalThis.__transformers_llguidance) { + return normalizeRuntime(globalThis.__transformers_llguidance); + } + + cachedRuntimePromise ??= loadBundledRuntime(); + return cachedRuntimePromise; +} diff --git a/packages/transformers/src/generation/grammar/tokenizer_bridge.js b/packages/transformers/src/generation/grammar/tokenizer_bridge.js new file mode 100644 index 000000000..04233650b --- /dev/null +++ b/packages/transformers/src/generation/grammar/tokenizer_bridge.js @@ -0,0 +1,51 @@ +/** + * @module generation/grammar/tokenizer_bridge + */ + +const encoder = new TextEncoder(); +const tokenizerCache = new WeakMap(); + +function decodeToken(tokenizer, token_id) { + try { + return tokenizer.decode([token_id], { skip_special_tokens: false, clean_up_tokenization_spaces: false }); + } catch { + return ''; + } +} + +/** + * Creates or returns the cached llguidance tokenizer handle for a Transformers.js tokenizer. + * + * @param {import('../../tokenization_utils.js').PreTrainedTokenizer} tokenizer The tokenizer to bridge. + * @param {Object} runtime The normalized llguidance runtime. + * @returns {Promise} + */ +export async function getLLGuidanceTokenizer(tokenizer, runtime) { + let runtimeCache = tokenizerCache.get(tokenizer); + if (!runtimeCache) { + runtimeCache = new WeakMap(); + tokenizerCache.set(tokenizer, runtimeCache); + } + + let cached = runtimeCache.get(runtime); + if (cached) { + return cached; + } + + cached = Promise.resolve().then(() => { + const vocab = tokenizer.get_vocab(); + const tokens = []; + for (const [token, id] of Object.entries(vocab)) { + tokens[id] = encoder.encode(decodeToken(tokenizer, id) || token); + } + + return runtime.createTokenizer({ + tokens, + eos_token_id: tokenizer.eos_token_id, + bos_token_id: tokenizer.bos_token_id, + special_token_ids: tokenizer.all_special_ids ?? [], + }); + }); + runtimeCache.set(runtime, cached); + return cached; +} diff --git a/packages/transformers/src/generation/logits_process.js b/packages/transformers/src/generation/logits_process.js index 647a30806..bc70e5469 100644 --- a/packages/transformers/src/generation/logits_process.js +++ b/packages/transformers/src/generation/logits_process.js @@ -6,6 +6,8 @@ import { Callable } from '../utils/generic.js'; import { Tensor } from '../utils/tensor.js'; import { max, log_softmax } from '../utils/maths.js'; +import { loadLLGuidanceRuntime } from './grammar/llguidance.js'; +import { getLLGuidanceTokenizer } from './grammar/tokenizer_bridge.js'; /** * Abstract base class for all logit processors that can be applied during generation. @@ -22,6 +24,25 @@ export class LogitsProcessor extends Callable { _call(input_ids, logits) { throw Error('`_call` should be implemented in a subclass'); } + + /** + * Optional hook called after a token has been selected by the sampler. + * + * @param {number} token_id The sampled token ID. + * @param {number} batch_idx The batch index that sampled the token. + * @param {bigint[][]} input_ids The input IDs after appending the sampled token. + */ + onTokenSampled(token_id, batch_idx, input_ids) {} + + /** + * Optional hook for processors that can terminate generation. + * + * @param {bigint[][]} input_ids The input IDs. + * @returns {boolean[]|null} + */ + shouldStop(input_ids) { + return null; + } } /** @@ -39,6 +60,25 @@ export class LogitsWarper extends Callable { _call(input_ids, logits) { throw Error('`_call` should be implemented in a subclass'); } + + /** + * Optional hook called after a token has been selected by the sampler. + * + * @param {number} token_id The sampled token ID. + * @param {number} batch_idx The batch index that sampled the token. + * @param {bigint[][]} input_ids The input IDs after appending the sampled token. + */ + onTokenSampled(token_id, batch_idx, input_ids) {} + + /** + * Optional hook for processors that can terminate generation. + * + * @param {bigint[][]} input_ids The input IDs. + * @returns {boolean[]|null} + */ + shouldStop(input_ids) { + return null; + } } /** @@ -88,11 +128,200 @@ export class LogitsProcessorList extends Callable { return toReturn; } + /** + * Calls post-sampling hooks on processors that need to update state after token selection. + * + * @param {number} token_id The sampled token ID. + * @param {number} batch_idx The batch index that sampled the token. + * @param {bigint[][]} input_ids The input IDs after appending the sampled token. + */ + onTokenSampled(token_id, batch_idx, input_ids) { + for (const processor of this.processors) { + processor.onTokenSampled?.(token_id, batch_idx, input_ids); + } + } + + /** + * Calls stopping hooks on processors that can terminate generation. + * + * @param {bigint[][]} input_ids The input IDs. + * @returns {boolean[]|null} + */ + shouldStop(input_ids) { + let stop = null; + for (const processor of this.processors) { + const processorStop = processor.shouldStop?.(input_ids); + if (!processorStop) { + continue; + } + stop ??= new Array(input_ids.length).fill(false); + for (let i = 0; i < stop.length; ++i) { + stop[i] ||= processorStop[i]; + } + } + return stop; + } + [Symbol.iterator]() { return this.processors.values(); } } +function validateResponseFormat(response_format) { + if (!response_format || typeof response_format !== 'object') { + throw new Error('`response_format` must be an object.'); + } + + if (response_format.type === 'json_object') { + return; + } + + if (response_format.type === 'json_schema') { + if (!response_format.json_schema || typeof response_format.json_schema !== 'object') { + throw new Error('`response_format.json_schema` must be an object for type "json_schema".'); + } + return; + } + + throw new Error('Unsupported `response_format.type`. Expected "json_object" or "json_schema".'); +} + +function callRuntimeMethod(object, camelCaseName, snakeCaseName, ...args) { + const method = object[camelCaseName] ?? object[snakeCaseName]; + if (typeof method !== 'function') { + throw new Error(`Invalid llguidance runtime object: missing ${camelCaseName}/${snakeCaseName}.`); + } + return method.call(object, ...args); +} + +function isAllowed(mask, token_id, vocab_size) { + if (mask.length >= vocab_size) { + return Boolean(mask[token_id]); + } + + return Boolean(mask[token_id >> 5] & (1 << (token_id & 31))); +} + +function isComputeAfterStopError(error) { + return error instanceof Error && error.message.includes('compute_mask() called after stop'); +} + +/** + * Logits processor backed by llguidance for JSON-constrained generation. + */ +export class SchemaConstrainedLogitsProcessor extends LogitsProcessor { + /** + * @param {Object} interpreter A llguidance interpreter/constraint instance. + * @param {number|null} eos_token_id The EOS token ID to force when the grammar stops. + */ + constructor(interpreter, eos_token_id = null) { + super(); + this.interpreter = interpreter; + this.eos_token_id = eos_token_id; + this.completed = false; + } + + /** + * Creates a constrained logits processor from an OpenAI-compatible response_format object. + * + * @param {Object} response_format The requested response format. + * @param {import('../tokenization_utils.js').PreTrainedTokenizer} tokenizer The tokenizer used for generation. + * @param {Object|null} runtime Optional llguidance runtime override. + * @returns {Promise} + */ + static async fromResponseFormat(response_format, tokenizer, runtime = null, eos_token_id = null) { + validateResponseFormat(response_format); + + const llguidance = await loadLLGuidanceRuntime(runtime); + const llguidanceTokenizer = llguidance.acceptsTokenizerObjects + ? tokenizer + : await getLLGuidanceTokenizer(tokenizer, llguidance); + const interpreter = await llguidance.createInterpreter({ + tokenizer: llguidanceTokenizer, + response_format, + }); + + return new SchemaConstrainedLogitsProcessor(interpreter, eos_token_id ?? tokenizer.eos_token_id); + } + + /** + * @param {Tensor} logits The logits to process. + */ + forceEOS(logits) { + if (!Number.isInteger(this.eos_token_id)) { + throw new Error('`response_format` reached a stop state, but no EOS token is available to end generation.'); + } + + logits.data.fill(-Infinity); + logits.data[this.eos_token_id] = 0; + return logits; + } + + /** + * @param {bigint[][]} input_ids The input IDs. + * @returns {boolean[]} + */ + shouldStop(input_ids) { + return new Array(input_ids.length).fill(this.completed); + } + + /** + * @param {bigint[][]} input_ids The input ids. + * @param {Tensor} logits The logits to process. + * @returns {Tensor} + */ + _call(input_ids, logits) { + if (logits.dims.at(0) !== 1) { + throw new Error('`response_format` currently supports batch_size=1 only.'); + } + + if (this.completed) { + return logits; + } + + let result; + try { + result = callRuntimeMethod(this.interpreter, 'computeMask', 'compute_mask'); + } catch (error) { + if (!isComputeAfterStopError(error)) { + throw error; + } + this.completed = true; + return logits; + } + + if (result.stop) { + this.completed = true; + return logits; + } + + const mask = result.mask ?? result; + const vocab_size = logits.dims.at(-1); + + for (let i = 0; i < vocab_size; ++i) { + if (!isAllowed(mask, i, vocab_size)) { + logits.data[i] = -Infinity; + } + } + + return logits; + } + + /** + * @param {number} token_id The sampled token ID. + */ + onTokenSampled(token_id) { + if (this.completed || token_id === this.eos_token_id) { + return; + } + + const result = callRuntimeMethod(this.interpreter, 'commitToken', 'commit_token', token_id); + if (result?.stop) { + this.completed = true; + } + } +} + // DEPRECATED: https://github.com/huggingface/transformers/pull/29485 // /** // * A logits processor that forces a specific token to be generated by the decoder. diff --git a/packages/transformers/src/models/modeling_utils.js b/packages/transformers/src/models/modeling_utils.js index ec9f17487..f4fc9b345 100644 --- a/packages/transformers/src/models/modeling_utils.js +++ b/packages/transformers/src/models/modeling_utils.js @@ -989,6 +989,10 @@ export class PreTrainedModel extends Callable { const logits = outputs.logits.slice(null, -1, null).to('float32'); const next_tokens_scores = prepared_logits_processor(all_input_ids, logits); + const processor_stop_before_sample = prepared_logits_processor.shouldStop(all_input_ids); + if (processor_stop_before_sample?.every((x) => x)) { + break; + } /** @type {[bigint][]} */ const generated_input_ids = []; @@ -1004,6 +1008,7 @@ export class PreTrainedModel extends Callable { // update generated ids, model inputs, and length for next step scores[batch_idx] += logProb; all_input_ids[batch_idx].push(bigint); + prepared_logits_processor.onTokenSampled(Number(newTokenId), batch_idx, all_input_ids); generated_input_ids.push([bigint]); // TODO: Support beam search @@ -1015,6 +1020,12 @@ export class PreTrainedModel extends Callable { } const stop = prepared_stopping_criteria(all_input_ids); + const processor_stop = prepared_logits_processor.shouldStop(all_input_ids); + if (processor_stop) { + for (let i = 0; i < stop.length; ++i) { + stop[i] ||= processor_stop[i]; + } + } if (stop.every((x) => x)) { break; } diff --git a/packages/transformers/src/pipelines/text-generation.js b/packages/transformers/src/pipelines/text-generation.js index fb9871ac9..c80e16c93 100644 --- a/packages/transformers/src/pipelines/text-generation.js +++ b/packages/transformers/src/pipelines/text-generation.js @@ -2,6 +2,7 @@ import { Pipeline } from './_base.js'; import { Tensor } from '../utils/tensor.js'; import { pick } from '../utils/core.js'; +import { LogitsProcessorList, SchemaConstrainedLogitsProcessor } from '../generation/logits_process.js'; /** * @typedef {import('./_base.js').TextPipelineConstructorArgs} TextPipelineConstructorArgs @@ -13,6 +14,17 @@ function isChat(x) { return Array.isArray(x) && x.every((x) => 'role' in x && 'content' in x); } +function trimToJSONPrefix(text) { + for (let i = text.length; i > 0; --i) { + const prefix = text.slice(0, i).trimEnd(); + try { + JSON.parse(prefix); + return prefix; + } catch {} + } + return text; +} + /** * @typedef {Object} TextGenerationSingleString * @property {string} generated_text The generated text. @@ -31,6 +43,7 @@ function isChat(x) { * @property {Object[]|null} [tools=null] A list of tools to expose to chat templates that support tool use. * @property {Record[]|null} [documents=null] A list of documents to expose to chat templates that support RAG. * @property {string|null} [chat_template=null] A specific chat template (or template name) to apply. + * @property {Object|null} [response_format=null] OpenAI-compatible structured output constraint. * @property {Object} [tokenizer_encode_kwargs] Additional keyword arguments to pass along to the encoding step of the tokenizer. * If the text input is a chat, it is passed to `apply_chat_template`. Otherwise, it is passed to the tokenizer's call function. * @typedef {import('../generation/parameters.js').GenerationFunctionParameters & TextGenerationSpecificParams} TextGenerationConfig @@ -108,6 +121,7 @@ export class TextGenerationPipeline tools, documents, chat_template, + response_format, tokenizer_encode_kwargs, ...generation_kwargs } = generate_kwargs; @@ -171,7 +185,28 @@ export class TextGenerationPipeline ...tokenizer_kwargs, }); - const outputTokenIds = /** @type {Tensor} */ ( + let response_format_eos_token_ids = null; + if (response_format) { + if (isBatched) { + throw new Error('`response_format` currently supports batch_size=1 only.'); + } + + const eos_token_id = generation_kwargs.eos_token_id ?? this.model.generation_config?.eos_token_id; + response_format_eos_token_ids = [ + ...(Array.isArray(eos_token_id) ? eos_token_id : [eos_token_id]), + this.tokenizer.eos_token_id, + ].filter(Number.isInteger); + const processor = await SchemaConstrainedLogitsProcessor.fromResponseFormat( + response_format, + this.tokenizer, + null, + response_format_eos_token_ids[0] ?? null, + ); + generation_kwargs.logits_processor ??= new LogitsProcessorList(); + generation_kwargs.logits_processor.push(processor); + } + + let outputTokenIds = /** @type {Tensor} */ ( await this.model.generate({ ...text_inputs, ...this._default_generation_config, @@ -179,6 +214,23 @@ export class TextGenerationPipeline }) ); + if (response_format_eos_token_ids) { + let numTrailingStopTokens = 0; + for (let i = outputTokenIds.data.length - 1; i >= 0; --i) { + if (!response_format_eos_token_ids.includes(Number(outputTokenIds.data[i]))) { + break; + } + ++numTrailingStopTokens; + } + + if (numTrailingStopTokens > 0) { + outputTokenIds = new Tensor(outputTokenIds.type, outputTokenIds.data.slice(0, -numTrailingStopTokens), [ + outputTokenIds.dims[0], + outputTokenIds.dims[1] - numTrailingStopTokens, + ]); + } + } + const decoded = this.tokenizer.batch_decode(outputTokenIds, { skip_special_tokens: true, }); @@ -201,6 +253,9 @@ export class TextGenerationPipeline // Trim the decoded text to only include the generated part decoded[i] = decoded[i].slice(promptLengths[textIndex]); } + if (response_format) { + decoded[i] = trimToJSONPrefix(decoded[i]); + } toReturn[textIndex].push( /** @type {TextGenerationSingle} */ ({ generated_text: isChatInput diff --git a/packages/transformers/tests/utils/generation.test.js b/packages/transformers/tests/utils/generation.test.js index 231985571..87b63defc 100644 --- a/packages/transformers/tests/utils/generation.test.js +++ b/packages/transformers/tests/utils/generation.test.js @@ -11,6 +11,7 @@ import { // Other TextStreamer, DynamicCache, + LogitsProcessor, random, full, } from "../../src/transformers.js"; @@ -208,6 +209,43 @@ describe("Generation parameters", () => { MAX_TEST_EXECUTION_TIME, ); + it( + "calls logits processor post-sample hook", + async () => { + class RecordingLogitsProcessor extends LogitsProcessor { + sampled = []; + + _call(input_ids, logits) { + return logits; + } + + onTokenSampled(token_id, batch_idx, input_ids) { + this.sampled.push({ + token_id, + batch_idx, + last_token_id: input_ids[batch_idx].at(-1), + }); + } + } + + const processor = new RecordingLogitsProcessor(); + const outputs = await generate(model, tokenizer, DUMMY_TEXT, { + max_new_tokens: 3, + logits_processor: [processor], + }); + + const generated_tokens = outputs.tolist()[0].slice(-3).map(Number); + expect(processor.sampled).toEqual( + generated_tokens.map((token_id) => ({ + token_id, + batch_idx: 0, + last_token_id: BigInt(token_id), + })), + ); + }, + MAX_TEST_EXECUTION_TIME, + ); + afterAll(async () => { await model?.dispose(); }, MAX_MODEL_DISPOSE_TIME); diff --git a/packages/transformers/tests/utils/logits_process.test.js b/packages/transformers/tests/utils/logits_process.test.js index 86ab06f2a..a988f6be8 100644 --- a/packages/transformers/tests/utils/logits_process.test.js +++ b/packages/transformers/tests/utils/logits_process.test.js @@ -2,6 +2,8 @@ import { // Pipelines pipeline, TextGenerationPipeline, + SchemaConstrainedLogitsProcessor, + Tensor, } from "../../src/transformers.js"; import { init } from "../init.js"; @@ -110,3 +112,125 @@ describe("Logits Processors", () => { }, MAX_MODEL_DISPOSE_TIME); }); }); + +describe("SchemaConstrainedLogitsProcessor", () => { + const tokenizer = { + get_vocab() { + return { "{": 0, "}": 1, foo: 2, bar: 3 }; + }, + decode(ids) { + return ["{", "}", "foo", "bar"][ids[0]]; + }, + eos_token_id: null, + bos_token_id: null, + all_special_ids: [], + }; + + it("applies llguidance masks and commits sampled tokens", async () => { + const committed = []; + const runtime = { + createTokenizer(config) { + expect(config.tokens.length).toEqual(4); + return { config }; + }, + createInterpreter({ tokenizer: llguidanceTokenizer, response_format }) { + expect(llguidanceTokenizer.config.tokens.length).toEqual(4); + expect(response_format).toEqual({ type: "json_object" }); + return { + computeMask() { + return { mask: [1, 0, 1, 0] }; + }, + commitToken(token_id) { + committed.push(token_id); + }, + }; + }, + }; + + const processor = await SchemaConstrainedLogitsProcessor.fromResponseFormat({ type: "json_object" }, tokenizer, runtime); + const logits = new Tensor("float32", new Float32Array([1, 2, 3, 4]), [1, 4]); + + processor([[0n]], logits); + processor.onTokenSampled(2); + + expect(Array.from(logits.data)).toEqual([1, -Infinity, 3, -Infinity]); + expect(committed).toEqual([2]); + }); + + it("passes callable tokenizers directly to llguidance", async () => { + const callableTokenizer = Object.assign(() => {}, tokenizer); + const runtime = { + acceptsTokenizerObjects: true, + createTokenizer() { + throw new Error("createTokenizer should not be called"); + }, + createInterpreter({ tokenizer: llguidanceTokenizer }) { + expect(llguidanceTokenizer).toBe(callableTokenizer); + expect(typeof llguidanceTokenizer).toEqual("function"); + expect(llguidanceTokenizer.get_vocab()).toEqual(tokenizer.get_vocab()); + return { + computeMask() { + return { mask: [1, 1, 1, 1] }; + }, + commitToken() {}, + }; + }, + }; + + await expect(SchemaConstrainedLogitsProcessor.fromResponseFormat({ type: "json_object" }, callableTokenizer, runtime)).resolves.toBeInstanceOf(SchemaConstrainedLogitsProcessor); + }); + + it("stops after llguidance reaches a stop state", async () => { + const runtime = { + createTokenizer(config) { + return { config }; + }, + createInterpreter() { + return { + computeMask() { + throw new Error("computeMask should not be called after stop"); + }, + commitToken() { + return { stop: true }; + }, + }; + }, + }; + + const processor = await SchemaConstrainedLogitsProcessor.fromResponseFormat({ type: "json_object" }, { ...tokenizer, eos_token_id: 1 }, runtime); + processor.onTokenSampled(2); + + const logits = new Tensor("float32", new Float32Array([1, 2, 3, 4]), [1, 4]); + processor([[0n]], logits); + + expect(Array.from(logits.data)).toEqual([1, 2, 3, 4]); + expect(processor.shouldStop([[0n]])).toEqual([true]); + }); + + it("stops when llguidance reports compute after stop", async () => { + const runtime = { + createTokenizer(config) { + return { config }; + }, + createInterpreter() { + return { + computeMask() { + throw new Error("computeMask failed: compute_mask() called after stop"); + }, + commitToken() {}, + }; + }, + }; + + const processor = await SchemaConstrainedLogitsProcessor.fromResponseFormat({ type: "json_object" }, { ...tokenizer, eos_token_id: 1 }, runtime); + const logits = new Tensor("float32", new Float32Array([1, 2, 3, 4]), [1, 4]); + processor([[0n]], logits); + + expect(Array.from(logits.data)).toEqual([1, 2, 3, 4]); + expect(processor.shouldStop([[0n]])).toEqual([true]); + }); + + it("validates response_format before loading llguidance", async () => { + await expect(SchemaConstrainedLogitsProcessor.fromResponseFormat({ type: "text" }, tokenizer)).rejects.toThrow("Unsupported `response_format.type`"); + }); +}); diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index bea34792e..bd219794d 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -23,6 +23,9 @@ importers: '@huggingface/tokenizers': specifier: ^0.1.3 version: 0.1.3 + llguidance: + specifier: 0.1.4 + version: 0.1.4 onnxruntime-node: specifier: 1.24.3 version: 1.24.3 @@ -1560,6 +1563,9 @@ packages: linkify-it@5.0.0: resolution: {integrity: sha512-5aHCbzQRADcdP+ATqnDuhhJ/MRIqDkZX5pyjFHRRysS8vZ5AbqGEoFIb6pYHPZ+L/OC2Lc+xT8uHVVR5CAK/wQ==} + llguidance@0.1.4: + resolution: {integrity: sha512-JYeeWyt+QxMEY1s55qWZXwwGVkCh2PWtiGzTbnG9YxPbGFfbs2Du9Z4wtifjs4cn2olEduePsCPV77V8LiHV8g==} + locate-path@5.0.0: resolution: {integrity: sha512-t7hw9pI+WvuwNJXwk5zVHpyhIqzg2qTlklJOf0mVxGSbe3Fp2VieZcduNYjaLDoy6p9uGpQEGWG87WpMKlNq8g==} engines: {node: '>=8'} @@ -3722,6 +3728,8 @@ snapshots: dependencies: uc.micro: 2.1.0 + llguidance@0.1.4: {} + locate-path@5.0.0: dependencies: p-locate: 4.1.0 diff --git a/test/response_format.js b/test/response_format.js new file mode 100644 index 000000000..2588b934f --- /dev/null +++ b/test/response_format.js @@ -0,0 +1,84 @@ +import { pipeline } from "../packages/transformers/src/transformers.js"; + +const colorMessages = [ + { + role: "user", + content: + "Return a JSON array of three short color names. Do not include any extra text.", + }, +]; + +const colorSchema = { + type: "array", + items: { type: "string" }, +}; + +const bookMessages = [ + { + role: "user", + content: + "Return a JSON array of three classic science fiction books. Include the title, author, and a short one-sentence summary for each book. Do not include any extra text.", + }, +]; + +const bookSchema = { + type: "array", + items: { + type: "object", + properties: { + title: { type: "string" }, + author: { type: "string" }, + summary: { type: "string" }, + }, + required: ["title", "author", "summary"], + additionalProperties: false, + }, +}; + +let progress = 0; +const pipe = await pipeline( + "text-generation", + "onnx-community/gemma-4-E2B-it-ONNX", + { + device: "webgpu", + dtype: "q4f16", + progress_callback: (i) => { + if (i.status === "progress_total") { + const p = Math.round(i.progress); + if (p !== progress) { + console.log(p); + progress = p; + } + } + }, + }, +); + +try { + for (const { label, messages, schema, max_new_tokens } of [ + { + label: "Colors", + messages: colorMessages, + schema: colorSchema, + max_new_tokens: 1024, + }, + { + label: "Books", + messages: bookMessages, + schema: bookSchema, + max_new_tokens: 2048, + }, + ]) { + const output = await pipe(messages, { + max_new_tokens, + response_format: { type: "json_schema", json_schema: schema }, + }); + + const generated = output[0].generated_text.at(-1).content; + console.log(`\n${label}:`); + console.log(generated); + console.log(JSON.parse(generated)); + } +} finally { + await pipe.dispose(); +} From a1c66f8ee63a5deeffdd7ebd664cafbc0df3fa50 Mon Sep 17 00:00:00 2001 From: Nico Martin Date: Thu, 9 Jul 2026 13:11:23 +0200 Subject: [PATCH 2/4] removed plan --- PLAN.local.md | 202 -------------------------------------------------- 1 file changed, 202 deletions(-) delete mode 100644 PLAN.local.md diff --git a/PLAN.local.md b/PLAN.local.md deleted file mode 100644 index 0deb8e5f8..000000000 --- a/PLAN.local.md +++ /dev/null @@ -1,202 +0,0 @@ -# Plan: `response_format` (JSON Schema constrained decoding) in Transformers.js - -## Goal - -Add OpenAI-compatible structured output support to Transformers.js, backed by -[llguidance](https://github.com/guidance-ai/llguidance) (Rust, compiled to WASM) for -grammar-constrained decoding. Guarantees schema-valid JSON output instead of -prompt-and-hope + post-hoc validation. - -## Target API - -```js -const pipe = await pipeline( - "text-generation", - "onnx-community/gemma-4-E2B-it-ONNX", - { - device: "webgpu", - dtype: "q4f16", - } -); - -const schema = { - type: "array", - items: { type: "string" }, -}; - -await pipe(messages, { - max_new_tokens: 1024, - response_format: { type: "json_schema", json_schema: schema }, -}); -``` - -- `response_format` is a new option accepted by the `pipeline()` call function (text-generation - pipeline first, generalizable later). -- Mirrors OpenAI's `response_format` shape so it's a drop-in mental model for JS devs. -- Internally resolves to attaching a constrained `LogitsProcessor` to the `generate()` call. -- `type: "json_object"` (schema-less, valid-JSON-only) should be supported too as a cheap - first milestone (no schema compilation needed, just JSON grammar). - -## Why llguidance - -Already covered in prior research — key points to keep in mind while implementing: - -- No pre-computation/startup cost (unlike Outlines), fast enough to compute masks on-the-fly - per step (~50μs on native CPU for a 128k vocab). -- Proven to run in a browser (merged into Chromium for `window.ai`'s `responseConstraint`), - though Chromium's actual wiring code is closed-source — we can't copy it, only use the - open-source llguidance engine itself. -- Tokenizer integration is a **one-time setup cost**, not a per-step cost: llguidance builds a - `TokTrie` from the full vocab (as raw bytes) once; per-step calls are just - `compute_mask()` → apply bitmask to logits → `commit_token()`. No decoding in the hot loop. - -## Architecture overview - -``` -response_format (user input) - │ - ▼ -schema → llguidance grammar compiler (WASM) ─┐ - │ │ (one-time per generate() call) - ▼ │ -tokenizer vocab → TokTrie (WASM) ─────────────┘ (one-time per model load, cached) - │ - ▼ -SchemaConstrainedLogitsProcessor (new LogitsProcessor subclass) - │ each decode step: - │ mask = interpreter.compute_mask() - │ logits[i] = -Infinity where mask bit is 0 - │ (after sampling) interpreter.commit_token(sampled_id) - ▼ -generate() loop (existing, unchanged) → LogitsProcessorList → sampler -``` - -## Phases - -### Phase 0 — Confirm/extend the plug-in point - -- Verify `generate()`'s public `logits_processor` option (appended after built-ins) is - sufficient, or needs a small extension to support processors that need a `commit_token`-style - callback *after* sampling (llguidance requires this — most existing `LogitsProcessor`s only - hook pre-sampling). -- Likely need a small addition to the generation loop: an optional post-sample hook - (e.g. `processor.onTokenSampled?.(tokenId)`), since `LogitsProcessorList` today assumes - stateless-after-`_call` processors. -- File: `src/generation/logits_process.js`, `src/models/modeling_utils.js` (generate loop). - -### Phase 1 — Spike: compile llguidance to WASM, measure feasibility - -- Compile llguidance's `parser` crate to `wasm32-unknown-unknown` (or `wasm32-wasi` if easier, - but prefer unknown-unknown for browser bundle compatibility). Check if a maintained - wasm-bindgen target already exists upstream before building our own. -- Measure and report: - - Bundle size (gzip) of the compiled `.wasm` — this is the #1 go/no-go metric for a - browser-shipped library. - - Per-step `compute_mask()` latency through the WASM boundary + JS marshaling (not just - native Rust numbers) for a realistic vocab size (Gemma/Llama ~256k, ~128k). - - One-time `TokTrie` build latency for a full model vocab. -- Decision gate: if bundle size or per-step latency is unacceptable, fall back to a - JSON-Schema-only hand-rolled JS grammar (Jsonformer-style: deterministic structural tokens, - masking only on leaf values) as a lighter-weight v1. Document the fallback trigger criteria - before starting Phase 2 so this isn't an open-ended detour later. - -### Phase 2 — Tokenizer → llguidance bridge (one-time per model) - -- Build an adapter that extracts, once per model/tokenizer load: - - `tokens: bytes[]` — every vocab entry as **raw UTF-8 bytes**, not the human-readable - decoded form. For byte-level BPE (GPT-2/Llama-style, e.g. `Ġgazed`), this means resolving - the byte-level encoding table to actual bytes, not just using `tokenizer.decode()`. - - `eos_token_id`, `bos_token_id`, `special_token_ids` — already available on Transformers.js - tokenizer objects, just need to be surfaced in the shape llguidance expects. -- Pass this to the WASM module once to construct the `LLTokenizer` / `TokTrie`. -- Cache the resulting handle per model instance (same caching pattern as `ModelRegistry`) — - rebuilding the trie per `generate()` call is wasted work if the model doesn't change. -- File: new `src/generation/grammar/tokenizer_bridge.js` (or similar). - -### Phase 3 — `SchemaConstrainedLogitsProcessor` - -```js -class SchemaConstrainedLogitsProcessor extends LogitsProcessor { - constructor(schema, llguidanceTokenizer) { - super(); - // compile schema -> llguidance grammar (LLInterpreter), once per generate() call - } - - _call(input_ids, logits) { - const { mask } = this.interpreter.compute_mask(); - // apply mask (bitset, vocab_size/32 elements) to logits typed array - return logits; - } - - onTokenSampled(tokenId) { - this.interpreter.commit_token(tokenId); - } -} -``` - -- Needs efficient bitmask application against Transformers.js's logits representation - (likely a `Tensor`/`Float32Array`) — avoid per-element JS loops if possible, use typed-array - ops. -- Handle the "grammar reached a stop state" signal from llguidance to force/allow EOS. -- File: `src/generation/logits_process.js`. - -### Phase 4 — Wire up `response_format` in the pipeline - -- Add `response_format` to the text-generation pipeline's call options - (`src/pipelines.js`, `TextGenerationPipeline`). -- Validation: - - `type: "json_object"` → JSON-only grammar, no schema compilation. - - `type: "json_schema"` → compile `json_schema` field via llguidance. - - Unsupported schema features → throw a clear error before generation starts (mirror Chrome's - `NotSupportedError` behavior — fail fast, not mid-generation). -- Translate into a `SchemaConstrainedLogitsProcessor` instance, append to the - `logits_processor` list passed into `generate()`. -- Decide default behavior on schema-includes-context-window cost: unlike Chrome, we probably - should **not** auto-inject the schema into the prompt by default (grammar constraint alone - should be sufficient and cheaper on context budget) — but expose an opt-in flag if empirically - output quality benefits from also showing the schema in-prompt (worth A/B testing once - working). - -### Phase 5 — Batch generation handling - -- llguidance's `Constraint`/`LLInterpreter` is single-sequence/stateful. -- For Transformers.js's batched `generate()`, need one interpreter instance **per batch item**, - each tracking its own grammar state and getting masked independently. -- Scope: v1 may explicitly restrict `response_format` to `batch_size === 1` and throw otherwise, - documented as a known limitation, with batched support as a fast-follow. - -### Phase 6 — Tests & docs - -- Unit tests: schema compilation, mask correctness on known small vocabs (reuse llguidance's - own JSON Schema Test Suite harness pattern if feasible). -- Integration test: run a small ONNX model end-to-end with a nested schema (object + array + - enum + required fields), assert `JSON.parse()` succeeds and matches schema on every run - (structural guarantee — should never flake). -- Perf test: measure generation-loop overhead with/without constrained decoding on WebGPU vs - WASM backend. -- Docs: new guide page + `response_format` API reference, modeled on OpenAI's docs since the - shape is intentionally familiar. - -## Open risks (carry into implementation, don't resolve in the plan) - -1. **WASM bundle size** — could be the single blocking issue; resolve in Phase 1 before - committing further engineering time. -2. **Post-sample hook** — `LogitsProcessorList`/`generate()` today may not have a clean seam for - `commit_token()`-style post-sampling callbacks; needs a small core API addition. -3. **Byte-level tokenizer edge cases** — getting the raw-bytes extraction wrong will silently - produce incorrect masks (accepting/rejecting wrong tokens) rather than an obvious crash — - needs solid test coverage, not just "it compiled." -4. **Batching** — v1 scoping to single-sequence generation is a real product limitation worth - flagging in the RFC/issue up front, not discovering mid-implementation. -5. **WebGPU/WASM boundary latency** — mask computation happening off the main compute path - (CPU, via WASM) needs to not stall the GPU forward pass; may need to run concurrently - (per llguidance's own recommendation: run `compute_mask()` while logits are being computed). - -## Suggested order of work for a first PR - -1. Phase 1 spike (throwaway branch, just prove bundle size + latency are acceptable). -2. Phase 0 + Phase 2 (core plug-in point + tokenizer bridge) — no user-facing API yet. -3. Phase 3 (`SchemaConstrainedLogitsProcessor`) with a manual/internal test using `generate()` - directly (no pipeline API yet). -4. Phase 4 (`response_format` on the pipeline) — first user-facing milestone. -5. Phase 5 (batching) and Phase 6 (tests/docs) as follow-ups, potentially separate PRs. \ No newline at end of file From b64ba3047ea221691117161f17d7409eea7ca82e Mon Sep 17 00:00:00 2001 From: Nico Martin Date: Fri, 10 Jul 2026 07:53:50 +0200 Subject: [PATCH 3/4] replaced response_schema with logits_processor --- packages/transformers/package.json | 1 - .../src/generation/grammar/llguidance.js | 54 ------ .../generation/grammar/tokenizer_bridge.js | 51 ------ .../src/generation/logits_process.js | 157 ------------------ .../transformers/src/models/modeling_utils.js | 6 +- .../src/pipelines/text-generation.js | 57 +------ .../tests/utils/generation.test.js | 33 +++- .../tests/utils/logits_process.test.js | 124 -------------- pnpm-lock.yaml | 32 +++- test/response_format.js | 84 ---------- 10 files changed, 62 insertions(+), 537 deletions(-) delete mode 100644 packages/transformers/src/generation/grammar/llguidance.js delete mode 100644 packages/transformers/src/generation/grammar/tokenizer_bridge.js delete mode 100644 test/response_format.js diff --git a/packages/transformers/package.json b/packages/transformers/package.json index 86f7819d7..8ea694b11 100644 --- a/packages/transformers/package.json +++ b/packages/transformers/package.json @@ -57,7 +57,6 @@ "dependencies": { "@huggingface/jinja": "^0.5.6", "@huggingface/tokenizers": "^0.1.3", - "llguidance": "0.1.4", "onnxruntime-node": "1.24.3", "onnxruntime-web": "1.26.0-dev.20260416-b7804b056c", "sharp": "^0.34.5" diff --git a/packages/transformers/src/generation/grammar/llguidance.js b/packages/transformers/src/generation/grammar/llguidance.js deleted file mode 100644 index fe5a1f23d..000000000 --- a/packages/transformers/src/generation/grammar/llguidance.js +++ /dev/null @@ -1,54 +0,0 @@ -/** - * @module generation/grammar/llguidance - */ - -let cachedRuntimePromise = null; -const normalizedRuntimeCache = new WeakMap(); - -function normalizeRuntime(runtime, options = {}) { - const cached = normalizedRuntimeCache.get(runtime); - if (cached) { - return cached; - } - - const createTokenizer = runtime.createTokenizer ?? runtime.create_tokenizer; - const createInterpreter = runtime.createInterpreter ?? runtime.create_interpreter; - - if (typeof createTokenizer !== 'function' || typeof createInterpreter !== 'function') { - throw new Error( - 'Invalid llguidance runtime: expected createTokenizer/createInterpreter functions exposed by the WASM module.', - ); - } - - const normalized = { ...runtime, ...options, createTokenizer, createInterpreter }; - normalizedRuntimeCache.set(runtime, normalized); - return normalized; -} - -async function loadBundledRuntime() { - const { loadBundledLLGuidance } = await import('llguidance'); - return normalizeRuntime(await loadBundledLLGuidance(), { acceptsTokenizerObjects: true }); -} - -/** - * Loads the llguidance WASM runtime. - * - * Consumers can provide a compatible runtime on - * `globalThis.__transformers_llguidance`; otherwise, the bundled `llguidance` - * package runtime is loaded lazily. - * - * @param {Object|null} runtime Optional runtime, primarily useful for tests or custom integrations. - * @returns {Promise} - */ -export async function loadLLGuidanceRuntime(runtime = null) { - if (runtime) { - return normalizeRuntime(runtime); - } - - if (globalThis.__transformers_llguidance) { - return normalizeRuntime(globalThis.__transformers_llguidance); - } - - cachedRuntimePromise ??= loadBundledRuntime(); - return cachedRuntimePromise; -} diff --git a/packages/transformers/src/generation/grammar/tokenizer_bridge.js b/packages/transformers/src/generation/grammar/tokenizer_bridge.js deleted file mode 100644 index 04233650b..000000000 --- a/packages/transformers/src/generation/grammar/tokenizer_bridge.js +++ /dev/null @@ -1,51 +0,0 @@ -/** - * @module generation/grammar/tokenizer_bridge - */ - -const encoder = new TextEncoder(); -const tokenizerCache = new WeakMap(); - -function decodeToken(tokenizer, token_id) { - try { - return tokenizer.decode([token_id], { skip_special_tokens: false, clean_up_tokenization_spaces: false }); - } catch { - return ''; - } -} - -/** - * Creates or returns the cached llguidance tokenizer handle for a Transformers.js tokenizer. - * - * @param {import('../../tokenization_utils.js').PreTrainedTokenizer} tokenizer The tokenizer to bridge. - * @param {Object} runtime The normalized llguidance runtime. - * @returns {Promise} - */ -export async function getLLGuidanceTokenizer(tokenizer, runtime) { - let runtimeCache = tokenizerCache.get(tokenizer); - if (!runtimeCache) { - runtimeCache = new WeakMap(); - tokenizerCache.set(tokenizer, runtimeCache); - } - - let cached = runtimeCache.get(runtime); - if (cached) { - return cached; - } - - cached = Promise.resolve().then(() => { - const vocab = tokenizer.get_vocab(); - const tokens = []; - for (const [token, id] of Object.entries(vocab)) { - tokens[id] = encoder.encode(decodeToken(tokenizer, id) || token); - } - - return runtime.createTokenizer({ - tokens, - eos_token_id: tokenizer.eos_token_id, - bos_token_id: tokenizer.bos_token_id, - special_token_ids: tokenizer.all_special_ids ?? [], - }); - }); - runtimeCache.set(runtime, cached); - return cached; -} diff --git a/packages/transformers/src/generation/logits_process.js b/packages/transformers/src/generation/logits_process.js index bc70e5469..4b32306c9 100644 --- a/packages/transformers/src/generation/logits_process.js +++ b/packages/transformers/src/generation/logits_process.js @@ -6,8 +6,6 @@ import { Callable } from '../utils/generic.js'; import { Tensor } from '../utils/tensor.js'; import { max, log_softmax } from '../utils/maths.js'; -import { loadLLGuidanceRuntime } from './grammar/llguidance.js'; -import { getLLGuidanceTokenizer } from './grammar/tokenizer_bridge.js'; /** * Abstract base class for all logit processors that can be applied during generation. @@ -167,161 +165,6 @@ export class LogitsProcessorList extends Callable { } } -function validateResponseFormat(response_format) { - if (!response_format || typeof response_format !== 'object') { - throw new Error('`response_format` must be an object.'); - } - - if (response_format.type === 'json_object') { - return; - } - - if (response_format.type === 'json_schema') { - if (!response_format.json_schema || typeof response_format.json_schema !== 'object') { - throw new Error('`response_format.json_schema` must be an object for type "json_schema".'); - } - return; - } - - throw new Error('Unsupported `response_format.type`. Expected "json_object" or "json_schema".'); -} - -function callRuntimeMethod(object, camelCaseName, snakeCaseName, ...args) { - const method = object[camelCaseName] ?? object[snakeCaseName]; - if (typeof method !== 'function') { - throw new Error(`Invalid llguidance runtime object: missing ${camelCaseName}/${snakeCaseName}.`); - } - return method.call(object, ...args); -} - -function isAllowed(mask, token_id, vocab_size) { - if (mask.length >= vocab_size) { - return Boolean(mask[token_id]); - } - - return Boolean(mask[token_id >> 5] & (1 << (token_id & 31))); -} - -function isComputeAfterStopError(error) { - return error instanceof Error && error.message.includes('compute_mask() called after stop'); -} - -/** - * Logits processor backed by llguidance for JSON-constrained generation. - */ -export class SchemaConstrainedLogitsProcessor extends LogitsProcessor { - /** - * @param {Object} interpreter A llguidance interpreter/constraint instance. - * @param {number|null} eos_token_id The EOS token ID to force when the grammar stops. - */ - constructor(interpreter, eos_token_id = null) { - super(); - this.interpreter = interpreter; - this.eos_token_id = eos_token_id; - this.completed = false; - } - - /** - * Creates a constrained logits processor from an OpenAI-compatible response_format object. - * - * @param {Object} response_format The requested response format. - * @param {import('../tokenization_utils.js').PreTrainedTokenizer} tokenizer The tokenizer used for generation. - * @param {Object|null} runtime Optional llguidance runtime override. - * @returns {Promise} - */ - static async fromResponseFormat(response_format, tokenizer, runtime = null, eos_token_id = null) { - validateResponseFormat(response_format); - - const llguidance = await loadLLGuidanceRuntime(runtime); - const llguidanceTokenizer = llguidance.acceptsTokenizerObjects - ? tokenizer - : await getLLGuidanceTokenizer(tokenizer, llguidance); - const interpreter = await llguidance.createInterpreter({ - tokenizer: llguidanceTokenizer, - response_format, - }); - - return new SchemaConstrainedLogitsProcessor(interpreter, eos_token_id ?? tokenizer.eos_token_id); - } - - /** - * @param {Tensor} logits The logits to process. - */ - forceEOS(logits) { - if (!Number.isInteger(this.eos_token_id)) { - throw new Error('`response_format` reached a stop state, but no EOS token is available to end generation.'); - } - - logits.data.fill(-Infinity); - logits.data[this.eos_token_id] = 0; - return logits; - } - - /** - * @param {bigint[][]} input_ids The input IDs. - * @returns {boolean[]} - */ - shouldStop(input_ids) { - return new Array(input_ids.length).fill(this.completed); - } - - /** - * @param {bigint[][]} input_ids The input ids. - * @param {Tensor} logits The logits to process. - * @returns {Tensor} - */ - _call(input_ids, logits) { - if (logits.dims.at(0) !== 1) { - throw new Error('`response_format` currently supports batch_size=1 only.'); - } - - if (this.completed) { - return logits; - } - - let result; - try { - result = callRuntimeMethod(this.interpreter, 'computeMask', 'compute_mask'); - } catch (error) { - if (!isComputeAfterStopError(error)) { - throw error; - } - this.completed = true; - return logits; - } - - if (result.stop) { - this.completed = true; - return logits; - } - - const mask = result.mask ?? result; - const vocab_size = logits.dims.at(-1); - - for (let i = 0; i < vocab_size; ++i) { - if (!isAllowed(mask, i, vocab_size)) { - logits.data[i] = -Infinity; - } - } - - return logits; - } - - /** - * @param {number} token_id The sampled token ID. - */ - onTokenSampled(token_id) { - if (this.completed || token_id === this.eos_token_id) { - return; - } - - const result = callRuntimeMethod(this.interpreter, 'commitToken', 'commit_token', token_id); - if (result?.stop) { - this.completed = true; - } - } -} - // DEPRECATED: https://github.com/huggingface/transformers/pull/29485 // /** // * A logits processor that forces a specific token to be generated by the decoder. diff --git a/packages/transformers/src/models/modeling_utils.js b/packages/transformers/src/models/modeling_utils.js index f4fc9b345..fbb3c68d1 100644 --- a/packages/transformers/src/models/modeling_utils.js +++ b/packages/transformers/src/models/modeling_utils.js @@ -539,7 +539,11 @@ export class PreTrainedModel extends Callable { } if (logits_processor !== null) { - processors.extend(logits_processor); + if (typeof logits_processor[Symbol.iterator] === 'function') { + processors.extend(logits_processor); + } else { + processors.push(logits_processor); + } } // `LogitNormalization` should always be the last logit processor, when present diff --git a/packages/transformers/src/pipelines/text-generation.js b/packages/transformers/src/pipelines/text-generation.js index c80e16c93..fb9871ac9 100644 --- a/packages/transformers/src/pipelines/text-generation.js +++ b/packages/transformers/src/pipelines/text-generation.js @@ -2,7 +2,6 @@ import { Pipeline } from './_base.js'; import { Tensor } from '../utils/tensor.js'; import { pick } from '../utils/core.js'; -import { LogitsProcessorList, SchemaConstrainedLogitsProcessor } from '../generation/logits_process.js'; /** * @typedef {import('./_base.js').TextPipelineConstructorArgs} TextPipelineConstructorArgs @@ -14,17 +13,6 @@ function isChat(x) { return Array.isArray(x) && x.every((x) => 'role' in x && 'content' in x); } -function trimToJSONPrefix(text) { - for (let i = text.length; i > 0; --i) { - const prefix = text.slice(0, i).trimEnd(); - try { - JSON.parse(prefix); - return prefix; - } catch {} - } - return text; -} - /** * @typedef {Object} TextGenerationSingleString * @property {string} generated_text The generated text. @@ -43,7 +31,6 @@ function trimToJSONPrefix(text) { * @property {Object[]|null} [tools=null] A list of tools to expose to chat templates that support tool use. * @property {Record[]|null} [documents=null] A list of documents to expose to chat templates that support RAG. * @property {string|null} [chat_template=null] A specific chat template (or template name) to apply. - * @property {Object|null} [response_format=null] OpenAI-compatible structured output constraint. * @property {Object} [tokenizer_encode_kwargs] Additional keyword arguments to pass along to the encoding step of the tokenizer. * If the text input is a chat, it is passed to `apply_chat_template`. Otherwise, it is passed to the tokenizer's call function. * @typedef {import('../generation/parameters.js').GenerationFunctionParameters & TextGenerationSpecificParams} TextGenerationConfig @@ -121,7 +108,6 @@ export class TextGenerationPipeline tools, documents, chat_template, - response_format, tokenizer_encode_kwargs, ...generation_kwargs } = generate_kwargs; @@ -185,28 +171,7 @@ export class TextGenerationPipeline ...tokenizer_kwargs, }); - let response_format_eos_token_ids = null; - if (response_format) { - if (isBatched) { - throw new Error('`response_format` currently supports batch_size=1 only.'); - } - - const eos_token_id = generation_kwargs.eos_token_id ?? this.model.generation_config?.eos_token_id; - response_format_eos_token_ids = [ - ...(Array.isArray(eos_token_id) ? eos_token_id : [eos_token_id]), - this.tokenizer.eos_token_id, - ].filter(Number.isInteger); - const processor = await SchemaConstrainedLogitsProcessor.fromResponseFormat( - response_format, - this.tokenizer, - null, - response_format_eos_token_ids[0] ?? null, - ); - generation_kwargs.logits_processor ??= new LogitsProcessorList(); - generation_kwargs.logits_processor.push(processor); - } - - let outputTokenIds = /** @type {Tensor} */ ( + const outputTokenIds = /** @type {Tensor} */ ( await this.model.generate({ ...text_inputs, ...this._default_generation_config, @@ -214,23 +179,6 @@ export class TextGenerationPipeline }) ); - if (response_format_eos_token_ids) { - let numTrailingStopTokens = 0; - for (let i = outputTokenIds.data.length - 1; i >= 0; --i) { - if (!response_format_eos_token_ids.includes(Number(outputTokenIds.data[i]))) { - break; - } - ++numTrailingStopTokens; - } - - if (numTrailingStopTokens > 0) { - outputTokenIds = new Tensor(outputTokenIds.type, outputTokenIds.data.slice(0, -numTrailingStopTokens), [ - outputTokenIds.dims[0], - outputTokenIds.dims[1] - numTrailingStopTokens, - ]); - } - } - const decoded = this.tokenizer.batch_decode(outputTokenIds, { skip_special_tokens: true, }); @@ -253,9 +201,6 @@ export class TextGenerationPipeline // Trim the decoded text to only include the generated part decoded[i] = decoded[i].slice(promptLengths[textIndex]); } - if (response_format) { - decoded[i] = trimToJSONPrefix(decoded[i]); - } toReturn[textIndex].push( /** @type {TextGenerationSingle} */ ({ generated_text: isChatInput diff --git a/packages/transformers/tests/utils/generation.test.js b/packages/transformers/tests/utils/generation.test.js index 87b63defc..b3589c7ce 100644 --- a/packages/transformers/tests/utils/generation.test.js +++ b/packages/transformers/tests/utils/generation.test.js @@ -231,7 +231,7 @@ describe("Generation parameters", () => { const processor = new RecordingLogitsProcessor(); const outputs = await generate(model, tokenizer, DUMMY_TEXT, { max_new_tokens: 3, - logits_processor: [processor], + logits_processor: processor, }); const generated_tokens = outputs.tolist()[0].slice(-3).map(Number); @@ -246,6 +246,37 @@ describe("Generation parameters", () => { MAX_TEST_EXECUTION_TIME, ); + it( + "supports logits processor stopping hook", + async () => { + class StopAfterTokenLogitsProcessor extends LogitsProcessor { + stopped = false; + + _call(input_ids, logits) { + return logits; + } + + onTokenSampled() { + this.stopped = true; + } + + shouldStop(input_ids) { + return new Array(input_ids.length).fill(this.stopped); + } + } + + const processor = new StopAfterTokenLogitsProcessor(); + const outputs = await generate(model, tokenizer, DUMMY_TEXT, { + max_new_tokens: 5, + logits_processor: processor, + }); + + // 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/packages/transformers/tests/utils/logits_process.test.js b/packages/transformers/tests/utils/logits_process.test.js index a988f6be8..86ab06f2a 100644 --- a/packages/transformers/tests/utils/logits_process.test.js +++ b/packages/transformers/tests/utils/logits_process.test.js @@ -2,8 +2,6 @@ import { // Pipelines pipeline, TextGenerationPipeline, - SchemaConstrainedLogitsProcessor, - Tensor, } from "../../src/transformers.js"; import { init } from "../init.js"; @@ -112,125 +110,3 @@ describe("Logits Processors", () => { }, MAX_MODEL_DISPOSE_TIME); }); }); - -describe("SchemaConstrainedLogitsProcessor", () => { - const tokenizer = { - get_vocab() { - return { "{": 0, "}": 1, foo: 2, bar: 3 }; - }, - decode(ids) { - return ["{", "}", "foo", "bar"][ids[0]]; - }, - eos_token_id: null, - bos_token_id: null, - all_special_ids: [], - }; - - it("applies llguidance masks and commits sampled tokens", async () => { - const committed = []; - const runtime = { - createTokenizer(config) { - expect(config.tokens.length).toEqual(4); - return { config }; - }, - createInterpreter({ tokenizer: llguidanceTokenizer, response_format }) { - expect(llguidanceTokenizer.config.tokens.length).toEqual(4); - expect(response_format).toEqual({ type: "json_object" }); - return { - computeMask() { - return { mask: [1, 0, 1, 0] }; - }, - commitToken(token_id) { - committed.push(token_id); - }, - }; - }, - }; - - const processor = await SchemaConstrainedLogitsProcessor.fromResponseFormat({ type: "json_object" }, tokenizer, runtime); - const logits = new Tensor("float32", new Float32Array([1, 2, 3, 4]), [1, 4]); - - processor([[0n]], logits); - processor.onTokenSampled(2); - - expect(Array.from(logits.data)).toEqual([1, -Infinity, 3, -Infinity]); - expect(committed).toEqual([2]); - }); - - it("passes callable tokenizers directly to llguidance", async () => { - const callableTokenizer = Object.assign(() => {}, tokenizer); - const runtime = { - acceptsTokenizerObjects: true, - createTokenizer() { - throw new Error("createTokenizer should not be called"); - }, - createInterpreter({ tokenizer: llguidanceTokenizer }) { - expect(llguidanceTokenizer).toBe(callableTokenizer); - expect(typeof llguidanceTokenizer).toEqual("function"); - expect(llguidanceTokenizer.get_vocab()).toEqual(tokenizer.get_vocab()); - return { - computeMask() { - return { mask: [1, 1, 1, 1] }; - }, - commitToken() {}, - }; - }, - }; - - await expect(SchemaConstrainedLogitsProcessor.fromResponseFormat({ type: "json_object" }, callableTokenizer, runtime)).resolves.toBeInstanceOf(SchemaConstrainedLogitsProcessor); - }); - - it("stops after llguidance reaches a stop state", async () => { - const runtime = { - createTokenizer(config) { - return { config }; - }, - createInterpreter() { - return { - computeMask() { - throw new Error("computeMask should not be called after stop"); - }, - commitToken() { - return { stop: true }; - }, - }; - }, - }; - - const processor = await SchemaConstrainedLogitsProcessor.fromResponseFormat({ type: "json_object" }, { ...tokenizer, eos_token_id: 1 }, runtime); - processor.onTokenSampled(2); - - const logits = new Tensor("float32", new Float32Array([1, 2, 3, 4]), [1, 4]); - processor([[0n]], logits); - - expect(Array.from(logits.data)).toEqual([1, 2, 3, 4]); - expect(processor.shouldStop([[0n]])).toEqual([true]); - }); - - it("stops when llguidance reports compute after stop", async () => { - const runtime = { - createTokenizer(config) { - return { config }; - }, - createInterpreter() { - return { - computeMask() { - throw new Error("computeMask failed: compute_mask() called after stop"); - }, - commitToken() {}, - }; - }, - }; - - const processor = await SchemaConstrainedLogitsProcessor.fromResponseFormat({ type: "json_object" }, { ...tokenizer, eos_token_id: 1 }, runtime); - const logits = new Tensor("float32", new Float32Array([1, 2, 3, 4]), [1, 4]); - processor([[0n]], logits); - - expect(Array.from(logits.data)).toEqual([1, 2, 3, 4]); - expect(processor.shouldStop([[0n]])).toEqual([true]); - }); - - it("validates response_format before loading llguidance", async () => { - await expect(SchemaConstrainedLogitsProcessor.fromResponseFormat({ type: "text" }, tokenizer)).rejects.toThrow("Unsupported `response_format.type`"); - }); -}); diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index bd219794d..3ed937a04 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -23,9 +23,6 @@ importers: '@huggingface/tokenizers': specifier: ^0.1.3 version: 0.1.3 - llguidance: - specifier: 0.1.4 - version: 0.1.4 onnxruntime-node: specifier: 1.24.3 version: 1.24.3 @@ -430,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==} @@ -803,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==} @@ -1563,9 +1584,6 @@ packages: linkify-it@5.0.0: resolution: {integrity: sha512-5aHCbzQRADcdP+ATqnDuhhJ/MRIqDkZX5pyjFHRRysS8vZ5AbqGEoFIb6pYHPZ+L/OC2Lc+xT8uHVVR5CAK/wQ==} - llguidance@0.1.4: - resolution: {integrity: sha512-JYeeWyt+QxMEY1s55qWZXwwGVkCh2PWtiGzTbnG9YxPbGFfbs2Du9Z4wtifjs4cn2olEduePsCPV77V8LiHV8g==} - locate-path@5.0.0: resolution: {integrity: sha512-t7hw9pI+WvuwNJXwk5zVHpyhIqzg2qTlklJOf0mVxGSbe3Fp2VieZcduNYjaLDoy6p9uGpQEGWG87WpMKlNq8g==} engines: {node: '>=8'} @@ -3728,8 +3746,6 @@ snapshots: dependencies: uc.micro: 2.1.0 - llguidance@0.1.4: {} - locate-path@5.0.0: dependencies: p-locate: 4.1.0 diff --git a/test/response_format.js b/test/response_format.js deleted file mode 100644 index 2588b934f..000000000 --- a/test/response_format.js +++ /dev/null @@ -1,84 +0,0 @@ -import { pipeline } from "../packages/transformers/src/transformers.js"; - -const colorMessages = [ - { - role: "user", - content: - "Return a JSON array of three short color names. Do not include any extra text.", - }, -]; - -const colorSchema = { - type: "array", - items: { type: "string" }, -}; - -const bookMessages = [ - { - role: "user", - content: - "Return a JSON array of three classic science fiction books. Include the title, author, and a short one-sentence summary for each book. Do not include any extra text.", - }, -]; - -const bookSchema = { - type: "array", - items: { - type: "object", - properties: { - title: { type: "string" }, - author: { type: "string" }, - summary: { type: "string" }, - }, - required: ["title", "author", "summary"], - additionalProperties: false, - }, -}; - -let progress = 0; -const pipe = await pipeline( - "text-generation", - "onnx-community/gemma-4-E2B-it-ONNX", - { - device: "webgpu", - dtype: "q4f16", - progress_callback: (i) => { - if (i.status === "progress_total") { - const p = Math.round(i.progress); - if (p !== progress) { - console.log(p); - progress = p; - } - } - }, - }, -); - -try { - for (const { label, messages, schema, max_new_tokens } of [ - { - label: "Colors", - messages: colorMessages, - schema: colorSchema, - max_new_tokens: 1024, - }, - { - label: "Books", - messages: bookMessages, - schema: bookSchema, - max_new_tokens: 2048, - }, - ]) { - const output = await pipe(messages, { - max_new_tokens, - response_format: { type: "json_schema", json_schema: schema }, - }); - - const generated = output[0].generated_text.at(-1).content; - console.log(`\n${label}:`); - console.log(generated); - console.log(JSON.parse(generated)); - } -} finally { - await pipe.dispose(); -} From 56a8114a89100ef4419fe7d40992246d0a7cec10 Mon Sep 17 00:00:00 2001 From: Nico Martin Date: Fri, 10 Jul 2026 10:48:42 +0200 Subject: [PATCH 4/4] clean up --- .../src/generation/logits_process.js | 71 ++----------------- .../transformers/src/models/modeling_utils.js | 25 +++---- .../tests/utils/generation.test.js | 65 ++++++++--------- 3 files changed, 47 insertions(+), 114 deletions(-) diff --git a/packages/transformers/src/generation/logits_process.js b/packages/transformers/src/generation/logits_process.js index 4b32306c9..7bb9545eb 100644 --- a/packages/transformers/src/generation/logits_process.js +++ b/packages/transformers/src/generation/logits_process.js @@ -22,25 +22,6 @@ export class LogitsProcessor extends Callable { _call(input_ids, logits) { throw Error('`_call` should be implemented in a subclass'); } - - /** - * Optional hook called after a token has been selected by the sampler. - * - * @param {number} token_id The sampled token ID. - * @param {number} batch_idx The batch index that sampled the token. - * @param {bigint[][]} input_ids The input IDs after appending the sampled token. - */ - onTokenSampled(token_id, batch_idx, input_ids) {} - - /** - * Optional hook for processors that can terminate generation. - * - * @param {bigint[][]} input_ids The input IDs. - * @returns {boolean[]|null} - */ - shouldStop(input_ids) { - return null; - } } /** @@ -58,25 +39,6 @@ export class LogitsWarper extends Callable { _call(input_ids, logits) { throw Error('`_call` should be implemented in a subclass'); } - - /** - * Optional hook called after a token has been selected by the sampler. - * - * @param {number} token_id The sampled token ID. - * @param {number} batch_idx The batch index that sampled the token. - * @param {bigint[][]} input_ids The input IDs after appending the sampled token. - */ - onTokenSampled(token_id, batch_idx, input_ids) {} - - /** - * Optional hook for processors that can terminate generation. - * - * @param {bigint[][]} input_ids The input IDs. - * @returns {boolean[]|null} - */ - shouldStop(input_ids) { - return null; - } } /** @@ -127,37 +89,16 @@ export class LogitsProcessorList extends Callable { } /** - * Calls post-sampling hooks on processors that need to update state after token selection. + * 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_id The sampled token ID. - * @param {number} batch_idx The batch index that sampled the token. - * @param {bigint[][]} input_ids The input IDs after appending the sampled token. + * @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. */ - onTokenSampled(token_id, batch_idx, input_ids) { + onTokensSampled(token_ids, input_ids) { for (const processor of this.processors) { - processor.onTokenSampled?.(token_id, batch_idx, input_ids); - } - } - - /** - * Calls stopping hooks on processors that can terminate generation. - * - * @param {bigint[][]} input_ids The input IDs. - * @returns {boolean[]|null} - */ - shouldStop(input_ids) { - let stop = null; - for (const processor of this.processors) { - const processorStop = processor.shouldStop?.(input_ids); - if (!processorStop) { - continue; - } - stop ??= new Array(input_ids.length).fill(false); - for (let i = 0; i < stop.length; ++i) { - stop[i] ||= processorStop[i]; - } + processor.onTokensSampled?.(token_ids, input_ids); } - return stop; } [Symbol.iterator]() { diff --git a/packages/transformers/src/models/modeling_utils.js b/packages/transformers/src/models/modeling_utils.js index fbb3c68d1..f5caec12c 100644 --- a/packages/transformers/src/models/modeling_utils.js +++ b/packages/transformers/src/models/modeling_utils.js @@ -539,11 +539,7 @@ export class PreTrainedModel extends Callable { } if (logits_processor !== null) { - if (typeof logits_processor[Symbol.iterator] === 'function') { - processors.extend(logits_processor); - } else { - processors.push(logits_processor); - } + processors.extend(logits_processor); } // `LogitNormalization` should always be the last logit processor, when present @@ -993,10 +989,6 @@ export class PreTrainedModel extends Callable { const logits = outputs.logits.slice(null, -1, null).to('float32'); const next_tokens_scores = prepared_logits_processor(all_input_ids, logits); - const processor_stop_before_sample = prepared_logits_processor.shouldStop(all_input_ids); - if (processor_stop_before_sample?.every((x) => x)) { - break; - } /** @type {[bigint][]} */ const generated_input_ids = []; @@ -1011,25 +1003,24 @@ 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); - prepared_logits_processor.onTokenSampled(Number(newTokenId), batch_idx, all_input_ids); 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); } const stop = prepared_stopping_criteria(all_input_ids); - const processor_stop = prepared_logits_processor.shouldStop(all_input_ids); - if (processor_stop) { - for (let i = 0; i < stop.length; ++i) { - stop[i] ||= processor_stop[i]; - } - } if (stop.every((x) => x)) { break; } diff --git a/packages/transformers/tests/utils/generation.test.js b/packages/transformers/tests/utils/generation.test.js index b3589c7ce..db991affc 100644 --- a/packages/transformers/tests/utils/generation.test.js +++ b/packages/transformers/tests/utils/generation.test.js @@ -12,6 +12,8 @@ import { TextStreamer, DynamicCache, LogitsProcessor, + LogitsProcessorList, + StoppingCriteria, random, full, } from "../../src/transformers.js"; @@ -210,65 +212,64 @@ describe("Generation parameters", () => { ); it( - "calls logits processor post-sample hook", + "calls logits processor post-sample hook after full batch step", async () => { class RecordingLogitsProcessor extends LogitsProcessor { - sampled = []; + snapshots = []; _call(input_ids, logits) { return logits; } - onTokenSampled(token_id, batch_idx, input_ids) { - this.sampled.push({ - token_id, - batch_idx, - last_token_id: input_ids[batch_idx].at(-1), + onTokensSampled(token_ids, input_ids) { + this.snapshots.push({ + token_ids, + lengths: input_ids.map((ids) => ids.length), }); } } const processor = new RecordingLogitsProcessor(); - const outputs = await generate(model, tokenizer, DUMMY_TEXT, { - max_new_tokens: 3, - logits_processor: processor, + 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()[0].slice(-3).map(Number); - expect(processor.sampled).toEqual( - generated_tokens.map((token_id) => ({ - token_id, - batch_idx: 0, - last_token_id: BigInt(token_id), - })), - ); + 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 logits processor stopping hook", + "supports custom stopping criteria", async () => { - class StopAfterTokenLogitsProcessor extends LogitsProcessor { - stopped = false; - - _call(input_ids, logits) { - return logits; - } - - onTokenSampled() { - this.stopped = true; + class StopAfterLengthCriteria extends StoppingCriteria { + constructor(max_length) { + super(); + this.max_length = max_length; } - shouldStop(input_ids) { - return new Array(input_ids.length).fill(this.stopped); + _call(input_ids) { + return input_ids.map((ids) => ids.length >= this.max_length); } } - const processor = new StopAfterTokenLogitsProcessor(); const outputs = await generate(model, tokenizer, DUMMY_TEXT, { max_new_tokens: 5, - logits_processor: processor, + stopping_criteria: new StopAfterLengthCriteria(3), }); // BOS + DUMMY_TEXT + exactly one generated token