From b179ef029cb176f76e6e64f346026eb5735745f3 Mon Sep 17 00:00:00 2001 From: Nico Martin Date: Thu, 9 Jul 2026 12:40:55 +0200 Subject: [PATCH 01/17] 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 02/17] 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 03/17] 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 04/17] 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 From 00deb593c17307d3cabb2f231f29707796b49fef Mon Sep 17 00:00:00 2001 From: Nico Martin Date: Fri, 10 Jul 2026 11:27:55 +0200 Subject: [PATCH 05/17] first POC --- packages/transformers-llguidance/README.md | 28 ++ .../transformers-llguidance/jest.config.mjs | 11 + packages/transformers-llguidance/package.json | 68 +++++ .../transformers-llguidance/scripts/build.mjs | 26 ++ .../transformers-llguidance/scripts/dev.mjs | 34 +++ packages/transformers-llguidance/src/index.ts | 250 ++++++++++++++++++ .../tests/llguidance-constraint.test.js | 52 ++++ .../transformers-llguidance/tsconfig.json | 21 ++ pnpm-lock.yaml | 54 ++-- 9 files changed, 520 insertions(+), 24 deletions(-) create mode 100644 packages/transformers-llguidance/README.md create mode 100644 packages/transformers-llguidance/jest.config.mjs create mode 100644 packages/transformers-llguidance/package.json create mode 100644 packages/transformers-llguidance/scripts/build.mjs create mode 100644 packages/transformers-llguidance/scripts/dev.mjs create mode 100644 packages/transformers-llguidance/src/index.ts create mode 100644 packages/transformers-llguidance/tests/llguidance-constraint.test.js create mode 100644 packages/transformers-llguidance/tsconfig.json diff --git a/packages/transformers-llguidance/README.md b/packages/transformers-llguidance/README.md new file mode 100644 index 000000000..f1776465c --- /dev/null +++ b/packages/transformers-llguidance/README.md @@ -0,0 +1,28 @@ +# @huggingface/transformers-llguidance + +Experimental constrained-generation helpers for Transformers.js. + +This package exports `LlguidanceConstraint`, which turns an llguidance response format into the `logits_processor` and `stopping_criteria` objects accepted by Transformers.js generation. + +```js +import { LlguidanceConstraint } from "@huggingface/transformers-llguidance"; + +const { logits_processor, stopping_criteria } = + await LlguidanceConstraint.fromResponseFormat(tokenizer, { + type: "json_schema", + json_schema: { + type: "object", + properties: { + answer: { type: "string" }, + }, + required: ["answer"], + additionalProperties: false, + }, + }); + +await model.generate({ + ...inputs, + logits_processor, + stopping_criteria, +}); +``` diff --git a/packages/transformers-llguidance/jest.config.mjs b/packages/transformers-llguidance/jest.config.mjs new file mode 100644 index 000000000..b905ca752 --- /dev/null +++ b/packages/transformers-llguidance/jest.config.mjs @@ -0,0 +1,11 @@ +/** @type {import('jest').Config} */ +export default { + clearMocks: true, + collectCoverage: true, + coverageDirectory: "coverage", + coveragePathIgnorePatterns: ["node_modules", "tests"], + coverageProvider: "v8", + roots: ["./tests/"], + testTimeout: 32000, + transform: {}, +}; diff --git a/packages/transformers-llguidance/package.json b/packages/transformers-llguidance/package.json new file mode 100644 index 000000000..27fe44062 --- /dev/null +++ b/packages/transformers-llguidance/package.json @@ -0,0 +1,68 @@ +{ + "name": "@huggingface/transformers-llguidance", + "version": "0.0.0", + "description": "llguidance integration helpers for Transformers.js constrained generation", + "main": "./dist/index.cjs", + "types": "./types/index.d.ts", + "type": "module", + "exports": { + "import": { + "types": "./types/index.d.ts", + "default": "./dist/index.js" + }, + "require": { + "types": "./types/index.d.ts", + "default": "./dist/index.cjs" + } + }, + "scripts": { + "format": "prettier --write . --ignore-path ../../.prettierignore", + "format:check": "prettier --check . --ignore-path ../../.prettierignore", + "typegen": "tsc --build", + "dev": "node scripts/dev.mjs", + "build": "node scripts/build.mjs && pnpm typegen", + "test": "node --experimental-vm-modules node_modules/jest/bin/jest.js --verbose" + }, + "repository": { + "type": "git", + "url": "git+https://github.com/huggingface/transformers.js.git" + }, + "keywords": [ + "transformers", + "transformers.js", + "huggingface", + "llguidance", + "constrained-generation" + ], + "author": "Hugging Face", + "license": "Apache-2.0", + "bugs": { + "url": "https://github.com/huggingface/transformers.js/issues" + }, + "homepage": "https://github.com/huggingface/transformers.js#readme", + "peerDependencies": { + "@huggingface/transformers": "^4.2.0" + }, + "devDependencies": { + "@huggingface/transformers": "workspace:*", + "@types/jest": "^30.0.0", + "@types/node": "^24.1.0", + "esbuild": "^0.27.2", + "jest": "^30.2.0", + "typescript": "5.9.3" + }, + "files": [ + "src", + "dist", + "types", + "README.md", + "LICENSE", + "!**/*.tsbuildinfo" + ], + "publishConfig": { + "access": "public" + }, + "dependencies": { + "llguidance": "^0.1.7" + } +} diff --git a/packages/transformers-llguidance/scripts/build.mjs b/packages/transformers-llguidance/scripts/build.mjs new file mode 100644 index 000000000..9b5ca9f5a --- /dev/null +++ b/packages/transformers-llguidance/scripts/build.mjs @@ -0,0 +1,26 @@ +import { build } from "esbuild"; +import { rmSync } from "node:fs"; + +rmSync("dist", { recursive: true, force: true }); + +const common = { + entryPoints: ["src/index.ts"], + bundle: true, + platform: "neutral", + target: "es2022", + sourcemap: true, + external: ["@huggingface/transformers", "llguidance"], +}; + +await Promise.all([ + build({ + ...common, + format: "esm", + outfile: "dist/index.js", + }), + build({ + ...common, + format: "cjs", + outfile: "dist/index.cjs", + }), +]); diff --git a/packages/transformers-llguidance/scripts/dev.mjs b/packages/transformers-llguidance/scripts/dev.mjs new file mode 100644 index 000000000..4b75dfcad --- /dev/null +++ b/packages/transformers-llguidance/scripts/dev.mjs @@ -0,0 +1,34 @@ +import { context } from "esbuild"; +import { rmSync } from "node:fs"; + +rmSync("dist", { recursive: true, force: true }); + +const common = { + entryPoints: ["src/index.ts"], + bundle: true, + platform: "neutral", + target: "es2022", + sourcemap: true, + external: ["@huggingface/transformers", "llguidance"], +}; + +const contexts = await Promise.all([ + context({ + ...common, + format: "esm", + outfile: "dist/index.js", + }), + context({ + ...common, + format: "cjs", + outfile: "dist/index.cjs", + }), +]); + +await Promise.all(contexts.map((ctx) => ctx.watch())); +console.log("Watching @huggingface/transformers-llguidance..."); + +process.on("SIGINT", async () => { + await Promise.all(contexts.map((ctx) => ctx.dispose())); + process.exit(0); +}); diff --git a/packages/transformers-llguidance/src/index.ts b/packages/transformers-llguidance/src/index.ts new file mode 100644 index 000000000..a578cff5a --- /dev/null +++ b/packages/transformers-llguidance/src/index.ts @@ -0,0 +1,250 @@ +import { LogitsProcessor, LogitsProcessorList, StoppingCriteria, type Tensor } from '@huggingface/transformers'; +import { type LLGuidanceResponseFormat, loadBundledLLGuidance } from 'llguidance'; + +export type ResponseFormat = LLGuidanceResponseFormat; + +type GuidanceMask = Uint32Array | Uint8Array | boolean[] | number[]; + +type GuidanceMaskResult = + | { mask: GuidanceMask; vocabSize?: number; stop?: false } + | { stop: true } + | { backtrack?: number; ffTokens?: number[] }; + +type GuidanceCommitResult = { + stop?: boolean; + backtrack?: number; + ffTokens?: number[]; +}; + +type GuidanceInterpreter = { + computeMask(): GuidanceMaskResult; + commitToken(tokenId: number): GuidanceCommitResult | undefined; +}; + +type LlguidanceState = { + completed: boolean; + interpreter: GuidanceInterpreter; + step: number; +}; + +export class LlguidanceConstraint { + static async fromResponseFormat(tokenizer: unknown, response_format: ResponseFormat) { + console.log('[LlguidanceConstraint] loading llguidance', { + response_format, + }); + + const runtime = await loadBundledLLGuidance(); + const interpreter = runtime.createInterpreter({ + tokenizer, + response_format, + }) as GuidanceInterpreter; + + console.log('[LlguidanceConstraint] interpreter created'); + + const state: LlguidanceState = { + completed: false, + interpreter, + step: 0, + }; + + const logits_processor = new LogitsProcessorList(); + logits_processor.push(new LlguidanceLogitsProcessor(state)); + + return { + logits_processor, + stopping_criteria: new LlguidanceStoppingCriteria(state), + }; + } +} + +class LlguidanceLogitsProcessor extends LogitsProcessor { + private state: LlguidanceState; + + constructor(state: LlguidanceState) { + super(); + this.state = state; + } + + _call(_inputIds: bigint[][], logits: Tensor) { + if (this.state.completed) { + console.log('[LlguidanceLogitsProcessor] skip completed', { + step: this.state.step, + }); + return logits; + } + + this.state.step++; + const vocabSize = logits.dims.at(-1); + console.log('[LlguidanceLogitsProcessor] compute mask', { + step: this.state.step, + logitsDims: logits.dims, + vocabSize, + }); + + let result: GuidanceMaskResult; + try { + result = this.state.interpreter.computeMask(); + } catch (error) { + if (!String((error as Error).message).includes('compute_mask() called after stop')) { + throw error; + } + this.state.completed = true; + console.log('[LlguidanceLogitsProcessor] compute after stop', { + step: this.state.step, + }); + return logits; + } + + console.log('[LlguidanceLogitsProcessor] mask result', { + step: this.state.step, + result: summarizeMaskResult(result, vocabSize), + }); + + if ('stop' in result && result.stop) { + this.state.completed = true; + console.log('[LlguidanceLogitsProcessor] stopped by mask', { + step: this.state.step, + }); + return logits; + } + + if ('mask' in result) { + const applied = applyMask(logits, result.mask, result.vocabSize ?? vocabSize); + console.log('[LlguidanceLogitsProcessor] mask applied', { + step: this.state.step, + ...applied, + }); + } + + return logits; + } + + onTokenSampled(tokenId: number, batchIdx: number, inputIds: bigint[][]) { + console.log('[LlguidanceLogitsProcessor] token sampled', { + step: this.state.step, + tokenId, + batchIdx, + inputLength: inputIds[batchIdx]?.length, + completed: this.state.completed, + }); + + if (this.state.completed) return; + + const result = this.state.interpreter.commitToken(tokenId); + console.log('[LlguidanceLogitsProcessor] token committed', { + step: this.state.step, + tokenId, + result, + }); + + if (result?.stop) { + this.state.completed = true; + console.log('[LlguidanceLogitsProcessor] stopped by commit', { + step: this.state.step, + tokenId, + }); + } + } + + onTokensSampled(tokenIds: number[], inputIds: bigint[][]) { + console.log('[LlguidanceLogitsProcessor] tokens sampled', { + step: this.state.step, + tokenIds, + inputLengths: inputIds.map((ids) => ids.length), + completed: this.state.completed, + }); + + for (let batchIdx = 0; batchIdx < tokenIds.length; ++batchIdx) { + this.onTokenSampled(tokenIds[batchIdx], batchIdx, inputIds); + } + } +} + +class LlguidanceStoppingCriteria extends StoppingCriteria { + private state: LlguidanceState; + + constructor(state: LlguidanceState) { + super(); + this.state = state; + } + + _call(inputIds: ArrayLike[]) { + const result = new Array(inputIds.length).fill(this.state.completed); + console.log('[LlguidanceStoppingCriteria] call', { + step: this.state.step, + completed: this.state.completed, + result, + inputLengths: inputIds.map((ids) => ids.length), + }); + return result; + } +} + +function summarizeMaskResult(result: GuidanceMaskResult, vocabSize?: number) { + if ('stop' in result && result.stop) { + return { stop: true }; + } + + if (!('mask' in result)) { + return result; + } + + return { + maskLength: result.mask.length, + vocabSize: result.vocabSize ?? vocabSize, + allowed: countAllowed(result.mask, result.vocabSize ?? vocabSize), + sampleAllowedTokenIds: sampleAllowedTokenIds(result.mask, result.vocabSize ?? vocabSize), + }; +} + +function countAllowed(mask: GuidanceMask, vocabSize?: number) { + if (!vocabSize) return undefined; + + let allowed = 0; + for (let tokenId = 0; tokenId < vocabSize; ++tokenId) { + if (isAllowed(mask, tokenId, vocabSize)) allowed++; + } + return allowed; +} + +function sampleAllowedTokenIds(mask: GuidanceMask, vocabSize?: number) { + if (!vocabSize) return []; + + const tokenIds: number[] = []; + for (let tokenId = 0; tokenId < vocabSize && tokenIds.length < 25; ++tokenId) { + if (isAllowed(mask, tokenId, vocabSize)) tokenIds.push(tokenId); + } + return tokenIds; +} + +function isAllowed(mask: GuidanceMask, tokenId: number, vocabSize: number) { + if (mask.length >= vocabSize) { + return Boolean(mask[tokenId]); + } + return Boolean(Number(mask[tokenId >> 5]) & (1 << (tokenId & 31))); +} + +function applyMask(logits: Tensor, mask: GuidanceMask, vocabSize?: number) { + if (!vocabSize) { + return { vocabSize, batchSize: 0, masked: 0, allowed: undefined }; + } + + const data = logits.data as Float32Array | Float64Array | number[]; + const batchSize = Math.max(1, data.length / vocabSize); + let masked = 0; + let allowed = 0; + + for (let batch = 0; batch < batchSize; ++batch) { + const offset = batch * vocabSize; + for (let tokenId = 0; tokenId < vocabSize; ++tokenId) { + if (!isAllowed(mask, tokenId, vocabSize)) { + data[offset + tokenId] = -Infinity; + masked++; + } else if (batch === 0) { + allowed++; + } + } + } + + return { vocabSize, batchSize, masked, allowed }; +} diff --git a/packages/transformers-llguidance/tests/llguidance-constraint.test.js b/packages/transformers-llguidance/tests/llguidance-constraint.test.js new file mode 100644 index 000000000..53e5c462d --- /dev/null +++ b/packages/transformers-llguidance/tests/llguidance-constraint.test.js @@ -0,0 +1,52 @@ +import { jest } from "@jest/globals"; +import { Tensor } from "@huggingface/transformers"; + +const computeMask = jest.fn(); +const commitToken = jest.fn(); +const createInterpreter = jest.fn(() => ({ computeMask, commitToken })); +const loadBundledLLGuidance = jest.fn(async () => ({ createInterpreter })); + +jest.unstable_mockModule("llguidance", () => ({ + loadBundledLLGuidance, +})); + +const { LlguidanceConstraint } = await import("../dist/index.js"); + +describe("LlguidanceConstraint", () => { + beforeEach(() => { + computeMask.mockReset(); + commitToken.mockReset(); + createInterpreter.mockClear(); + loadBundledLLGuidance.mockClear(); + }); + + it("loads llguidance and applies masks", async () => { + computeMask.mockReturnValue({ mask: [true, false, true, false], vocabSize: 4 }); + + const tokenizer = { name: "tokenizer" }; + const response_format = { type: "json_schema", json_schema: { type: "object" } }; + const { logits_processor } = await LlguidanceConstraint.fromResponseFormat(tokenizer, response_format); + const logits = new Tensor("float32", new Float32Array([1, 2, 3, 4, 5, 6, 7, 8]), [2, 4]); + + logits_processor([[0n], [0n]], logits); + + expect(loadBundledLLGuidance).toHaveBeenCalledTimes(1); + expect(createInterpreter).toHaveBeenCalledWith({ tokenizer, response_format }); + expect(Array.from(logits.data)).toEqual([1, -Infinity, 3, -Infinity, 5, -Infinity, 7, -Infinity]); + }); + + it("commits sampled tokens and stops when llguidance stops", async () => { + computeMask.mockReturnValue({ mask: [true, true], vocabSize: 2 }); + commitToken.mockReturnValueOnce(undefined).mockReturnValueOnce({ stop: true }); + + const { logits_processor, stopping_criteria } = await LlguidanceConstraint.fromResponseFormat({}, { type: "json_object" }); + + expect(stopping_criteria([[0n]])).toEqual([false]); + + logits_processor.onTokensSampled([0, 1], [[0n, 0n], [0n, 1n]]); + + expect(commitToken).toHaveBeenCalledWith(0); + expect(commitToken).toHaveBeenCalledWith(1); + expect(stopping_criteria([[0n, 0n], [0n, 1n]])).toEqual([true, true]); + }); +}); diff --git a/packages/transformers-llguidance/tsconfig.json b/packages/transformers-llguidance/tsconfig.json new file mode 100644 index 000000000..0a88cabdd --- /dev/null +++ b/packages/transformers-llguidance/tsconfig.json @@ -0,0 +1,21 @@ +{ + "include": ["src/**/*"], + "compilerOptions": { + "target": "esnext", + "module": "esnext", + "moduleResolution": "bundler", + "outDir": "types", + "rootDir": "src", + "strict": true, + "skipLibCheck": true, + "declaration": true, + "declarationMap": true, + "noEmit": false, + "emitDeclarationOnly": true, + "esModuleInterop": true, + "composite": true + }, + "typeAcquisition": { + "include": ["jest"] + } +} diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 3ed937a04..b0585df63 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -58,6 +58,31 @@ importers: specifier: 5.9.3 version: 5.9.3 + packages/transformers-llguidance: + dependencies: + llguidance: + specifier: ^0.1.7 + version: 0.1.7 + devDependencies: + '@huggingface/transformers': + specifier: workspace:* + version: link:../transformers + '@types/jest': + specifier: ^30.0.0 + version: 30.0.0 + '@types/node': + specifier: ^24.1.0 + version: 24.10.9 + esbuild: + specifier: ^0.27.2 + version: 0.27.2 + jest: + specifier: ^30.2.0 + version: 30.2.0(@types/node@24.10.9) + typescript: + specifier: 5.9.3 + version: 5.9.3 + packages: '@babel/code-frame@7.28.6': @@ -427,105 +452,89 @@ 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==} @@ -816,49 +825,41 @@ 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==} @@ -1584,6 +1585,9 @@ packages: linkify-it@5.0.0: resolution: {integrity: sha512-5aHCbzQRADcdP+ATqnDuhhJ/MRIqDkZX5pyjFHRRysS8vZ5AbqGEoFIb6pYHPZ+L/OC2Lc+xT8uHVVR5CAK/wQ==} + llguidance@0.1.7: + resolution: {integrity: sha512-r55h3hDmkq313cFlqbKyx4IebqpreyrrnMJip7IR+ocwPGm7O6ZoDVCIAEg2wd6k7pQ5wWyuSzV9rSefqKVs6Q==} + locate-path@5.0.0: resolution: {integrity: sha512-t7hw9pI+WvuwNJXwk5zVHpyhIqzg2qTlklJOf0mVxGSbe3Fp2VieZcduNYjaLDoy6p9uGpQEGWG87WpMKlNq8g==} engines: {node: '>=8'} @@ -3746,6 +3750,8 @@ snapshots: dependencies: uc.micro: 2.1.0 + llguidance@0.1.7: {} + locate-path@5.0.0: dependencies: p-locate: 4.1.0 From 8b7dc30b1111b6f6dc639eeb03adac86fc4cb628 Mon Sep 17 00:00:00 2001 From: Nico Martin Date: Fri, 10 Jul 2026 12:03:01 +0200 Subject: [PATCH 06/17] added LlguidanceConstraint --- .../transformers-llguidance/scripts/dev.mjs | 31 +++++ packages/transformers-llguidance/src/index.ts | 126 +++++++++++++++--- .../tests/llguidance-constraint.test.js | 6 + packages/transformers/src/transformers.js | 2 + 4 files changed, 148 insertions(+), 17 deletions(-) diff --git a/packages/transformers-llguidance/scripts/dev.mjs b/packages/transformers-llguidance/scripts/dev.mjs index 4b75dfcad..4c5186e26 100644 --- a/packages/transformers-llguidance/scripts/dev.mjs +++ b/packages/transformers-llguidance/scripts/dev.mjs @@ -1,7 +1,30 @@ import { context } from "esbuild"; +import { spawn } from "node:child_process"; import { rmSync } from "node:fs"; rmSync("dist", { recursive: true, force: true }); +rmSync("types", { recursive: true, force: true }); + +const watchLogger = { + name: "watch-logger", + setup(build) { + let startTime = 0; + + build.onStart(() => { + startTime = performance.now(); + console.log(`[transformers-llguidance] rebuilding ${build.initialOptions.outfile}...`); + }); + + build.onEnd((result) => { + const duration = (performance.now() - startTime).toFixed(2); + if (result.errors.length > 0) { + console.log(`[transformers-llguidance] rebuild failed in ${duration}ms`); + } else { + console.log(`[transformers-llguidance] rebuilt ${build.initialOptions.outfile} in ${duration}ms`); + } + }); + }, +}; const common = { entryPoints: ["src/index.ts"], @@ -10,6 +33,7 @@ const common = { target: "es2022", sourcemap: true, external: ["@huggingface/transformers", "llguidance"], + plugins: [watchLogger], }; const contexts = await Promise.all([ @@ -26,9 +50,16 @@ const contexts = await Promise.all([ ]); await Promise.all(contexts.map((ctx) => ctx.watch())); + +const tscWatch = spawn("tsc", ["--build", "--watch", "--preserveWatchOutput"], { + stdio: "inherit", + shell: true, +}); + console.log("Watching @huggingface/transformers-llguidance..."); process.on("SIGINT", async () => { + tscWatch.kill(); await Promise.all(contexts.map((ctx) => ctx.dispose())); process.exit(0); }); diff --git a/packages/transformers-llguidance/src/index.ts b/packages/transformers-llguidance/src/index.ts index a578cff5a..7578d85dc 100644 --- a/packages/transformers-llguidance/src/index.ts +++ b/packages/transformers-llguidance/src/index.ts @@ -1,8 +1,27 @@ -import { LogitsProcessor, LogitsProcessorList, StoppingCriteria, type Tensor } from '@huggingface/transformers'; -import { type LLGuidanceResponseFormat, loadBundledLLGuidance } from 'llguidance'; +import { + LogitsProcessor, + LogitsProcessorList, + StoppingCriteria, + env, + loadWasmBinary, + loadWasmFactory, + logger, + type Tensor, +} from '@huggingface/transformers'; +import { type LLGuidanceResponseFormat, type LoadBundledLLGuidanceOptions, loadBundledLLGuidance } from 'llguidance'; export type ResponseFormat = LLGuidanceResponseFormat; +export type LlguidanceLoadOptions = LoadBundledLLGuidanceOptions & { + /** Whether to pre-load and cache llguidance WASM assets. Defaults to env.useWasmCache. */ + useWasmCache?: boolean; +}; + +const LLGUIDANCE_VERSION = '0.1.7'; +const LLGUIDANCE_WASM_BASE = `https://cdn.jsdelivr.net/npm/llguidance@${LLGUIDANCE_VERSION}/wasm/`; +const DEFAULT_LLGUIDANCE_WASM_URL = `${LLGUIDANCE_WASM_BASE}llguidance_wasm_bg.wasm`; +const DEFAULT_LLGUIDANCE_WASM_FACTORY_URL = `${LLGUIDANCE_WASM_BASE}llguidance_wasm.js`; + type GuidanceMask = Uint32Array | Uint8Array | boolean[] | number[]; type GuidanceMaskResult = @@ -28,18 +47,22 @@ type LlguidanceState = { }; export class LlguidanceConstraint { - static async fromResponseFormat(tokenizer: unknown, response_format: ResponseFormat) { - console.log('[LlguidanceConstraint] loading llguidance', { + static async fromResponseFormat( + tokenizer: unknown, + response_format: ResponseFormat, + loadOptions: LlguidanceLoadOptions = {}, + ) { + logger.debug('[LlguidanceConstraint] loading llguidance', { response_format, }); - const runtime = await loadBundledLLGuidance(); + const runtime = await loadLLGuidance(loadOptions); const interpreter = runtime.createInterpreter({ tokenizer, response_format, }) as GuidanceInterpreter; - console.log('[LlguidanceConstraint] interpreter created'); + logger.debug('[LlguidanceConstraint] interpreter created'); const state: LlguidanceState = { completed: false, @@ -57,6 +80,75 @@ export class LlguidanceConstraint { } } +async function loadLLGuidance(loadOptions: LlguidanceLoadOptions) { + const { useWasmCache = env.useWasmCache, ...options } = loadOptions; + if (!useWasmCache || isNodeLikeRuntime() || options.wasmFactory) { + return loadBundledLLGuidance(options); + } + + const wasmSource = options.wasm ?? options.wasmUrl ?? DEFAULT_LLGUIDANCE_WASM_URL; + const wasmFactorySource = options.wasmFactoryUrl ?? DEFAULT_LLGUIDANCE_WASM_FACTORY_URL; + const cachedOptions = { ...options }; + + const [wasm, wasmFactoryUrl] = await Promise.all([ + loadCacheableWasm(wasmSource), + loadCacheableWasmFactory(wasmFactorySource), + ]); + + if (wasm) { + cachedOptions.wasm = wasm; + delete cachedOptions.wasmUrl; + } + + if (wasm && wasmFactoryUrl) { + cachedOptions.wasmFactoryUrl = wasmFactoryUrl; + } + + return loadBundledLLGuidance(cachedOptions); +} + +async function loadCacheableWasm(source: LoadBundledLLGuidanceOptions['wasm']) { + const url = toCacheableURL(source); + if (!url) return null; + + try { + return await loadWasmBinary(url); + } catch (error) { + logger.warn('Failed to pre-load llguidance WASM binary:', error); + return null; + } +} + +async function loadCacheableWasmFactory(source: LoadBundledLLGuidanceOptions['wasmFactoryUrl']) { + const url = toCacheableURL(source); + if (!url) return null; + + try { + return await loadWasmFactory(url); + } catch (error) { + logger.warn('Failed to pre-load llguidance WASM factory:', error); + return null; + } +} + +function toCacheableURL(source: unknown) { + if (typeof source === 'string') { + return isBlobURL(source) ? null : new URL(source, globalThis.location?.href).href; + } + if (source instanceof URL) { + return isBlobURL(source.href) ? null : source.href; + } + return null; +} + +function isBlobURL(url: string) { + return url.startsWith('blob:'); +} + +function isNodeLikeRuntime() { + return typeof process !== 'undefined' && Boolean(process.versions?.node); +} + class LlguidanceLogitsProcessor extends LogitsProcessor { private state: LlguidanceState; @@ -67,7 +159,7 @@ class LlguidanceLogitsProcessor extends LogitsProcessor { _call(_inputIds: bigint[][], logits: Tensor) { if (this.state.completed) { - console.log('[LlguidanceLogitsProcessor] skip completed', { + logger.debug('[LlguidanceLogitsProcessor] skip completed', { step: this.state.step, }); return logits; @@ -75,7 +167,7 @@ class LlguidanceLogitsProcessor extends LogitsProcessor { this.state.step++; const vocabSize = logits.dims.at(-1); - console.log('[LlguidanceLogitsProcessor] compute mask', { + logger.debug('[LlguidanceLogitsProcessor] compute mask', { step: this.state.step, logitsDims: logits.dims, vocabSize, @@ -89,20 +181,20 @@ class LlguidanceLogitsProcessor extends LogitsProcessor { throw error; } this.state.completed = true; - console.log('[LlguidanceLogitsProcessor] compute after stop', { + logger.debug('[LlguidanceLogitsProcessor] compute after stop', { step: this.state.step, }); return logits; } - console.log('[LlguidanceLogitsProcessor] mask result', { + logger.debug('[LlguidanceLogitsProcessor] mask result', { step: this.state.step, result: summarizeMaskResult(result, vocabSize), }); if ('stop' in result && result.stop) { this.state.completed = true; - console.log('[LlguidanceLogitsProcessor] stopped by mask', { + logger.debug('[LlguidanceLogitsProcessor] stopped by mask', { step: this.state.step, }); return logits; @@ -110,7 +202,7 @@ class LlguidanceLogitsProcessor extends LogitsProcessor { if ('mask' in result) { const applied = applyMask(logits, result.mask, result.vocabSize ?? vocabSize); - console.log('[LlguidanceLogitsProcessor] mask applied', { + logger.debug('[LlguidanceLogitsProcessor] mask applied', { step: this.state.step, ...applied, }); @@ -120,7 +212,7 @@ class LlguidanceLogitsProcessor extends LogitsProcessor { } onTokenSampled(tokenId: number, batchIdx: number, inputIds: bigint[][]) { - console.log('[LlguidanceLogitsProcessor] token sampled', { + logger.debug('[LlguidanceLogitsProcessor] token sampled', { step: this.state.step, tokenId, batchIdx, @@ -131,7 +223,7 @@ class LlguidanceLogitsProcessor extends LogitsProcessor { if (this.state.completed) return; const result = this.state.interpreter.commitToken(tokenId); - console.log('[LlguidanceLogitsProcessor] token committed', { + logger.debug('[LlguidanceLogitsProcessor] token committed', { step: this.state.step, tokenId, result, @@ -139,7 +231,7 @@ class LlguidanceLogitsProcessor extends LogitsProcessor { if (result?.stop) { this.state.completed = true; - console.log('[LlguidanceLogitsProcessor] stopped by commit', { + logger.debug('[LlguidanceLogitsProcessor] stopped by commit', { step: this.state.step, tokenId, }); @@ -147,7 +239,7 @@ class LlguidanceLogitsProcessor extends LogitsProcessor { } onTokensSampled(tokenIds: number[], inputIds: bigint[][]) { - console.log('[LlguidanceLogitsProcessor] tokens sampled', { + logger.debug('[LlguidanceLogitsProcessor] tokens sampled', { step: this.state.step, tokenIds, inputLengths: inputIds.map((ids) => ids.length), @@ -170,7 +262,7 @@ class LlguidanceStoppingCriteria extends StoppingCriteria { _call(inputIds: ArrayLike[]) { const result = new Array(inputIds.length).fill(this.state.completed); - console.log('[LlguidanceStoppingCriteria] call', { + logger.debug('[LlguidanceStoppingCriteria] call', { step: this.state.step, completed: this.state.completed, result, diff --git a/packages/transformers-llguidance/tests/llguidance-constraint.test.js b/packages/transformers-llguidance/tests/llguidance-constraint.test.js index 53e5c462d..c9f32a42d 100644 --- a/packages/transformers-llguidance/tests/llguidance-constraint.test.js +++ b/packages/transformers-llguidance/tests/llguidance-constraint.test.js @@ -35,6 +35,12 @@ describe("LlguidanceConstraint", () => { expect(Array.from(logits.data)).toEqual([1, -Infinity, 3, -Infinity, 5, -Infinity, 7, -Infinity]); }); + it("passes explicit llguidance load options", async () => { + await LlguidanceConstraint.fromResponseFormat({}, { type: "json_object" }, { useWasmCache: false, wasmUrl: "custom.wasm" }); + + expect(loadBundledLLGuidance).toHaveBeenCalledWith({ wasmUrl: "custom.wasm" }); + }); + it("commits sampled tokens and stops when llguidance stops", async () => { computeMask.mockReturnValue({ mask: [true, true], vocabSize: 2 }); commitToken.mockReturnValueOnce(undefined).mockReturnValueOnce({ stop: true }); diff --git a/packages/transformers/src/transformers.js b/packages/transformers/src/transformers.js index ef01569ef..7d6168020 100644 --- a/packages/transformers/src/transformers.js +++ b/packages/transformers/src/transformers.js @@ -52,6 +52,8 @@ export { load_video, RawVideo, RawVideoFrame } from './utils/video.js'; export * from './utils/tensor.js'; export { softmax, log_softmax, dot, cos_sim } from './utils/maths.js'; export { random } from './utils/random.js'; +export { logger } from './utils/logger.js'; +export { loadWasmBinary, loadWasmFactory } from './backends/utils/cacheWasm.js'; export { DynamicCache } from './cache_utils.js'; From 097dd6555ce66a28ff3ea46e05e0d59c41c87522 Mon Sep 17 00:00:00 2001 From: Nico Martin Date: Fri, 10 Jul 2026 12:09:25 +0200 Subject: [PATCH 07/17] added regex example --- packages/transformers-llguidance/README.md | 20 ++++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/packages/transformers-llguidance/README.md b/packages/transformers-llguidance/README.md index f1776465c..e167b245f 100644 --- a/packages/transformers-llguidance/README.md +++ b/packages/transformers-llguidance/README.md @@ -26,3 +26,23 @@ await model.generate({ stopping_criteria, }); ``` + +## Regex constraints + +Use `type: "regex"` to constrain generation to a regular expression. For example, this only allows ISO-like dates in `YYYY-MM-DD` format: + +```js +import { LlguidanceConstraint } from "@huggingface/transformers-llguidance"; + +const { logits_processor, stopping_criteria } = + await LlguidanceConstraint.fromResponseFormat(tokenizer, { + type: "regex", + regex: "\\d{4}-\\d{2}-\\d{2}", + }); + +const output = await model.generate({ + ...inputs, + logits_processor, + stopping_criteria, +}); +``` From e695be3707e12e12c730faf816f6547d1b8bdd51 Mon Sep 17 00:00:00 2001 From: Nico Martin Date: Sat, 11 Jul 2026 08:22:41 +0200 Subject: [PATCH 08/17] clean up and added more unit tests --- .../src/LlguidanceConstraint.ts | 166 +++++++++ .../transformers-llguidance/src/constants.ts | 4 + packages/transformers-llguidance/src/index.ts | 344 +----------------- .../transformers-llguidance/src/utils/mask.ts | 72 ++++ .../src/utils/runtime.ts | 3 + .../src/utils/types.ts | 23 ++ .../transformers-llguidance/src/utils/wasm.ts | 75 ++++ .../tests/llguidance-constraint.test.js | 73 +++- .../tests/wasm-loading.test.js | 123 +++++++ 9 files changed, 539 insertions(+), 344 deletions(-) create mode 100644 packages/transformers-llguidance/src/LlguidanceConstraint.ts create mode 100644 packages/transformers-llguidance/src/constants.ts create mode 100644 packages/transformers-llguidance/src/utils/mask.ts create mode 100644 packages/transformers-llguidance/src/utils/runtime.ts create mode 100644 packages/transformers-llguidance/src/utils/types.ts create mode 100644 packages/transformers-llguidance/src/utils/wasm.ts create mode 100644 packages/transformers-llguidance/tests/wasm-loading.test.js diff --git a/packages/transformers-llguidance/src/LlguidanceConstraint.ts b/packages/transformers-llguidance/src/LlguidanceConstraint.ts new file mode 100644 index 000000000..bee7b774e --- /dev/null +++ b/packages/transformers-llguidance/src/LlguidanceConstraint.ts @@ -0,0 +1,166 @@ +import { LogitsProcessor, LogitsProcessorList, StoppingCriteria, logger, type Tensor } from '@huggingface/transformers'; +import { type LLGuidanceResponseFormat } from 'llguidance'; + +import { applyMask, summarizeMaskResult } from './utils/mask'; +import { type GuidanceInterpreter, type GuidanceMaskResult, type LlguidanceState } from './utils/types'; +import { type LlguidanceLoadOptions, loadLLGuidance } from './utils/wasm'; + +export type ResponseFormat = LLGuidanceResponseFormat; +export type { LlguidanceLoadOptions }; + +export class LlguidanceConstraint { + static async fromResponseFormat( + tokenizer: unknown, + response_format: ResponseFormat, + loadOptions: LlguidanceLoadOptions = {}, + ) { + logger.debug('[LlguidanceConstraint] loading llguidance', { + response_format, + }); + + const runtime = await loadLLGuidance(loadOptions); + const interpreter = runtime.createInterpreter({ + tokenizer, + response_format, + }) as GuidanceInterpreter; + + logger.debug('[LlguidanceConstraint] interpreter created'); + + const state: LlguidanceState = { + completed: false, + interpreter, + step: 0, + }; + + const logits_processor = new LogitsProcessorList(); + logits_processor.push(new LlguidanceLogitsProcessor(state)); + + return { + logits_processor, + stopping_criteria: new LlguidanceStoppingCriteria(state), + }; + } +} + +class LlguidanceLogitsProcessor extends LogitsProcessor { + private state: LlguidanceState; + + constructor(state: LlguidanceState) { + super(); + this.state = state; + } + + _call(_inputIds: bigint[][], logits: Tensor) { + if (this.state.completed) { + logger.debug('[LlguidanceLogitsProcessor] skip completed', { + step: this.state.step, + }); + return logits; + } + + this.state.step++; + const vocabSize = logits.dims.at(-1); + logger.debug('[LlguidanceLogitsProcessor] compute mask', { + step: this.state.step, + logitsDims: logits.dims, + vocabSize, + }); + + let result: GuidanceMaskResult; + try { + result = this.state.interpreter.computeMask(); + } catch (error) { + if (!String((error as Error).message).includes('compute_mask() called after stop')) { + throw error; + } + this.state.completed = true; + logger.debug('[LlguidanceLogitsProcessor] compute after stop', { + step: this.state.step, + }); + return logits; + } + + logger.debug('[LlguidanceLogitsProcessor] mask result', { + step: this.state.step, + result: summarizeMaskResult(result, vocabSize), + }); + + if ('stop' in result && result.stop) { + this.state.completed = true; + logger.debug('[LlguidanceLogitsProcessor] stopped by mask', { + step: this.state.step, + }); + return logits; + } + + if ('mask' in result) { + const applied = applyMask(logits, result.mask, result.vocabSize ?? vocabSize); + logger.debug('[LlguidanceLogitsProcessor] mask applied', { + step: this.state.step, + ...applied, + }); + } + + return logits; + } + + onTokenSampled(tokenId: number, batchIdx: number, inputIds: bigint[][]) { + logger.debug('[LlguidanceLogitsProcessor] token sampled', { + step: this.state.step, + tokenId, + batchIdx, + inputLength: inputIds[batchIdx]?.length, + completed: this.state.completed, + }); + + if (this.state.completed) return; + + const result = this.state.interpreter.commitToken(tokenId); + logger.debug('[LlguidanceLogitsProcessor] token committed', { + step: this.state.step, + tokenId, + result, + }); + + if (result?.stop) { + this.state.completed = true; + logger.debug('[LlguidanceLogitsProcessor] stopped by commit', { + step: this.state.step, + tokenId, + }); + } + } + + onTokensSampled(tokenIds: number[], inputIds: bigint[][]) { + logger.debug('[LlguidanceLogitsProcessor] tokens sampled', { + step: this.state.step, + tokenIds, + inputLengths: inputIds.map((ids) => ids.length), + completed: this.state.completed, + }); + + for (let batchIdx = 0; batchIdx < tokenIds.length; ++batchIdx) { + this.onTokenSampled(tokenIds[batchIdx], batchIdx, inputIds); + } + } +} + +class LlguidanceStoppingCriteria extends StoppingCriteria { + private state: LlguidanceState; + + constructor(state: LlguidanceState) { + super(); + this.state = state; + } + + _call(inputIds: ArrayLike[]) { + const result = new Array(inputIds.length).fill(this.state.completed); + logger.debug('[LlguidanceStoppingCriteria] call', { + step: this.state.step, + completed: this.state.completed, + result, + inputLengths: inputIds.map((ids) => ids.length), + }); + return result; + } +} diff --git a/packages/transformers-llguidance/src/constants.ts b/packages/transformers-llguidance/src/constants.ts new file mode 100644 index 000000000..773108fac --- /dev/null +++ b/packages/transformers-llguidance/src/constants.ts @@ -0,0 +1,4 @@ +export const LLGUIDANCE_VERSION = '0.1.7'; +export const LLGUIDANCE_WASM_BASE = `https://cdn.jsdelivr.net/npm/llguidance@${LLGUIDANCE_VERSION}/wasm/`; +export const DEFAULT_LLGUIDANCE_WASM_URL = `${LLGUIDANCE_WASM_BASE}llguidance_wasm_bg.wasm`; +export const DEFAULT_LLGUIDANCE_WASM_FACTORY_URL = `${LLGUIDANCE_WASM_BASE}llguidance_wasm.js`; diff --git a/packages/transformers-llguidance/src/index.ts b/packages/transformers-llguidance/src/index.ts index 7578d85dc..a870a56f9 100644 --- a/packages/transformers-llguidance/src/index.ts +++ b/packages/transformers-llguidance/src/index.ts @@ -1,342 +1,2 @@ -import { - LogitsProcessor, - LogitsProcessorList, - StoppingCriteria, - env, - loadWasmBinary, - loadWasmFactory, - logger, - type Tensor, -} from '@huggingface/transformers'; -import { type LLGuidanceResponseFormat, type LoadBundledLLGuidanceOptions, loadBundledLLGuidance } from 'llguidance'; - -export type ResponseFormat = LLGuidanceResponseFormat; - -export type LlguidanceLoadOptions = LoadBundledLLGuidanceOptions & { - /** Whether to pre-load and cache llguidance WASM assets. Defaults to env.useWasmCache. */ - useWasmCache?: boolean; -}; - -const LLGUIDANCE_VERSION = '0.1.7'; -const LLGUIDANCE_WASM_BASE = `https://cdn.jsdelivr.net/npm/llguidance@${LLGUIDANCE_VERSION}/wasm/`; -const DEFAULT_LLGUIDANCE_WASM_URL = `${LLGUIDANCE_WASM_BASE}llguidance_wasm_bg.wasm`; -const DEFAULT_LLGUIDANCE_WASM_FACTORY_URL = `${LLGUIDANCE_WASM_BASE}llguidance_wasm.js`; - -type GuidanceMask = Uint32Array | Uint8Array | boolean[] | number[]; - -type GuidanceMaskResult = - | { mask: GuidanceMask; vocabSize?: number; stop?: false } - | { stop: true } - | { backtrack?: number; ffTokens?: number[] }; - -type GuidanceCommitResult = { - stop?: boolean; - backtrack?: number; - ffTokens?: number[]; -}; - -type GuidanceInterpreter = { - computeMask(): GuidanceMaskResult; - commitToken(tokenId: number): GuidanceCommitResult | undefined; -}; - -type LlguidanceState = { - completed: boolean; - interpreter: GuidanceInterpreter; - step: number; -}; - -export class LlguidanceConstraint { - static async fromResponseFormat( - tokenizer: unknown, - response_format: ResponseFormat, - loadOptions: LlguidanceLoadOptions = {}, - ) { - logger.debug('[LlguidanceConstraint] loading llguidance', { - response_format, - }); - - const runtime = await loadLLGuidance(loadOptions); - const interpreter = runtime.createInterpreter({ - tokenizer, - response_format, - }) as GuidanceInterpreter; - - logger.debug('[LlguidanceConstraint] interpreter created'); - - const state: LlguidanceState = { - completed: false, - interpreter, - step: 0, - }; - - const logits_processor = new LogitsProcessorList(); - logits_processor.push(new LlguidanceLogitsProcessor(state)); - - return { - logits_processor, - stopping_criteria: new LlguidanceStoppingCriteria(state), - }; - } -} - -async function loadLLGuidance(loadOptions: LlguidanceLoadOptions) { - const { useWasmCache = env.useWasmCache, ...options } = loadOptions; - if (!useWasmCache || isNodeLikeRuntime() || options.wasmFactory) { - return loadBundledLLGuidance(options); - } - - const wasmSource = options.wasm ?? options.wasmUrl ?? DEFAULT_LLGUIDANCE_WASM_URL; - const wasmFactorySource = options.wasmFactoryUrl ?? DEFAULT_LLGUIDANCE_WASM_FACTORY_URL; - const cachedOptions = { ...options }; - - const [wasm, wasmFactoryUrl] = await Promise.all([ - loadCacheableWasm(wasmSource), - loadCacheableWasmFactory(wasmFactorySource), - ]); - - if (wasm) { - cachedOptions.wasm = wasm; - delete cachedOptions.wasmUrl; - } - - if (wasm && wasmFactoryUrl) { - cachedOptions.wasmFactoryUrl = wasmFactoryUrl; - } - - return loadBundledLLGuidance(cachedOptions); -} - -async function loadCacheableWasm(source: LoadBundledLLGuidanceOptions['wasm']) { - const url = toCacheableURL(source); - if (!url) return null; - - try { - return await loadWasmBinary(url); - } catch (error) { - logger.warn('Failed to pre-load llguidance WASM binary:', error); - return null; - } -} - -async function loadCacheableWasmFactory(source: LoadBundledLLGuidanceOptions['wasmFactoryUrl']) { - const url = toCacheableURL(source); - if (!url) return null; - - try { - return await loadWasmFactory(url); - } catch (error) { - logger.warn('Failed to pre-load llguidance WASM factory:', error); - return null; - } -} - -function toCacheableURL(source: unknown) { - if (typeof source === 'string') { - return isBlobURL(source) ? null : new URL(source, globalThis.location?.href).href; - } - if (source instanceof URL) { - return isBlobURL(source.href) ? null : source.href; - } - return null; -} - -function isBlobURL(url: string) { - return url.startsWith('blob:'); -} - -function isNodeLikeRuntime() { - return typeof process !== 'undefined' && Boolean(process.versions?.node); -} - -class LlguidanceLogitsProcessor extends LogitsProcessor { - private state: LlguidanceState; - - constructor(state: LlguidanceState) { - super(); - this.state = state; - } - - _call(_inputIds: bigint[][], logits: Tensor) { - if (this.state.completed) { - logger.debug('[LlguidanceLogitsProcessor] skip completed', { - step: this.state.step, - }); - return logits; - } - - this.state.step++; - const vocabSize = logits.dims.at(-1); - logger.debug('[LlguidanceLogitsProcessor] compute mask', { - step: this.state.step, - logitsDims: logits.dims, - vocabSize, - }); - - let result: GuidanceMaskResult; - try { - result = this.state.interpreter.computeMask(); - } catch (error) { - if (!String((error as Error).message).includes('compute_mask() called after stop')) { - throw error; - } - this.state.completed = true; - logger.debug('[LlguidanceLogitsProcessor] compute after stop', { - step: this.state.step, - }); - return logits; - } - - logger.debug('[LlguidanceLogitsProcessor] mask result', { - step: this.state.step, - result: summarizeMaskResult(result, vocabSize), - }); - - if ('stop' in result && result.stop) { - this.state.completed = true; - logger.debug('[LlguidanceLogitsProcessor] stopped by mask', { - step: this.state.step, - }); - return logits; - } - - if ('mask' in result) { - const applied = applyMask(logits, result.mask, result.vocabSize ?? vocabSize); - logger.debug('[LlguidanceLogitsProcessor] mask applied', { - step: this.state.step, - ...applied, - }); - } - - return logits; - } - - onTokenSampled(tokenId: number, batchIdx: number, inputIds: bigint[][]) { - logger.debug('[LlguidanceLogitsProcessor] token sampled', { - step: this.state.step, - tokenId, - batchIdx, - inputLength: inputIds[batchIdx]?.length, - completed: this.state.completed, - }); - - if (this.state.completed) return; - - const result = this.state.interpreter.commitToken(tokenId); - logger.debug('[LlguidanceLogitsProcessor] token committed', { - step: this.state.step, - tokenId, - result, - }); - - if (result?.stop) { - this.state.completed = true; - logger.debug('[LlguidanceLogitsProcessor] stopped by commit', { - step: this.state.step, - tokenId, - }); - } - } - - onTokensSampled(tokenIds: number[], inputIds: bigint[][]) { - logger.debug('[LlguidanceLogitsProcessor] tokens sampled', { - step: this.state.step, - tokenIds, - inputLengths: inputIds.map((ids) => ids.length), - completed: this.state.completed, - }); - - for (let batchIdx = 0; batchIdx < tokenIds.length; ++batchIdx) { - this.onTokenSampled(tokenIds[batchIdx], batchIdx, inputIds); - } - } -} - -class LlguidanceStoppingCriteria extends StoppingCriteria { - private state: LlguidanceState; - - constructor(state: LlguidanceState) { - super(); - this.state = state; - } - - _call(inputIds: ArrayLike[]) { - const result = new Array(inputIds.length).fill(this.state.completed); - logger.debug('[LlguidanceStoppingCriteria] call', { - step: this.state.step, - completed: this.state.completed, - result, - inputLengths: inputIds.map((ids) => ids.length), - }); - return result; - } -} - -function summarizeMaskResult(result: GuidanceMaskResult, vocabSize?: number) { - if ('stop' in result && result.stop) { - return { stop: true }; - } - - if (!('mask' in result)) { - return result; - } - - return { - maskLength: result.mask.length, - vocabSize: result.vocabSize ?? vocabSize, - allowed: countAllowed(result.mask, result.vocabSize ?? vocabSize), - sampleAllowedTokenIds: sampleAllowedTokenIds(result.mask, result.vocabSize ?? vocabSize), - }; -} - -function countAllowed(mask: GuidanceMask, vocabSize?: number) { - if (!vocabSize) return undefined; - - let allowed = 0; - for (let tokenId = 0; tokenId < vocabSize; ++tokenId) { - if (isAllowed(mask, tokenId, vocabSize)) allowed++; - } - return allowed; -} - -function sampleAllowedTokenIds(mask: GuidanceMask, vocabSize?: number) { - if (!vocabSize) return []; - - const tokenIds: number[] = []; - for (let tokenId = 0; tokenId < vocabSize && tokenIds.length < 25; ++tokenId) { - if (isAllowed(mask, tokenId, vocabSize)) tokenIds.push(tokenId); - } - return tokenIds; -} - -function isAllowed(mask: GuidanceMask, tokenId: number, vocabSize: number) { - if (mask.length >= vocabSize) { - return Boolean(mask[tokenId]); - } - return Boolean(Number(mask[tokenId >> 5]) & (1 << (tokenId & 31))); -} - -function applyMask(logits: Tensor, mask: GuidanceMask, vocabSize?: number) { - if (!vocabSize) { - return { vocabSize, batchSize: 0, masked: 0, allowed: undefined }; - } - - const data = logits.data as Float32Array | Float64Array | number[]; - const batchSize = Math.max(1, data.length / vocabSize); - let masked = 0; - let allowed = 0; - - for (let batch = 0; batch < batchSize; ++batch) { - const offset = batch * vocabSize; - for (let tokenId = 0; tokenId < vocabSize; ++tokenId) { - if (!isAllowed(mask, tokenId, vocabSize)) { - data[offset + tokenId] = -Infinity; - masked++; - } else if (batch === 0) { - allowed++; - } - } - } - - return { vocabSize, batchSize, masked, allowed }; -} +export { LlguidanceConstraint } from './LlguidanceConstraint'; +export type { LlguidanceLoadOptions, ResponseFormat } from './LlguidanceConstraint'; diff --git a/packages/transformers-llguidance/src/utils/mask.ts b/packages/transformers-llguidance/src/utils/mask.ts new file mode 100644 index 000000000..f1ee29ba2 --- /dev/null +++ b/packages/transformers-llguidance/src/utils/mask.ts @@ -0,0 +1,72 @@ +import { type Tensor } from '@huggingface/transformers'; + +import { type GuidanceMask, type GuidanceMaskResult } from './types'; + +export function summarizeMaskResult(result: GuidanceMaskResult, vocabSize?: number) { + if ('stop' in result && result.stop) { + return { stop: true }; + } + + if (!('mask' in result)) { + return result; + } + + return { + maskLength: result.mask.length, + vocabSize: result.vocabSize ?? vocabSize, + allowed: countAllowed(result.mask, result.vocabSize ?? vocabSize), + sampleAllowedTokenIds: sampleAllowedTokenIds(result.mask, result.vocabSize ?? vocabSize), + }; +} + +export function applyMask(logits: Tensor, mask: GuidanceMask, vocabSize?: number) { + if (!vocabSize) { + return { vocabSize, batchSize: 0, masked: 0, allowed: undefined }; + } + + const data = logits.data as Float32Array | Float64Array | number[]; + const batchSize = Math.max(1, data.length / vocabSize); + let masked = 0; + let allowed = 0; + + for (let batch = 0; batch < batchSize; ++batch) { + const offset = batch * vocabSize; + for (let tokenId = 0; tokenId < vocabSize; ++tokenId) { + if (!isAllowed(mask, tokenId, vocabSize)) { + data[offset + tokenId] = -Infinity; + masked++; + } else if (batch === 0) { + allowed++; + } + } + } + + return { vocabSize, batchSize, masked, allowed }; +} + +function countAllowed(mask: GuidanceMask, vocabSize?: number) { + if (!vocabSize) return undefined; + + let allowed = 0; + for (let tokenId = 0; tokenId < vocabSize; ++tokenId) { + if (isAllowed(mask, tokenId, vocabSize)) allowed++; + } + return allowed; +} + +function sampleAllowedTokenIds(mask: GuidanceMask, vocabSize?: number) { + if (!vocabSize) return []; + + const tokenIds: number[] = []; + for (let tokenId = 0; tokenId < vocabSize && tokenIds.length < 25; ++tokenId) { + if (isAllowed(mask, tokenId, vocabSize)) tokenIds.push(tokenId); + } + return tokenIds; +} + +function isAllowed(mask: GuidanceMask, tokenId: number, vocabSize: number) { + if (mask.length >= vocabSize) { + return Boolean(mask[tokenId]); + } + return Boolean(Number(mask[tokenId >> 5]) & (1 << (tokenId & 31))); +} diff --git a/packages/transformers-llguidance/src/utils/runtime.ts b/packages/transformers-llguidance/src/utils/runtime.ts new file mode 100644 index 000000000..f2ba25566 --- /dev/null +++ b/packages/transformers-llguidance/src/utils/runtime.ts @@ -0,0 +1,3 @@ +export function isNodeLikeRuntime() { + return typeof process !== 'undefined' && Boolean(process.versions?.node); +} diff --git a/packages/transformers-llguidance/src/utils/types.ts b/packages/transformers-llguidance/src/utils/types.ts new file mode 100644 index 000000000..3024c3601 --- /dev/null +++ b/packages/transformers-llguidance/src/utils/types.ts @@ -0,0 +1,23 @@ +export type GuidanceMask = Uint32Array | Uint8Array | boolean[] | number[]; + +export type GuidanceMaskResult = + | { mask: GuidanceMask; vocabSize?: number; stop?: false } + | { stop: true } + | { backtrack?: number; ffTokens?: number[] }; + +export type GuidanceCommitResult = { + stop?: boolean; + backtrack?: number; + ffTokens?: number[]; +}; + +export type GuidanceInterpreter = { + computeMask(): GuidanceMaskResult; + commitToken(tokenId: number): GuidanceCommitResult | undefined; +}; + +export type LlguidanceState = { + completed: boolean; + interpreter: GuidanceInterpreter; + step: number; +}; diff --git a/packages/transformers-llguidance/src/utils/wasm.ts b/packages/transformers-llguidance/src/utils/wasm.ts new file mode 100644 index 000000000..1e0324104 --- /dev/null +++ b/packages/transformers-llguidance/src/utils/wasm.ts @@ -0,0 +1,75 @@ +import { env, loadWasmBinary, loadWasmFactory, logger } from '@huggingface/transformers'; +import { type LoadBundledLLGuidanceOptions, loadBundledLLGuidance } from 'llguidance'; + +import { DEFAULT_LLGUIDANCE_WASM_FACTORY_URL, DEFAULT_LLGUIDANCE_WASM_URL } from '../constants'; +import { isNodeLikeRuntime } from './runtime'; + +export type LlguidanceLoadOptions = LoadBundledLLGuidanceOptions & { + /** Whether to pre-load and cache llguidance WASM assets. Defaults to env.useWasmCache. */ + useWasmCache?: boolean; +}; + +export async function loadLLGuidance(loadOptions: LlguidanceLoadOptions) { + const { useWasmCache = env.useWasmCache, ...options } = loadOptions; + if (!useWasmCache || isNodeLikeRuntime() || options.wasmFactory) { + return loadBundledLLGuidance(options); + } + + const wasmSource = options.wasm ?? options.wasmUrl ?? DEFAULT_LLGUIDANCE_WASM_URL; + const wasmFactorySource = options.wasmFactoryUrl ?? DEFAULT_LLGUIDANCE_WASM_FACTORY_URL; + const cachedOptions = { ...options }; + + const [wasm, wasmFactoryUrl] = await Promise.all([ + loadCacheableWasm(wasmSource), + loadCacheableWasmFactory(wasmFactorySource), + ]); + + if (wasm) { + cachedOptions.wasm = wasm; + delete cachedOptions.wasmUrl; + } + + if (wasm && wasmFactoryUrl) { + cachedOptions.wasmFactoryUrl = wasmFactoryUrl; + } + + return loadBundledLLGuidance(cachedOptions); +} + +async function loadCacheableWasm(source: LoadBundledLLGuidanceOptions['wasm']) { + const url = toCacheableURL(source); + if (!url) return null; + + try { + return await loadWasmBinary(url); + } catch (error) { + logger.warn('Failed to pre-load llguidance WASM binary:', error); + return null; + } +} + +async function loadCacheableWasmFactory(source: LoadBundledLLGuidanceOptions['wasmFactoryUrl']) { + const url = toCacheableURL(source); + if (!url) return null; + + try { + return await loadWasmFactory(url); + } catch (error) { + logger.warn('Failed to pre-load llguidance WASM factory:', error); + return null; + } +} + +function toCacheableURL(source: unknown) { + if (typeof source === 'string') { + return isBlobURL(source) ? null : new URL(source, globalThis.location?.href).href; + } + if (source instanceof URL) { + return isBlobURL(source.href) ? null : source.href; + } + return null; +} + +function isBlobURL(url: string) { + return url.startsWith('blob:'); +} diff --git a/packages/transformers-llguidance/tests/llguidance-constraint.test.js b/packages/transformers-llguidance/tests/llguidance-constraint.test.js index c9f32a42d..b12f66259 100644 --- a/packages/transformers-llguidance/tests/llguidance-constraint.test.js +++ b/packages/transformers-llguidance/tests/llguidance-constraint.test.js @@ -20,6 +20,7 @@ describe("LlguidanceConstraint", () => { loadBundledLLGuidance.mockClear(); }); + // Verifies the core integration path: llguidance creates an interpreter and its mask is applied across batches. it("loads llguidance and applies masks", async () => { computeMask.mockReturnValue({ mask: [true, false, true, false], vocabSize: 4 }); @@ -35,12 +36,14 @@ describe("LlguidanceConstraint", () => { expect(Array.from(logits.data)).toEqual([1, -Infinity, 3, -Infinity, 5, -Infinity, 7, -Infinity]); }); + // Ensures caller-provided load options are forwarded without the cache-only control flag. it("passes explicit llguidance load options", async () => { await LlguidanceConstraint.fromResponseFormat({}, { type: "json_object" }, { useWasmCache: false, wasmUrl: "custom.wasm" }); expect(loadBundledLLGuidance).toHaveBeenCalledWith({ wasmUrl: "custom.wasm" }); }); + // Confirms sampled tokens are committed back to llguidance so stopping criteria can end generation. it("commits sampled tokens and stops when llguidance stops", async () => { computeMask.mockReturnValue({ mask: [true, true], vocabSize: 2 }); commitToken.mockReturnValueOnce(undefined).mockReturnValueOnce({ stop: true }); @@ -49,10 +52,76 @@ describe("LlguidanceConstraint", () => { expect(stopping_criteria([[0n]])).toEqual([false]); - logits_processor.onTokensSampled([0, 1], [[0n, 0n], [0n, 1n]]); + logits_processor.onTokensSampled( + [0, 1], + [ + [0n, 0n], + [0n, 1n], + ], + ); expect(commitToken).toHaveBeenCalledWith(0); expect(commitToken).toHaveBeenCalledWith(1); - expect(stopping_criteria([[0n, 0n], [0n, 1n]])).toEqual([true, true]); + expect( + stopping_criteria([ + [0n, 0n], + [0n, 1n], + ]), + ).toEqual([true, true]); + }); + + // Covers llguidance's compact bitset mask format, not just boolean arrays. + it("applies packed uint32 masks", async () => { + computeMask.mockReturnValue({ mask: new Uint32Array([0b0101]), vocabSize: 4 }); + + const { logits_processor } = await LlguidanceConstraint.fromResponseFormat({}, { type: "json_object" }); + const logits = new Tensor("float32", new Float32Array([1, 2, 3, 4]), [1, 4]); + + logits_processor([[0n]], logits); + + expect(Array.from(logits.data)).toEqual([1, -Infinity, 3, -Infinity]); + }); + + // Ensures a stop response from computeMask marks the shared state complete without committing more tokens. + it("stops generation when computeMask returns stop", async () => { + computeMask.mockReturnValue({ stop: true }); + + const { logits_processor, stopping_criteria } = await LlguidanceConstraint.fromResponseFormat({}, { type: "json_object" }); + const logits = new Tensor("float32", new Float32Array([1, 2]), [1, 2]); + + logits_processor([[0n]], logits); + logits_processor.onTokensSampled([0], [[0n]]); + + expect(Array.from(logits.data)).toEqual([1, 2]); + expect(commitToken).not.toHaveBeenCalled(); + expect(stopping_criteria([[0n]])).toEqual([true]); + }); + + // Handles llguidance's post-stop compute error as a normal completion signal. + it("treats compute_mask after stop as completed", async () => { + computeMask.mockImplementation(() => { + throw new Error("compute_mask() called after stop"); + }); + + const { logits_processor, stopping_criteria } = await LlguidanceConstraint.fromResponseFormat({}, { type: "json_object" }); + const logits = new Tensor("float32", new Float32Array([1, 2]), [1, 2]); + + logits_processor([[0n]], logits); + + expect(Array.from(logits.data)).toEqual([1, 2]); + expect(stopping_criteria([[0n]])).toEqual([true]); + }); + + // Keeps unexpected interpreter failures visible instead of swallowing real bugs. + it("rethrows unexpected computeMask errors", async () => { + const error = new Error("unexpected"); + computeMask.mockImplementation(() => { + throw error; + }); + + const { logits_processor } = await LlguidanceConstraint.fromResponseFormat({}, { type: "json_object" }); + const logits = new Tensor("float32", new Float32Array([1, 2]), [1, 2]); + + expect(() => logits_processor([[0n]], logits)).toThrow(error); }); }); diff --git a/packages/transformers-llguidance/tests/wasm-loading.test.js b/packages/transformers-llguidance/tests/wasm-loading.test.js new file mode 100644 index 000000000..853a75da9 --- /dev/null +++ b/packages/transformers-llguidance/tests/wasm-loading.test.js @@ -0,0 +1,123 @@ +import { jest } from "@jest/globals"; + +const createInterpreter = jest.fn(() => ({ + computeMask: jest.fn(() => ({ stop: true })), + commitToken: jest.fn(), +})); +const loadBundledLLGuidance = jest.fn(async () => ({ createInterpreter })); +const loadWasmBinary = jest.fn(); +const loadWasmFactory = jest.fn(); +const logger = { + debug: jest.fn(), + warn: jest.fn(), +}; + +class LogitsProcessor {} +class LogitsProcessorList extends Array {} +class StoppingCriteria {} + +jest.unstable_mockModule("@huggingface/transformers", () => ({ + LogitsProcessor, + LogitsProcessorList, + StoppingCriteria, + env: { useWasmCache: true }, + loadWasmBinary, + loadWasmFactory, + logger, +})); + +jest.unstable_mockModule("llguidance", () => ({ + loadBundledLLGuidance, +})); + +const originalProcess = globalThis.process; +const originalLocationDescriptor = Object.getOwnPropertyDescriptor(globalThis, "location"); +const { LlguidanceConstraint } = await import("../dist/index.js"); + +describe("llguidance WASM loading", () => { + beforeEach(() => { + delete globalThis.process; + Object.defineProperty(globalThis, "location", { + configurable: true, + value: { href: "https://example.test/app/" }, + }); + + createInterpreter.mockClear(); + loadBundledLLGuidance.mockClear(); + loadWasmBinary.mockReset(); + loadWasmFactory.mockReset(); + logger.debug.mockClear(); + logger.warn.mockClear(); + }); + + afterEach(() => { + globalThis.process = originalProcess; + if (originalLocationDescriptor) { + Object.defineProperty(globalThis, "location", originalLocationDescriptor); + } else { + delete globalThis.location; + } + }); + + // Exercises browser-like cache preloading with the package's default CDN asset URLs. + it("preloads default WASM assets when cache is enabled outside Node", async () => { + const wasm = new Uint8Array([1, 2, 3]); + loadWasmBinary.mockResolvedValue(wasm); + loadWasmFactory.mockResolvedValue("blob:factory-url"); + + await LlguidanceConstraint.fromResponseFormat({}, { type: "json_object" }); + + expect(loadWasmBinary).toHaveBeenCalledWith("https://cdn.jsdelivr.net/npm/llguidance@0.1.7/wasm/llguidance_wasm_bg.wasm"); + expect(loadWasmFactory).toHaveBeenCalledWith("https://cdn.jsdelivr.net/npm/llguidance@0.1.7/wasm/llguidance_wasm.js"); + expect(loadBundledLLGuidance).toHaveBeenCalledWith({ + wasm, + wasmFactoryUrl: "blob:factory-url", + }); + }); + + // Verifies custom relative and absolute URLs are normalized before using the transformers cache helpers. + it("resolves custom cacheable WASM URLs", async () => { + const wasm = new Uint8Array([4, 5, 6]); + loadWasmBinary.mockResolvedValue(wasm); + loadWasmFactory.mockResolvedValue("blob:custom-factory-url"); + + await LlguidanceConstraint.fromResponseFormat( + {}, + { type: "json_object" }, + { + wasmUrl: "assets/custom.wasm", + wasmFactoryUrl: new URL("https://cdn.test/custom-factory.js"), + }, + ); + + expect(loadWasmBinary).toHaveBeenCalledWith("https://example.test/app/assets/custom.wasm"); + expect(loadWasmFactory).toHaveBeenCalledWith("https://cdn.test/custom-factory.js"); + expect(loadBundledLLGuidance).toHaveBeenCalledWith({ + wasm, + wasmFactoryUrl: "blob:custom-factory-url", + }); + }); + + // Ensures a failed preload does not prevent llguidance from loading with the original options. + it("falls back to uncached options when WASM preload fails", async () => { + const error = new Error("network failure"); + loadWasmBinary.mockRejectedValue(error); + loadWasmFactory.mockResolvedValue("blob:factory-url"); + + await LlguidanceConstraint.fromResponseFormat({}, { type: "json_object" }, { wasmUrl: "custom.wasm" }); + + expect(logger.warn).toHaveBeenCalledWith("Failed to pre-load llguidance WASM binary:", error); + expect(loadBundledLLGuidance).toHaveBeenCalledWith({ wasmUrl: "custom.wasm" }); + }); + + // Avoids cache preloading when the caller already supplied an initialized factory. + it("skips preloading when a WASM factory is provided", async () => { + const wasmFactory = jest.fn(); + + await LlguidanceConstraint.fromResponseFormat({}, { type: "json_object" }, { wasmFactory }); + + expect(loadWasmBinary).not.toHaveBeenCalled(); + expect(loadWasmFactory).not.toHaveBeenCalled(); + expect(loadBundledLLGuidance).toHaveBeenCalledWith({ wasmFactory }); + }); +}); From a7ab7f3327bc038e2bf8fed40702636081f2f143 Mon Sep 17 00:00:00 2001 From: Nico Martin Date: Sat, 11 Jul 2026 08:41:38 +0200 Subject: [PATCH 09/17] clean up --- packages/transformers-llguidance/src/utils/mask.ts | 1 - 1 file changed, 1 deletion(-) diff --git a/packages/transformers-llguidance/src/utils/mask.ts b/packages/transformers-llguidance/src/utils/mask.ts index f1ee29ba2..d73df19e3 100644 --- a/packages/transformers-llguidance/src/utils/mask.ts +++ b/packages/transformers-llguidance/src/utils/mask.ts @@ -1,5 +1,4 @@ import { type Tensor } from '@huggingface/transformers'; - import { type GuidanceMask, type GuidanceMaskResult } from './types'; export function summarizeMaskResult(result: GuidanceMaskResult, vocabSize?: number) { From 1eb0175f53a45b1544479be3d42742bea1bc376b Mon Sep 17 00:00:00 2001 From: Nico Martin Date: Thu, 30 Jul 2026 14:34:02 +0200 Subject: [PATCH 10/17] performance improvements --- packages/transformers-llguidance/package.json | 2 +- .../src/LlguidanceConstraint.ts | 175 ++++++++++++------ .../transformers-llguidance/src/constants.ts | 2 +- packages/transformers-llguidance/src/index.ts | 2 +- .../transformers-llguidance/src/utils/mask.ts | 78 ++++++-- .../src/utils/types.ts | 12 +- .../tests/wasm-loading.test.js | 4 +- pnpm-lock.yaml | 10 +- 8 files changed, 204 insertions(+), 81 deletions(-) diff --git a/packages/transformers-llguidance/package.json b/packages/transformers-llguidance/package.json index 27fe44062..fee459238 100644 --- a/packages/transformers-llguidance/package.json +++ b/packages/transformers-llguidance/package.json @@ -63,6 +63,6 @@ "access": "public" }, "dependencies": { - "llguidance": "^0.1.7" + "llguidance": "^0.1.8" } } diff --git a/packages/transformers-llguidance/src/LlguidanceConstraint.ts b/packages/transformers-llguidance/src/LlguidanceConstraint.ts index bee7b774e..1d7e5c4cd 100644 --- a/packages/transformers-llguidance/src/LlguidanceConstraint.ts +++ b/packages/transformers-llguidance/src/LlguidanceConstraint.ts @@ -1,12 +1,24 @@ -import { LogitsProcessor, LogitsProcessorList, StoppingCriteria, logger, type Tensor } from '@huggingface/transformers'; +import { + env, + LogitsProcessor, + LogitsProcessorList, + StoppingCriteria, + logger, + type Tensor, +} from '@huggingface/transformers'; import { type LLGuidanceResponseFormat } from 'llguidance'; -import { applyMask, summarizeMaskResult } from './utils/mask'; -import { type GuidanceInterpreter, type GuidanceMaskResult, type LlguidanceState } from './utils/types'; +import { applyMask, forceToken, summarizeMaskResult } from './utils/mask'; +import { + type GuidanceInterpreter, + type GuidanceMaskResult, + type LlguidanceState, + type LlguidanceStats, +} from './utils/types'; import { type LlguidanceLoadOptions, loadLLGuidance } from './utils/wasm'; export type ResponseFormat = LLGuidanceResponseFormat; -export type { LlguidanceLoadOptions }; +export type { LlguidanceLoadOptions, LlguidanceStats }; export class LlguidanceConstraint { static async fromResponseFormat( @@ -30,6 +42,7 @@ export class LlguidanceConstraint { completed: false, interpreter, step: 0, + stats: { steps: 0, computeMaskMs: 0, applyMaskMs: 0, commitTokenMs: 0 }, }; const logits_processor = new LogitsProcessorList(); @@ -38,6 +51,7 @@ export class LlguidanceConstraint { return { logits_processor, stopping_criteria: new LlguidanceStoppingCriteria(state), + stats: state.stats, }; } } @@ -52,92 +66,136 @@ class LlguidanceLogitsProcessor extends LogitsProcessor { _call(_inputIds: bigint[][], logits: Tensor) { if (this.state.completed) { - logger.debug('[LlguidanceLogitsProcessor] skip completed', { - step: this.state.step, - }); + if (isDebugEnabled()) { + logger.debug('[LlguidanceLogitsProcessor] skip completed', { + step: this.state.step, + }); + } return logits; } this.state.step++; + this.state.stats.steps = this.state.step; const vocabSize = logits.dims.at(-1); - logger.debug('[LlguidanceLogitsProcessor] compute mask', { - step: this.state.step, - logitsDims: logits.dims, - vocabSize, - }); + if (isDebugEnabled()) { + logger.debug('[LlguidanceLogitsProcessor] compute mask', { + step: this.state.step, + logitsDims: logits.dims, + vocabSize, + }); + } let result: GuidanceMaskResult; + const maskStart = performance.now(); try { result = this.state.interpreter.computeMask(); + this.state.stats.computeMaskMs += performance.now() - maskStart; } catch (error) { + this.state.stats.computeMaskMs += performance.now() - maskStart; if (!String((error as Error).message).includes('compute_mask() called after stop')) { throw error; } this.state.completed = true; - logger.debug('[LlguidanceLogitsProcessor] compute after stop', { - step: this.state.step, - }); + if (isDebugEnabled()) { + logger.debug('[LlguidanceLogitsProcessor] compute after stop', { + step: this.state.step, + }); + } return logits; } - logger.debug('[LlguidanceLogitsProcessor] mask result', { - step: this.state.step, - result: summarizeMaskResult(result, vocabSize), - }); + if (isDebugEnabled()) { + logger.debug('[LlguidanceLogitsProcessor] mask result', { + step: this.state.step, + result: summarizeMaskResult(result, vocabSize), + }); + } if ('stop' in result && result.stop) { this.state.completed = true; - logger.debug('[LlguidanceLogitsProcessor] stopped by mask', { - step: this.state.step, - }); + if (isDebugEnabled()) { + logger.debug('[LlguidanceLogitsProcessor] stopped by mask', { + step: this.state.step, + }); + } + return logits; + } + + if ('ffTokens' in result && result.ffTokens?.length) { + // Fast-forward splice: the grammar forces the next token, so ban + // everything else instead of letting the model sample unconstrained. + forceToken(logits, result.ffTokens[0], vocabSize); + if (isDebugEnabled()) { + logger.debug('[LlguidanceLogitsProcessor] forced splice token', { + step: this.state.step, + ffTokens: Array.from(result.ffTokens), + backtrack: 'backtrack' in result ? result.backtrack : undefined, + }); + } return logits; } if ('mask' in result) { - const applied = applyMask(logits, result.mask, result.vocabSize ?? vocabSize); - logger.debug('[LlguidanceLogitsProcessor] mask applied', { - step: this.state.step, - ...applied, - }); + const applyStart = performance.now(); + if (isDebugEnabled()) { + const applied = applyMask(logits, result.mask, result.vocabSize ?? vocabSize, true); + logger.debug('[LlguidanceLogitsProcessor] mask applied', { + step: this.state.step, + ...applied, + }); + } else { + applyMask(logits, result.mask, result.vocabSize ?? vocabSize); + } + this.state.stats.applyMaskMs += performance.now() - applyStart; } return logits; } onTokenSampled(tokenId: number, batchIdx: number, inputIds: bigint[][]) { - logger.debug('[LlguidanceLogitsProcessor] token sampled', { - step: this.state.step, - tokenId, - batchIdx, - inputLength: inputIds[batchIdx]?.length, - completed: this.state.completed, - }); + if (isDebugEnabled()) { + logger.debug('[LlguidanceLogitsProcessor] token sampled', { + step: this.state.step, + tokenId, + batchIdx, + inputLength: inputIds[batchIdx]?.length, + completed: this.state.completed, + }); + } if (this.state.completed) return; + const commitStart = performance.now(); const result = this.state.interpreter.commitToken(tokenId); - logger.debug('[LlguidanceLogitsProcessor] token committed', { - step: this.state.step, - tokenId, - result, - }); - - if (result?.stop) { - this.state.completed = true; - logger.debug('[LlguidanceLogitsProcessor] stopped by commit', { + this.state.stats.commitTokenMs += performance.now() - commitStart; + if (isDebugEnabled()) { + logger.debug('[LlguidanceLogitsProcessor] token committed', { step: this.state.step, tokenId, + result, }); } + + if (result?.stop) { + this.state.completed = true; + if (isDebugEnabled()) { + logger.debug('[LlguidanceLogitsProcessor] stopped by commit', { + step: this.state.step, + tokenId, + }); + } + } } onTokensSampled(tokenIds: number[], inputIds: bigint[][]) { - logger.debug('[LlguidanceLogitsProcessor] tokens sampled', { - step: this.state.step, - tokenIds, - inputLengths: inputIds.map((ids) => ids.length), - completed: this.state.completed, - }); + if (isDebugEnabled()) { + logger.debug('[LlguidanceLogitsProcessor] tokens sampled', { + step: this.state.step, + tokenIds, + inputLengths: inputIds.map((ids) => ids.length), + completed: this.state.completed, + }); + } for (let batchIdx = 0; batchIdx < tokenIds.length; ++batchIdx) { this.onTokenSampled(tokenIds[batchIdx], batchIdx, inputIds); @@ -155,12 +213,19 @@ class LlguidanceStoppingCriteria extends StoppingCriteria { _call(inputIds: ArrayLike[]) { const result = new Array(inputIds.length).fill(this.state.completed); - logger.debug('[LlguidanceStoppingCriteria] call', { - step: this.state.step, - completed: this.state.completed, - result, - inputLengths: inputIds.map((ids) => ids.length), - }); + if (isDebugEnabled()) { + logger.debug('[LlguidanceStoppingCriteria] call', { + step: this.state.step, + completed: this.state.completed, + result, + inputLengths: inputIds.map((ids) => ids.length), + }); + } return result; } } + +function isDebugEnabled() { + // Keep compatibility with Transformers.js releases that predate the LogLevel export. + return env.logLevel <= 10; +} diff --git a/packages/transformers-llguidance/src/constants.ts b/packages/transformers-llguidance/src/constants.ts index 773108fac..2b68f0381 100644 --- a/packages/transformers-llguidance/src/constants.ts +++ b/packages/transformers-llguidance/src/constants.ts @@ -1,4 +1,4 @@ -export const LLGUIDANCE_VERSION = '0.1.7'; +export const LLGUIDANCE_VERSION = '0.1.8'; export const LLGUIDANCE_WASM_BASE = `https://cdn.jsdelivr.net/npm/llguidance@${LLGUIDANCE_VERSION}/wasm/`; export const DEFAULT_LLGUIDANCE_WASM_URL = `${LLGUIDANCE_WASM_BASE}llguidance_wasm_bg.wasm`; export const DEFAULT_LLGUIDANCE_WASM_FACTORY_URL = `${LLGUIDANCE_WASM_BASE}llguidance_wasm.js`; diff --git a/packages/transformers-llguidance/src/index.ts b/packages/transformers-llguidance/src/index.ts index a870a56f9..2a7dd1267 100644 --- a/packages/transformers-llguidance/src/index.ts +++ b/packages/transformers-llguidance/src/index.ts @@ -1,2 +1,2 @@ export { LlguidanceConstraint } from './LlguidanceConstraint'; -export type { LlguidanceLoadOptions, ResponseFormat } from './LlguidanceConstraint'; +export type { LlguidanceLoadOptions, LlguidanceStats, ResponseFormat } from './LlguidanceConstraint'; diff --git a/packages/transformers-llguidance/src/utils/mask.ts b/packages/transformers-llguidance/src/utils/mask.ts index d73df19e3..a63136c04 100644 --- a/packages/transformers-llguidance/src/utils/mask.ts +++ b/packages/transformers-llguidance/src/utils/mask.ts @@ -1,6 +1,8 @@ import { type Tensor } from '@huggingface/transformers'; import { type GuidanceMask, type GuidanceMaskResult } from './types'; +type LogitsData = Float32Array | Float64Array | number[]; + export function summarizeMaskResult(result: GuidanceMaskResult, vocabSize?: number) { if ('stop' in result && result.stop) { return { stop: true }; @@ -18,29 +20,77 @@ export function summarizeMaskResult(result: GuidanceMaskResult, vocabSize?: numb }; } -export function applyMask(logits: Tensor, mask: GuidanceMask, vocabSize?: number) { +export function applyMask(logits: Tensor, mask: GuidanceMask, vocabSize?: number, includeSummary = false) { if (!vocabSize) { - return { vocabSize, batchSize: 0, masked: 0, allowed: undefined }; + return includeSummary ? { vocabSize, batchSize: 0, masked: 0, allowed: undefined } : undefined; } - const data = logits.data as Float32Array | Float64Array | number[]; - const batchSize = Math.max(1, data.length / vocabSize); - let masked = 0; - let allowed = 0; + const data = logits.data as LogitsData; + const stride = (logits.dims?.at?.(-1) as number) || vocabSize; + const bound = Math.min(vocabSize, stride); + const batchSize = Math.max(1, Math.floor(data.length / stride)); + const packed = mask.length < vocabSize; for (let batch = 0; batch < batchSize; ++batch) { - const offset = batch * vocabSize; - for (let tokenId = 0; tokenId < vocabSize; ++tokenId) { - if (!isAllowed(mask, tokenId, vocabSize)) { - data[offset + tokenId] = -Infinity; - masked++; - } else if (batch === 0) { - allowed++; + const offset = batch * stride; + if (packed) { + applyPackedMask(data, mask, offset, bound); + } else { + for (let tokenId = 0; tokenId < bound; ++tokenId) { + if (!mask[tokenId]) data[offset + tokenId] = -Infinity; } } + if (stride > vocabSize) { + // Logits padded beyond the grammar vocab can never be committed. + data.fill(-Infinity, offset + vocabSize, offset + stride); + } + } + + if (!includeSummary) return undefined; + + const allowed = countAllowed(mask, vocabSize) ?? 0; + return { vocabSize, batchSize, masked: (vocabSize - allowed) * batchSize, allowed }; +} + +// Forces a single token by banning everything else, e.g. for fast-forward splices. +export function forceToken(logits: Tensor, tokenId: number, vocabSize?: number) { + if (!vocabSize) return; + + const data = logits.data as LogitsData; + const stride = (logits.dims?.at?.(-1) as number) || vocabSize; + const batchSize = Math.max(1, Math.floor(data.length / stride)); + + for (let batch = 0; batch < batchSize; ++batch) { + const offset = batch * stride; + const kept = data[offset + tokenId]; + data.fill(-Infinity, offset, offset + stride); + data[offset + tokenId] = Number.isFinite(kept) ? kept : 0; + } +} + +// Hot path: runs once per generated token over the whole vocab. Grammar masks +// are skewed — most 32-token words are either fully allowed (skip) or fully +// banned (memset) — so per-bit work only happens on the few mixed words. +function applyPackedMask(data: LogitsData, mask: GuidanceMask, offset: number, vocabSize: number) { + const numWords = vocabSize >>> 5; + for (let word = 0; word < numWords; ++word) { + const bits = (mask[word] as number) | 0; + if (bits === -1) continue; + const base = offset + (word << 5); + if (bits === 0) { + data.fill(-Infinity, base, base + 32); + continue; + } + for (let bit = 0; bit < 32; ++bit) { + if (!(bits & (1 << bit))) data[base + bit] = -Infinity; + } } - return { vocabSize, batchSize, masked, allowed }; + for (let tokenId = numWords << 5; tokenId < vocabSize; ++tokenId) { + if (!((mask[tokenId >>> 5] as number) & (1 << (tokenId & 31)))) { + data[offset + tokenId] = -Infinity; + } + } } function countAllowed(mask: GuidanceMask, vocabSize?: number) { diff --git a/packages/transformers-llguidance/src/utils/types.ts b/packages/transformers-llguidance/src/utils/types.ts index 3024c3601..d3ba9f24c 100644 --- a/packages/transformers-llguidance/src/utils/types.ts +++ b/packages/transformers-llguidance/src/utils/types.ts @@ -3,12 +3,12 @@ export type GuidanceMask = Uint32Array | Uint8Array | boolean[] | number[]; export type GuidanceMaskResult = | { mask: GuidanceMask; vocabSize?: number; stop?: false } | { stop: true } - | { backtrack?: number; ffTokens?: number[] }; + | { backtrack?: number; ffTokens?: number[] | Uint32Array }; export type GuidanceCommitResult = { stop?: boolean; backtrack?: number; - ffTokens?: number[]; + ffTokens?: number[] | Uint32Array; }; export type GuidanceInterpreter = { @@ -16,8 +16,16 @@ export type GuidanceInterpreter = { commitToken(tokenId: number): GuidanceCommitResult | undefined; }; +export type LlguidanceStats = { + steps: number; + computeMaskMs: number; + applyMaskMs: number; + commitTokenMs: number; +}; + export type LlguidanceState = { completed: boolean; interpreter: GuidanceInterpreter; step: number; + stats: LlguidanceStats; }; diff --git a/packages/transformers-llguidance/tests/wasm-loading.test.js b/packages/transformers-llguidance/tests/wasm-loading.test.js index 853a75da9..c218a5228 100644 --- a/packages/transformers-llguidance/tests/wasm-loading.test.js +++ b/packages/transformers-llguidance/tests/wasm-loading.test.js @@ -67,8 +67,8 @@ describe("llguidance WASM loading", () => { await LlguidanceConstraint.fromResponseFormat({}, { type: "json_object" }); - expect(loadWasmBinary).toHaveBeenCalledWith("https://cdn.jsdelivr.net/npm/llguidance@0.1.7/wasm/llguidance_wasm_bg.wasm"); - expect(loadWasmFactory).toHaveBeenCalledWith("https://cdn.jsdelivr.net/npm/llguidance@0.1.7/wasm/llguidance_wasm.js"); + expect(loadWasmBinary).toHaveBeenCalledWith("https://cdn.jsdelivr.net/npm/llguidance@0.1.8/wasm/llguidance_wasm_bg.wasm"); + expect(loadWasmFactory).toHaveBeenCalledWith("https://cdn.jsdelivr.net/npm/llguidance@0.1.8/wasm/llguidance_wasm.js"); expect(loadBundledLLGuidance).toHaveBeenCalledWith({ wasm, wasmFactoryUrl: "blob:factory-url", diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index b0585df63..379ccfc6b 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -61,8 +61,8 @@ importers: packages/transformers-llguidance: dependencies: llguidance: - specifier: ^0.1.7 - version: 0.1.7 + specifier: ^0.1.8 + version: 0.1.8 devDependencies: '@huggingface/transformers': specifier: workspace:* @@ -1585,8 +1585,8 @@ packages: linkify-it@5.0.0: resolution: {integrity: sha512-5aHCbzQRADcdP+ATqnDuhhJ/MRIqDkZX5pyjFHRRysS8vZ5AbqGEoFIb6pYHPZ+L/OC2Lc+xT8uHVVR5CAK/wQ==} - llguidance@0.1.7: - resolution: {integrity: sha512-r55h3hDmkq313cFlqbKyx4IebqpreyrrnMJip7IR+ocwPGm7O6ZoDVCIAEg2wd6k7pQ5wWyuSzV9rSefqKVs6Q==} + llguidance@0.1.8: + resolution: {integrity: sha512-n44NMtq2ChgUGkcy6G7o8TUp90kyK2TergBDZ2nvxUN87vtkmCm/WxEWG0hl3Pav1BhqxybKAR6DP+9oib0Law==} locate-path@5.0.0: resolution: {integrity: sha512-t7hw9pI+WvuwNJXwk5zVHpyhIqzg2qTlklJOf0mVxGSbe3Fp2VieZcduNYjaLDoy6p9uGpQEGWG87WpMKlNq8g==} @@ -3750,7 +3750,7 @@ snapshots: dependencies: uc.micro: 2.1.0 - llguidance@0.1.7: {} + llguidance@0.1.8: {} locate-path@5.0.0: dependencies: From c7d04003394b286c95250513dfd5c0df25c1629e Mon Sep 17 00:00:00 2001 From: Nico Martin Date: Fri, 31 Jul 2026 10:04:26 +0200 Subject: [PATCH 11/17] updated llguidamce to 0.2.0 --- packages/transformers-llguidance/package.json | 2 +- .../src/LlguidanceConstraint.ts | 54 ++++---- .../transformers-llguidance/src/constants.ts | 4 - packages/transformers-llguidance/src/index.ts | 2 +- .../src/utils/runtime.ts | 3 - .../src/utils/types.ts | 12 ++ .../transformers-llguidance/src/utils/wasm.ts | 75 ----------- .../tests/llguidance-constraint.test.js | 46 ++++--- .../tests/wasm-loading.test.js | 123 ------------------ pnpm-lock.yaml | 10 +- 10 files changed, 74 insertions(+), 257 deletions(-) delete mode 100644 packages/transformers-llguidance/src/constants.ts delete mode 100644 packages/transformers-llguidance/src/utils/runtime.ts delete mode 100644 packages/transformers-llguidance/src/utils/wasm.ts delete mode 100644 packages/transformers-llguidance/tests/wasm-loading.test.js diff --git a/packages/transformers-llguidance/package.json b/packages/transformers-llguidance/package.json index fee459238..66fe660c9 100644 --- a/packages/transformers-llguidance/package.json +++ b/packages/transformers-llguidance/package.json @@ -63,6 +63,6 @@ "access": "public" }, "dependencies": { - "llguidance": "^0.1.8" + "llguidance": "0.2.0" } } diff --git a/packages/transformers-llguidance/src/LlguidanceConstraint.ts b/packages/transformers-llguidance/src/LlguidanceConstraint.ts index 1d7e5c4cd..4c917f73a 100644 --- a/packages/transformers-llguidance/src/LlguidanceConstraint.ts +++ b/packages/transformers-llguidance/src/LlguidanceConstraint.ts @@ -6,7 +6,7 @@ import { logger, type Tensor, } from '@huggingface/transformers'; -import { type LLGuidanceResponseFormat } from 'llguidance'; +import { loadBundledLLGuidance, type JSONSchema, type LLGuidanceTokenizerSource } from 'llguidance'; import { applyMask, forceToken, summarizeMaskResult } from './utils/mask'; import { @@ -15,22 +15,20 @@ import { type LlguidanceState, type LlguidanceStats, } from './utils/types'; -import { type LlguidanceLoadOptions, loadLLGuidance } from './utils/wasm'; -export type ResponseFormat = LLGuidanceResponseFormat; -export type { LlguidanceLoadOptions, LlguidanceStats }; +export type ResponseFormat = + | { type: 'json_object' } + | { type: 'json_schema'; json_schema: JSONSchema } + | { type: 'regex'; regex: string }; +export type { LlguidanceStats }; export class LlguidanceConstraint { - static async fromResponseFormat( - tokenizer: unknown, - response_format: ResponseFormat, - loadOptions: LlguidanceLoadOptions = {}, - ) { + static async fromResponseFormat(tokenizer: LLGuidanceTokenizerSource, response_format: ResponseFormat) { logger.debug('[LlguidanceConstraint] loading llguidance', { response_format, }); - const runtime = await loadLLGuidance(loadOptions); + const runtime = await loadBundledLLGuidance(); const interpreter = runtime.createInterpreter({ tokenizer, response_format, @@ -42,7 +40,15 @@ export class LlguidanceConstraint { completed: false, interpreter, step: 0, - stats: { steps: 0, computeMaskMs: 0, applyMaskMs: 0, commitTokenMs: 0 }, + stats: { + steps: 0, + computeMaskMs: 0, + applyMaskMs: 0, + commitTokenMs: 0, + trieNodesVisited: 0, + sharedJsonMaskCacheHits: 0, + sharedJsonMaskCacheMisses: 0, + }, }; const logits_processor = new LogitsProcessorList(); @@ -87,21 +93,19 @@ class LlguidanceLogitsProcessor extends LogitsProcessor { let result: GuidanceMaskResult; const maskStart = performance.now(); - try { + const maskWords = vocabSize === undefined ? undefined : Math.ceil(vocabSize / 32); + if (this.state.interpreter.computeMaskInto && maskWords !== undefined) { + this.state.maskBuffer ??= new Uint32Array(maskWords); + result = this.state.interpreter.computeMaskInto(this.state.maskBuffer); + } else { result = this.state.interpreter.computeMask(); - this.state.stats.computeMaskMs += performance.now() - maskStart; - } catch (error) { - this.state.stats.computeMaskMs += performance.now() - maskStart; - if (!String((error as Error).message).includes('compute_mask() called after stop')) { - throw error; - } - this.state.completed = true; - if (isDebugEnabled()) { - logger.debug('[LlguidanceLogitsProcessor] compute after stop', { - step: this.state.step, - }); - } - return logits; + } + this.state.stats.computeMaskMs += performance.now() - maskStart; + const maskInstrumentation = this.state.interpreter.maskInstrumentation?.(); + if (maskInstrumentation) { + this.state.stats.trieNodesVisited = maskInstrumentation.trieNodesVisited; + this.state.stats.sharedJsonMaskCacheHits = maskInstrumentation.sharedJsonMaskCacheHits; + this.state.stats.sharedJsonMaskCacheMisses = maskInstrumentation.sharedJsonMaskCacheMisses; } if (isDebugEnabled()) { diff --git a/packages/transformers-llguidance/src/constants.ts b/packages/transformers-llguidance/src/constants.ts deleted file mode 100644 index 2b68f0381..000000000 --- a/packages/transformers-llguidance/src/constants.ts +++ /dev/null @@ -1,4 +0,0 @@ -export const LLGUIDANCE_VERSION = '0.1.8'; -export const LLGUIDANCE_WASM_BASE = `https://cdn.jsdelivr.net/npm/llguidance@${LLGUIDANCE_VERSION}/wasm/`; -export const DEFAULT_LLGUIDANCE_WASM_URL = `${LLGUIDANCE_WASM_BASE}llguidance_wasm_bg.wasm`; -export const DEFAULT_LLGUIDANCE_WASM_FACTORY_URL = `${LLGUIDANCE_WASM_BASE}llguidance_wasm.js`; diff --git a/packages/transformers-llguidance/src/index.ts b/packages/transformers-llguidance/src/index.ts index 2a7dd1267..5f709e319 100644 --- a/packages/transformers-llguidance/src/index.ts +++ b/packages/transformers-llguidance/src/index.ts @@ -1,2 +1,2 @@ export { LlguidanceConstraint } from './LlguidanceConstraint'; -export type { LlguidanceLoadOptions, LlguidanceStats, ResponseFormat } from './LlguidanceConstraint'; +export type { LlguidanceStats, ResponseFormat } from './LlguidanceConstraint'; diff --git a/packages/transformers-llguidance/src/utils/runtime.ts b/packages/transformers-llguidance/src/utils/runtime.ts deleted file mode 100644 index f2ba25566..000000000 --- a/packages/transformers-llguidance/src/utils/runtime.ts +++ /dev/null @@ -1,3 +0,0 @@ -export function isNodeLikeRuntime() { - return typeof process !== 'undefined' && Boolean(process.versions?.node); -} diff --git a/packages/transformers-llguidance/src/utils/types.ts b/packages/transformers-llguidance/src/utils/types.ts index d3ba9f24c..325dc69b3 100644 --- a/packages/transformers-llguidance/src/utils/types.ts +++ b/packages/transformers-llguidance/src/utils/types.ts @@ -11,8 +11,16 @@ export type GuidanceCommitResult = { ffTokens?: number[] | Uint32Array; }; +export type GuidanceMaskInstrumentation = { + trieNodesVisited: number; + sharedJsonMaskCacheHits: number; + sharedJsonMaskCacheMisses: number; +}; + export type GuidanceInterpreter = { computeMask(): GuidanceMaskResult; + computeMaskInto?(target: Uint32Array): GuidanceMaskResult; + maskInstrumentation?(): GuidanceMaskInstrumentation; commitToken(tokenId: number): GuidanceCommitResult | undefined; }; @@ -21,11 +29,15 @@ export type LlguidanceStats = { computeMaskMs: number; applyMaskMs: number; commitTokenMs: number; + trieNodesVisited: number; + sharedJsonMaskCacheHits: number; + sharedJsonMaskCacheMisses: number; }; export type LlguidanceState = { completed: boolean; interpreter: GuidanceInterpreter; + maskBuffer?: Uint32Array; step: number; stats: LlguidanceStats; }; diff --git a/packages/transformers-llguidance/src/utils/wasm.ts b/packages/transformers-llguidance/src/utils/wasm.ts deleted file mode 100644 index 1e0324104..000000000 --- a/packages/transformers-llguidance/src/utils/wasm.ts +++ /dev/null @@ -1,75 +0,0 @@ -import { env, loadWasmBinary, loadWasmFactory, logger } from '@huggingface/transformers'; -import { type LoadBundledLLGuidanceOptions, loadBundledLLGuidance } from 'llguidance'; - -import { DEFAULT_LLGUIDANCE_WASM_FACTORY_URL, DEFAULT_LLGUIDANCE_WASM_URL } from '../constants'; -import { isNodeLikeRuntime } from './runtime'; - -export type LlguidanceLoadOptions = LoadBundledLLGuidanceOptions & { - /** Whether to pre-load and cache llguidance WASM assets. Defaults to env.useWasmCache. */ - useWasmCache?: boolean; -}; - -export async function loadLLGuidance(loadOptions: LlguidanceLoadOptions) { - const { useWasmCache = env.useWasmCache, ...options } = loadOptions; - if (!useWasmCache || isNodeLikeRuntime() || options.wasmFactory) { - return loadBundledLLGuidance(options); - } - - const wasmSource = options.wasm ?? options.wasmUrl ?? DEFAULT_LLGUIDANCE_WASM_URL; - const wasmFactorySource = options.wasmFactoryUrl ?? DEFAULT_LLGUIDANCE_WASM_FACTORY_URL; - const cachedOptions = { ...options }; - - const [wasm, wasmFactoryUrl] = await Promise.all([ - loadCacheableWasm(wasmSource), - loadCacheableWasmFactory(wasmFactorySource), - ]); - - if (wasm) { - cachedOptions.wasm = wasm; - delete cachedOptions.wasmUrl; - } - - if (wasm && wasmFactoryUrl) { - cachedOptions.wasmFactoryUrl = wasmFactoryUrl; - } - - return loadBundledLLGuidance(cachedOptions); -} - -async function loadCacheableWasm(source: LoadBundledLLGuidanceOptions['wasm']) { - const url = toCacheableURL(source); - if (!url) return null; - - try { - return await loadWasmBinary(url); - } catch (error) { - logger.warn('Failed to pre-load llguidance WASM binary:', error); - return null; - } -} - -async function loadCacheableWasmFactory(source: LoadBundledLLGuidanceOptions['wasmFactoryUrl']) { - const url = toCacheableURL(source); - if (!url) return null; - - try { - return await loadWasmFactory(url); - } catch (error) { - logger.warn('Failed to pre-load llguidance WASM factory:', error); - return null; - } -} - -function toCacheableURL(source: unknown) { - if (typeof source === 'string') { - return isBlobURL(source) ? null : new URL(source, globalThis.location?.href).href; - } - if (source instanceof URL) { - return isBlobURL(source.href) ? null : source.href; - } - return null; -} - -function isBlobURL(url: string) { - return url.startsWith('blob:'); -} diff --git a/packages/transformers-llguidance/tests/llguidance-constraint.test.js b/packages/transformers-llguidance/tests/llguidance-constraint.test.js index b12f66259..4beb3b117 100644 --- a/packages/transformers-llguidance/tests/llguidance-constraint.test.js +++ b/packages/transformers-llguidance/tests/llguidance-constraint.test.js @@ -2,8 +2,14 @@ import { jest } from "@jest/globals"; import { Tensor } from "@huggingface/transformers"; const computeMask = jest.fn(); +const computeMaskInto = jest.fn(); const commitToken = jest.fn(); -const createInterpreter = jest.fn(() => ({ computeMask, commitToken })); +let supportsComputeMaskInto = false; +const createInterpreter = jest.fn(() => ({ + computeMask, + ...(supportsComputeMaskInto ? { computeMaskInto } : {}), + commitToken, +})); const loadBundledLLGuidance = jest.fn(async () => ({ createInterpreter })); jest.unstable_mockModule("llguidance", () => ({ @@ -15,6 +21,8 @@ const { LlguidanceConstraint } = await import("../dist/index.js"); describe("LlguidanceConstraint", () => { beforeEach(() => { computeMask.mockReset(); + computeMaskInto.mockReset(); + supportsComputeMaskInto = false; commitToken.mockReset(); createInterpreter.mockClear(); loadBundledLLGuidance.mockClear(); @@ -32,15 +40,28 @@ describe("LlguidanceConstraint", () => { logits_processor([[0n], [0n]], logits); expect(loadBundledLLGuidance).toHaveBeenCalledTimes(1); + expect(loadBundledLLGuidance).toHaveBeenCalledWith(); expect(createInterpreter).toHaveBeenCalledWith({ tokenizer, response_format }); expect(Array.from(logits.data)).toEqual([1, -Infinity, 3, -Infinity, 5, -Infinity, 7, -Infinity]); }); - // Ensures caller-provided load options are forwarded without the cache-only control flag. - it("passes explicit llguidance load options", async () => { - await LlguidanceConstraint.fromResponseFormat({}, { type: "json_object" }, { useWasmCache: false, wasmUrl: "custom.wasm" }); + it("reuses a packed mask buffer when the runtime supports computeMaskInto", async () => { + supportsComputeMaskInto = true; + computeMaskInto.mockImplementation((target) => { + target[0] = 0b0101; + return { mask: target, vocabSize: 4 }; + }); - expect(loadBundledLLGuidance).toHaveBeenCalledWith({ wasmUrl: "custom.wasm" }); + const { logits_processor } = await LlguidanceConstraint.fromResponseFormat({}, { type: "json_object" }); + const first = new Tensor("float32", new Float32Array([1, 2, 3, 4]), [1, 4]); + const second = new Tensor("float32", new Float32Array([5, 6, 7, 8]), [1, 4]); + logits_processor([[0n]], first); + logits_processor([[0n]], second); + + expect(computeMask).not.toHaveBeenCalled(); + expect(computeMaskInto).toHaveBeenCalledTimes(2); + expect(computeMaskInto.mock.calls[0][0]).toBe(computeMaskInto.mock.calls[1][0]); + expect(Array.from(second.data)).toEqual([5, -Infinity, 7, -Infinity]); }); // Confirms sampled tokens are committed back to llguidance so stopping criteria can end generation. @@ -97,21 +118,6 @@ describe("LlguidanceConstraint", () => { expect(stopping_criteria([[0n]])).toEqual([true]); }); - // Handles llguidance's post-stop compute error as a normal completion signal. - it("treats compute_mask after stop as completed", async () => { - computeMask.mockImplementation(() => { - throw new Error("compute_mask() called after stop"); - }); - - const { logits_processor, stopping_criteria } = await LlguidanceConstraint.fromResponseFormat({}, { type: "json_object" }); - const logits = new Tensor("float32", new Float32Array([1, 2]), [1, 2]); - - logits_processor([[0n]], logits); - - expect(Array.from(logits.data)).toEqual([1, 2]); - expect(stopping_criteria([[0n]])).toEqual([true]); - }); - // Keeps unexpected interpreter failures visible instead of swallowing real bugs. it("rethrows unexpected computeMask errors", async () => { const error = new Error("unexpected"); diff --git a/packages/transformers-llguidance/tests/wasm-loading.test.js b/packages/transformers-llguidance/tests/wasm-loading.test.js deleted file mode 100644 index c218a5228..000000000 --- a/packages/transformers-llguidance/tests/wasm-loading.test.js +++ /dev/null @@ -1,123 +0,0 @@ -import { jest } from "@jest/globals"; - -const createInterpreter = jest.fn(() => ({ - computeMask: jest.fn(() => ({ stop: true })), - commitToken: jest.fn(), -})); -const loadBundledLLGuidance = jest.fn(async () => ({ createInterpreter })); -const loadWasmBinary = jest.fn(); -const loadWasmFactory = jest.fn(); -const logger = { - debug: jest.fn(), - warn: jest.fn(), -}; - -class LogitsProcessor {} -class LogitsProcessorList extends Array {} -class StoppingCriteria {} - -jest.unstable_mockModule("@huggingface/transformers", () => ({ - LogitsProcessor, - LogitsProcessorList, - StoppingCriteria, - env: { useWasmCache: true }, - loadWasmBinary, - loadWasmFactory, - logger, -})); - -jest.unstable_mockModule("llguidance", () => ({ - loadBundledLLGuidance, -})); - -const originalProcess = globalThis.process; -const originalLocationDescriptor = Object.getOwnPropertyDescriptor(globalThis, "location"); -const { LlguidanceConstraint } = await import("../dist/index.js"); - -describe("llguidance WASM loading", () => { - beforeEach(() => { - delete globalThis.process; - Object.defineProperty(globalThis, "location", { - configurable: true, - value: { href: "https://example.test/app/" }, - }); - - createInterpreter.mockClear(); - loadBundledLLGuidance.mockClear(); - loadWasmBinary.mockReset(); - loadWasmFactory.mockReset(); - logger.debug.mockClear(); - logger.warn.mockClear(); - }); - - afterEach(() => { - globalThis.process = originalProcess; - if (originalLocationDescriptor) { - Object.defineProperty(globalThis, "location", originalLocationDescriptor); - } else { - delete globalThis.location; - } - }); - - // Exercises browser-like cache preloading with the package's default CDN asset URLs. - it("preloads default WASM assets when cache is enabled outside Node", async () => { - const wasm = new Uint8Array([1, 2, 3]); - loadWasmBinary.mockResolvedValue(wasm); - loadWasmFactory.mockResolvedValue("blob:factory-url"); - - await LlguidanceConstraint.fromResponseFormat({}, { type: "json_object" }); - - expect(loadWasmBinary).toHaveBeenCalledWith("https://cdn.jsdelivr.net/npm/llguidance@0.1.8/wasm/llguidance_wasm_bg.wasm"); - expect(loadWasmFactory).toHaveBeenCalledWith("https://cdn.jsdelivr.net/npm/llguidance@0.1.8/wasm/llguidance_wasm.js"); - expect(loadBundledLLGuidance).toHaveBeenCalledWith({ - wasm, - wasmFactoryUrl: "blob:factory-url", - }); - }); - - // Verifies custom relative and absolute URLs are normalized before using the transformers cache helpers. - it("resolves custom cacheable WASM URLs", async () => { - const wasm = new Uint8Array([4, 5, 6]); - loadWasmBinary.mockResolvedValue(wasm); - loadWasmFactory.mockResolvedValue("blob:custom-factory-url"); - - await LlguidanceConstraint.fromResponseFormat( - {}, - { type: "json_object" }, - { - wasmUrl: "assets/custom.wasm", - wasmFactoryUrl: new URL("https://cdn.test/custom-factory.js"), - }, - ); - - expect(loadWasmBinary).toHaveBeenCalledWith("https://example.test/app/assets/custom.wasm"); - expect(loadWasmFactory).toHaveBeenCalledWith("https://cdn.test/custom-factory.js"); - expect(loadBundledLLGuidance).toHaveBeenCalledWith({ - wasm, - wasmFactoryUrl: "blob:custom-factory-url", - }); - }); - - // Ensures a failed preload does not prevent llguidance from loading with the original options. - it("falls back to uncached options when WASM preload fails", async () => { - const error = new Error("network failure"); - loadWasmBinary.mockRejectedValue(error); - loadWasmFactory.mockResolvedValue("blob:factory-url"); - - await LlguidanceConstraint.fromResponseFormat({}, { type: "json_object" }, { wasmUrl: "custom.wasm" }); - - expect(logger.warn).toHaveBeenCalledWith("Failed to pre-load llguidance WASM binary:", error); - expect(loadBundledLLGuidance).toHaveBeenCalledWith({ wasmUrl: "custom.wasm" }); - }); - - // Avoids cache preloading when the caller already supplied an initialized factory. - it("skips preloading when a WASM factory is provided", async () => { - const wasmFactory = jest.fn(); - - await LlguidanceConstraint.fromResponseFormat({}, { type: "json_object" }, { wasmFactory }); - - expect(loadWasmBinary).not.toHaveBeenCalled(); - expect(loadWasmFactory).not.toHaveBeenCalled(); - expect(loadBundledLLGuidance).toHaveBeenCalledWith({ wasmFactory }); - }); -}); diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 379ccfc6b..5bb74e840 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -61,8 +61,8 @@ importers: packages/transformers-llguidance: dependencies: llguidance: - specifier: ^0.1.8 - version: 0.1.8 + specifier: 0.2.0 + version: 0.2.0 devDependencies: '@huggingface/transformers': specifier: workspace:* @@ -1585,8 +1585,8 @@ packages: linkify-it@5.0.0: resolution: {integrity: sha512-5aHCbzQRADcdP+ATqnDuhhJ/MRIqDkZX5pyjFHRRysS8vZ5AbqGEoFIb6pYHPZ+L/OC2Lc+xT8uHVVR5CAK/wQ==} - llguidance@0.1.8: - resolution: {integrity: sha512-n44NMtq2ChgUGkcy6G7o8TUp90kyK2TergBDZ2nvxUN87vtkmCm/WxEWG0hl3Pav1BhqxybKAR6DP+9oib0Law==} + llguidance@0.2.0: + resolution: {integrity: sha512-BvN0WlQLLZC9kkRpz7k4DoAVpplcaJKcKrRVI68tCwPZlWtrGv/kriFyjKwP7y0qBct4M7vmOWEqXJTDFsMB6w==} locate-path@5.0.0: resolution: {integrity: sha512-t7hw9pI+WvuwNJXwk5zVHpyhIqzg2qTlklJOf0mVxGSbe3Fp2VieZcduNYjaLDoy6p9uGpQEGWG87WpMKlNq8g==} @@ -3750,7 +3750,7 @@ snapshots: dependencies: uc.micro: 2.1.0 - llguidance@0.1.8: {} + llguidance@0.2.0: {} locate-path@5.0.0: dependencies: From 4d399a7ae6e4f9b3fe4fa9cee9c60e4d6b50dc2c Mon Sep 17 00:00:00 2001 From: Nico Martin Date: Fri, 31 Jul 2026 11:07:43 +0200 Subject: [PATCH 12/17] clean up --- packages/transformers-llguidance/README.md | 64 ++++--- .../src/LlguidanceConstraint.ts | 168 +++++++++++++----- .../transformers-llguidance/src/utils/mask.ts | 112 ++++++------ .../src/utils/types.ts | 31 +--- .../tests/llguidance-constraint.test.js | 150 ++++++++++------ packages/transformers/src/transformers.js | 1 - 6 files changed, 317 insertions(+), 209 deletions(-) diff --git a/packages/transformers-llguidance/README.md b/packages/transformers-llguidance/README.md index e167b245f..55b1e2fee 100644 --- a/packages/transformers-llguidance/README.md +++ b/packages/transformers-llguidance/README.md @@ -7,26 +7,35 @@ This package exports `LlguidanceConstraint`, which turns an llguidance response ```js import { LlguidanceConstraint } from "@huggingface/transformers-llguidance"; -const { logits_processor, stopping_criteria } = - await LlguidanceConstraint.fromResponseFormat(tokenizer, { - type: "json_schema", - json_schema: { - type: "object", - properties: { - answer: { type: "string" }, - }, - required: ["answer"], - additionalProperties: false, +const constraint = await LlguidanceConstraint.fromResponseFormat(tokenizer, { + type: "json_schema", + json_schema: { + type: "object", + properties: { + answer: { type: "string" }, }, - }); - -await model.generate({ - ...inputs, - logits_processor, - stopping_criteria, + required: ["answer"], + additionalProperties: false, + }, }); + +try { + await model.generate({ + ...inputs, + logits_processor: constraint.logits_processor, + stopping_criteria: constraint.stopping_criteria, + }); +} finally { + constraint.dispose(); +} ``` +Constraints currently support a single generated sequence at a time. Generation throws when the logical batch size is not `1` rather than sharing mutable grammar state across sequences. + +The interpreter is automatically disposed when llguidance reaches a terminal state. Call `dispose()` in a `finally` block as shown above to also release resources when generation ends for another reason, such as `max_new_tokens` or cancellation. + +If llguidance reports acceptance before sampling, the constraint forces the tokenizer's EOS token so no unconstrained token is appended. Compatible tokenizer objects must expose an EOS ID such as `eos_token_id` or `eosTokenId`; generation fails closed if acceptance occurs without one. + ## Regex constraints Use `type: "regex"` to constrain generation to a regular expression. For example, this only allows ISO-like dates in `YYYY-MM-DD` format: @@ -34,15 +43,18 @@ Use `type: "regex"` to constrain generation to a regular expression. For example ```js import { LlguidanceConstraint } from "@huggingface/transformers-llguidance"; -const { logits_processor, stopping_criteria } = - await LlguidanceConstraint.fromResponseFormat(tokenizer, { - type: "regex", - regex: "\\d{4}-\\d{2}-\\d{2}", - }); - -const output = await model.generate({ - ...inputs, - logits_processor, - stopping_criteria, +const constraint = await LlguidanceConstraint.fromResponseFormat(tokenizer, { + type: "regex", + regex: "\\d{4}-\\d{2}-\\d{2}", }); + +try { + const output = await model.generate({ + ...inputs, + logits_processor: constraint.logits_processor, + stopping_criteria: constraint.stopping_criteria, + }); +} finally { + constraint.dispose(); +} ``` diff --git a/packages/transformers-llguidance/src/LlguidanceConstraint.ts b/packages/transformers-llguidance/src/LlguidanceConstraint.ts index 4c917f73a..9fb0e560d 100644 --- a/packages/transformers-llguidance/src/LlguidanceConstraint.ts +++ b/packages/transformers-llguidance/src/LlguidanceConstraint.ts @@ -6,15 +6,15 @@ import { logger, type Tensor, } from '@huggingface/transformers'; -import { loadBundledLLGuidance, type JSONSchema, type LLGuidanceTokenizerSource } from 'llguidance'; - -import { applyMask, forceToken, summarizeMaskResult } from './utils/mask'; import { - type GuidanceInterpreter, - type GuidanceMaskResult, - type LlguidanceState, - type LlguidanceStats, -} from './utils/types'; + loadBundledLLGuidance, + type JSONSchema, + type LLGuidanceMaskResult, + type LLGuidanceTokenizerSource, +} from 'llguidance'; + +import { applyMask, forceTokens, summarizeMaskResult } from './utils/mask'; +import { type LlguidanceState, type LlguidanceStats } from './utils/types'; export type ResponseFormat = | { type: 'json_object' } @@ -29,15 +29,18 @@ export class LlguidanceConstraint { }); const runtime = await loadBundledLLGuidance(); + const eosTokenIds = getEosTokenIds(tokenizer); const interpreter = runtime.createInterpreter({ tokenizer, response_format, - }) as GuidanceInterpreter; + }); logger.debug('[LlguidanceConstraint] interpreter created'); const state: LlguidanceState = { completed: false, + disposed: false, + eosTokenIds, interpreter, step: 0, stats: { @@ -58,6 +61,7 @@ export class LlguidanceConstraint { logits_processor, stopping_criteria: new LlguidanceStoppingCriteria(state), stats: state.stats, + dispose: () => disposeState(state), }; } } @@ -71,6 +75,8 @@ class LlguidanceLogitsProcessor extends LogitsProcessor { } _call(_inputIds: bigint[][], logits: Tensor) { + assertSingleSequence(_inputIds.length, this.state); + if (this.state.completed) { if (isDebugEnabled()) { logger.debug('[LlguidanceLogitsProcessor] skip completed', { @@ -79,10 +85,16 @@ class LlguidanceLogitsProcessor extends LogitsProcessor { } return logits; } + if (this.state.disposed) { + throw new Error('LlguidanceConstraint has been disposed.'); + } this.state.step++; this.state.stats.steps = this.state.step; const vocabSize = logits.dims.at(-1); + if (vocabSize === undefined || !Number.isInteger(vocabSize) || vocabSize <= 0) { + throw failState(this.state, 'LlguidanceConstraint requires logits with a vocabulary dimension.'); + } if (isDebugEnabled()) { logger.debug('[LlguidanceLogitsProcessor] compute mask', { step: this.state.step, @@ -91,14 +103,17 @@ class LlguidanceLogitsProcessor extends LogitsProcessor { }); } - let result: GuidanceMaskResult; + let result: LLGuidanceMaskResult; const maskStart = performance.now(); - const maskWords = vocabSize === undefined ? undefined : Math.ceil(vocabSize / 32); - if (this.state.interpreter.computeMaskInto && maskWords !== undefined) { - this.state.maskBuffer ??= new Uint32Array(maskWords); + const maskWords = Math.ceil(vocabSize / 32); + if (this.state.maskBuffer?.length !== maskWords) { + this.state.maskBuffer = new Uint32Array(maskWords); + } + try { result = this.state.interpreter.computeMaskInto(this.state.maskBuffer); - } else { - result = this.state.interpreter.computeMask(); + } catch (error) { + disposeState(this.state); + throw error; } this.state.stats.computeMaskMs += performance.now() - maskStart; const maskInstrumentation = this.state.interpreter.maskInstrumentation?.(); @@ -111,29 +126,26 @@ class LlguidanceLogitsProcessor extends LogitsProcessor { if (isDebugEnabled()) { logger.debug('[LlguidanceLogitsProcessor] mask result', { step: this.state.step, - result: summarizeMaskResult(result, vocabSize), + result: summarizeMaskResult(result), }); } if ('stop' in result && result.stop) { + this.state.stats.stopReason = result.reason; + if (result.reason === 'dead_end') { + throw failState(this.state, 'llguidance reached a dead end before satisfying the constraint.'); + } this.state.completed = true; - if (isDebugEnabled()) { - logger.debug('[LlguidanceLogitsProcessor] stopped by mask', { - step: this.state.step, - }); + try { + forceTokens(logits, this.state.eosTokenIds); + } catch (error) { + disposeState(this.state); + throw error; } - return logits; - } - - if ('ffTokens' in result && result.ffTokens?.length) { - // Fast-forward splice: the grammar forces the next token, so ban - // everything else instead of letting the model sample unconstrained. - forceToken(logits, result.ffTokens[0], vocabSize); + disposeState(this.state); if (isDebugEnabled()) { - logger.debug('[LlguidanceLogitsProcessor] forced splice token', { + logger.debug('[LlguidanceLogitsProcessor] stopped by mask', { step: this.state.step, - ffTokens: Array.from(result.ffTokens), - backtrack: 'backtrack' in result ? result.backtrack : undefined, }); } return logits; @@ -141,14 +153,19 @@ class LlguidanceLogitsProcessor extends LogitsProcessor { if ('mask' in result) { const applyStart = performance.now(); - if (isDebugEnabled()) { - const applied = applyMask(logits, result.mask, result.vocabSize ?? vocabSize, true); - logger.debug('[LlguidanceLogitsProcessor] mask applied', { - step: this.state.step, - ...applied, - }); - } else { - applyMask(logits, result.mask, result.vocabSize ?? vocabSize); + try { + if (isDebugEnabled()) { + const applied = applyMask(logits, result.mask, result.vocabSize, true); + logger.debug('[LlguidanceLogitsProcessor] mask applied', { + step: this.state.step, + ...applied, + }); + } else { + applyMask(logits, result.mask, result.vocabSize); + } + } catch (error) { + disposeState(this.state); + throw error; } this.state.stats.applyMaskMs += performance.now() - applyStart; } @@ -168,9 +185,18 @@ class LlguidanceLogitsProcessor extends LogitsProcessor { } if (this.state.completed) return; + if (this.state.disposed) { + throw new Error('LlguidanceConstraint has been disposed.'); + } const commitStart = performance.now(); - const result = this.state.interpreter.commitToken(tokenId); + let result; + try { + result = this.state.interpreter.commitToken(tokenId); + } catch (error) { + disposeState(this.state); + throw error; + } this.state.stats.commitTokenMs += performance.now() - commitStart; if (isDebugEnabled()) { logger.debug('[LlguidanceLogitsProcessor] token committed', { @@ -180,8 +206,17 @@ class LlguidanceLogitsProcessor extends LogitsProcessor { }); } - if (result?.stop) { + if (result.backtrack > 0) { + throw failState( + this.state, + `llguidance requested backtracking by ${result.backtrack} token(s), which Transformers.js does not support.`, + ); + } + + if (result.stop) { this.state.completed = true; + this.state.stats.stopReason = 'accepted'; + disposeState(this.state); if (isDebugEnabled()) { logger.debug('[LlguidanceLogitsProcessor] stopped by commit', { step: this.state.step, @@ -192,6 +227,8 @@ class LlguidanceLogitsProcessor extends LogitsProcessor { } onTokensSampled(tokenIds: number[], inputIds: bigint[][]) { + assertSingleSequence(tokenIds.length, this.state); + assertSingleSequence(inputIds.length, this.state); if (isDebugEnabled()) { logger.debug('[LlguidanceLogitsProcessor] tokens sampled', { step: this.state.step, @@ -201,9 +238,7 @@ class LlguidanceLogitsProcessor extends LogitsProcessor { }); } - for (let batchIdx = 0; batchIdx < tokenIds.length; ++batchIdx) { - this.onTokenSampled(tokenIds[batchIdx], batchIdx, inputIds); - } + this.onTokenSampled(tokenIds[0], 0, inputIds); } } @@ -216,6 +251,7 @@ class LlguidanceStoppingCriteria extends StoppingCriteria { } _call(inputIds: ArrayLike[]) { + assertSingleSequence(inputIds.length, this.state); const result = new Array(inputIds.length).fill(this.state.completed); if (isDebugEnabled()) { logger.debug('[LlguidanceStoppingCriteria] call', { @@ -231,5 +267,49 @@ class LlguidanceStoppingCriteria extends StoppingCriteria { function isDebugEnabled() { // Keep compatibility with Transformers.js releases that predate the LogLevel export. - return env.logLevel <= 10; + return env.logLevel <= DEBUG_LOG_LEVEL; +} + +const DEBUG_LOG_LEVEL = 10; + +function assertSingleSequence(batchSize: number, state: LlguidanceState) { + if (batchSize !== 1) { + throw failState(state, `LlguidanceConstraint currently supports batch size 1; received ${batchSize}.`); + } +} + +function failState(state: LlguidanceState, message: string) { + disposeState(state); + return new Error(message); +} + +function disposeState(state: LlguidanceState) { + if (state.disposed) return; + state.disposed = true; + state.maskBuffer = undefined; + state.interpreter.dispose(); +} + +function getEosTokenIds(tokenizer: LLGuidanceTokenizerSource) { + const source = tokenizer as Record; + const values = [ + source.eos_token_id, + source.eosTokenId, + source.eos_token_ids, + source.eosTokenIds, + source.eos_token, + source.eos_tokens, + ]; + + if (typeof source.eosToken === 'function') { + values.push(source.eosToken.call(tokenizer)); + } + + return [ + ...new Set( + values + .flat() + .filter((value): value is number => typeof value === 'number' && Number.isInteger(value) && value >= 0), + ), + ]; } diff --git a/packages/transformers-llguidance/src/utils/mask.ts b/packages/transformers-llguidance/src/utils/mask.ts index a63136c04..66706999a 100644 --- a/packages/transformers-llguidance/src/utils/mask.ts +++ b/packages/transformers-llguidance/src/utils/mask.ts @@ -1,45 +1,49 @@ import { type Tensor } from '@huggingface/transformers'; -import { type GuidanceMask, type GuidanceMaskResult } from './types'; +import type { LLGuidanceMaskResult } from 'llguidance'; type LogitsData = Float32Array | Float64Array | number[]; -export function summarizeMaskResult(result: GuidanceMaskResult, vocabSize?: number) { - if ('stop' in result && result.stop) { - return { stop: true }; - } - +export function summarizeMaskResult(result: LLGuidanceMaskResult): + | { stop: true; reason: 'accepted' | 'dead_end' } + | { + maskLength: number; + vocabSize: number; + allowed: number; + sampleAllowedTokenIds: number[]; + } { if (!('mask' in result)) { - return result; + return { stop: true, reason: result.reason }; } return { maskLength: result.mask.length, - vocabSize: result.vocabSize ?? vocabSize, - allowed: countAllowed(result.mask, result.vocabSize ?? vocabSize), - sampleAllowedTokenIds: sampleAllowedTokenIds(result.mask, result.vocabSize ?? vocabSize), + vocabSize: result.vocabSize, + allowed: countAllowed(result.mask, result.vocabSize), + sampleAllowedTokenIds: sampleAllowedTokenIds(result.mask, result.vocabSize), }; } -export function applyMask(logits: Tensor, mask: GuidanceMask, vocabSize?: number, includeSummary = false) { - if (!vocabSize) { - return includeSummary ? { vocabSize, batchSize: 0, masked: 0, allowed: undefined } : undefined; +export function applyMask(logits: Tensor, mask: Uint32Array, vocabSize: number, includeSummary = false) { + if (!Number.isInteger(vocabSize) || vocabSize <= 0) { + throw new Error(`llguidance returned an invalid vocabulary size: ${vocabSize}.`); } const data = logits.data as LogitsData; - const stride = (logits.dims?.at?.(-1) as number) || vocabSize; - const bound = Math.min(vocabSize, stride); - const batchSize = Math.max(1, Math.floor(data.length / stride)); - const packed = mask.length < vocabSize; + const stride = logits.dims.at(-1); + if (stride === undefined || !Number.isInteger(stride) || stride <= 0) { + throw new Error('LlguidanceConstraint requires logits with a vocabulary dimension.'); + } + if (vocabSize > stride) { + throw new Error(`llguidance vocabulary size ${vocabSize} exceeds logits vocabulary size ${stride}.`); + } + if (mask.length < Math.ceil(vocabSize / 32)) { + throw new Error(`llguidance returned a mask that is too short for vocabulary size ${vocabSize}.`); + } + const batchSize = Math.floor(data.length / stride); for (let batch = 0; batch < batchSize; ++batch) { const offset = batch * stride; - if (packed) { - applyPackedMask(data, mask, offset, bound); - } else { - for (let tokenId = 0; tokenId < bound; ++tokenId) { - if (!mask[tokenId]) data[offset + tokenId] = -Infinity; - } - } + applyPackedMask(data, mask, offset, vocabSize); if (stride > vocabSize) { // Logits padded beyond the grammar vocab can never be committed. data.fill(-Infinity, offset + vocabSize, offset + stride); @@ -48,33 +52,40 @@ export function applyMask(logits: Tensor, mask: GuidanceMask, vocabSize?: number if (!includeSummary) return undefined; - const allowed = countAllowed(mask, vocabSize) ?? 0; + const allowed = countAllowed(mask, vocabSize); return { vocabSize, batchSize, masked: (vocabSize - allowed) * batchSize, allowed }; } -// Forces a single token by banning everything else, e.g. for fast-forward splices. -export function forceToken(logits: Tensor, tokenId: number, vocabSize?: number) { - if (!vocabSize) return; - +export function forceTokens(logits: Tensor, tokenIds: number[]) { const data = logits.data as LogitsData; - const stride = (logits.dims?.at?.(-1) as number) || vocabSize; - const batchSize = Math.max(1, Math.floor(data.length / stride)); + const vocabSize = logits.dims.at(-1); + if (vocabSize === undefined || !Number.isInteger(vocabSize) || vocabSize <= 0) { + throw new Error('LlguidanceConstraint requires logits with a vocabulary dimension.'); + } + if (tokenIds.length === 0) { + throw new Error('LlguidanceConstraint cannot stop on acceptance because the tokenizer has no EOS token ID.'); + } - for (let batch = 0; batch < batchSize; ++batch) { - const offset = batch * stride; - const kept = data[offset + tokenId]; - data.fill(-Infinity, offset, offset + stride); - data[offset + tokenId] = Number.isFinite(kept) ? kept : 0; + const kept = tokenIds.map((tokenId) => { + if (!Number.isInteger(tokenId) || tokenId < 0 || tokenId >= vocabSize) { + throw new Error(`Tokenizer EOS token ID ${tokenId} is outside logits vocabulary size ${vocabSize}.`); + } + return data[tokenId]; + }); + + data.fill(-Infinity); + for (let i = 0; i < tokenIds.length; ++i) { + data[tokenIds[i]] = Number.isFinite(kept[i]) ? kept[i] : 0; } } // Hot path: runs once per generated token over the whole vocab. Grammar masks -// are skewed — most 32-token words are either fully allowed (skip) or fully -// banned (memset) — so per-bit work only happens on the few mixed words. -function applyPackedMask(data: LogitsData, mask: GuidanceMask, offset: number, vocabSize: number) { +// are skewed: most 32-token words are either fully allowed (skip) or fully +// banned (memset), so per-bit work only happens on the few mixed words. +function applyPackedMask(data: LogitsData, mask: Uint32Array, offset: number, vocabSize: number) { const numWords = vocabSize >>> 5; for (let word = 0; word < numWords; ++word) { - const bits = (mask[word] as number) | 0; + const bits = mask[word] | 0; if (bits === -1) continue; const base = offset + (word << 5); if (bits === 0) { @@ -87,35 +98,28 @@ function applyPackedMask(data: LogitsData, mask: GuidanceMask, offset: number, v } for (let tokenId = numWords << 5; tokenId < vocabSize; ++tokenId) { - if (!((mask[tokenId >>> 5] as number) & (1 << (tokenId & 31)))) { + if (!(mask[tokenId >>> 5] & (1 << (tokenId & 31)))) { data[offset + tokenId] = -Infinity; } } } -function countAllowed(mask: GuidanceMask, vocabSize?: number) { - if (!vocabSize) return undefined; - +function countAllowed(mask: Uint32Array, vocabSize: number) { let allowed = 0; for (let tokenId = 0; tokenId < vocabSize; ++tokenId) { - if (isAllowed(mask, tokenId, vocabSize)) allowed++; + if (isAllowed(mask, tokenId)) allowed++; } return allowed; } -function sampleAllowedTokenIds(mask: GuidanceMask, vocabSize?: number) { - if (!vocabSize) return []; - +function sampleAllowedTokenIds(mask: Uint32Array, vocabSize: number) { const tokenIds: number[] = []; for (let tokenId = 0; tokenId < vocabSize && tokenIds.length < 25; ++tokenId) { - if (isAllowed(mask, tokenId, vocabSize)) tokenIds.push(tokenId); + if (isAllowed(mask, tokenId)) tokenIds.push(tokenId); } return tokenIds; } -function isAllowed(mask: GuidanceMask, tokenId: number, vocabSize: number) { - if (mask.length >= vocabSize) { - return Boolean(mask[tokenId]); - } - return Boolean(Number(mask[tokenId >> 5]) & (1 << (tokenId & 31))); +function isAllowed(mask: Uint32Array, tokenId: number) { + return Boolean(mask[tokenId >> 5] & (1 << (tokenId & 31))); } diff --git a/packages/transformers-llguidance/src/utils/types.ts b/packages/transformers-llguidance/src/utils/types.ts index 325dc69b3..b5c3c1470 100644 --- a/packages/transformers-llguidance/src/utils/types.ts +++ b/packages/transformers-llguidance/src/utils/types.ts @@ -1,28 +1,4 @@ -export type GuidanceMask = Uint32Array | Uint8Array | boolean[] | number[]; - -export type GuidanceMaskResult = - | { mask: GuidanceMask; vocabSize?: number; stop?: false } - | { stop: true } - | { backtrack?: number; ffTokens?: number[] | Uint32Array }; - -export type GuidanceCommitResult = { - stop?: boolean; - backtrack?: number; - ffTokens?: number[] | Uint32Array; -}; - -export type GuidanceMaskInstrumentation = { - trieNodesVisited: number; - sharedJsonMaskCacheHits: number; - sharedJsonMaskCacheMisses: number; -}; - -export type GuidanceInterpreter = { - computeMask(): GuidanceMaskResult; - computeMaskInto?(target: Uint32Array): GuidanceMaskResult; - maskInstrumentation?(): GuidanceMaskInstrumentation; - commitToken(tokenId: number): GuidanceCommitResult | undefined; -}; +import type { LLGuidanceInterpreterHandle } from 'llguidance'; export type LlguidanceStats = { steps: number; @@ -32,11 +8,14 @@ export type LlguidanceStats = { trieNodesVisited: number; sharedJsonMaskCacheHits: number; sharedJsonMaskCacheMisses: number; + stopReason?: 'accepted' | 'dead_end'; }; export type LlguidanceState = { completed: boolean; - interpreter: GuidanceInterpreter; + disposed: boolean; + eosTokenIds: number[]; + interpreter: LLGuidanceInterpreterHandle; maskBuffer?: Uint32Array; step: number; stats: LlguidanceStats; diff --git a/packages/transformers-llguidance/tests/llguidance-constraint.test.js b/packages/transformers-llguidance/tests/llguidance-constraint.test.js index 4beb3b117..ceda77c84 100644 --- a/packages/transformers-llguidance/tests/llguidance-constraint.test.js +++ b/packages/transformers-llguidance/tests/llguidance-constraint.test.js @@ -4,11 +4,12 @@ import { Tensor } from "@huggingface/transformers"; const computeMask = jest.fn(); const computeMaskInto = jest.fn(); const commitToken = jest.fn(); -let supportsComputeMaskInto = false; +const disposeInterpreter = jest.fn(); const createInterpreter = jest.fn(() => ({ computeMask, - ...(supportsComputeMaskInto ? { computeMaskInto } : {}), + computeMaskInto, commitToken, + dispose: disposeInterpreter, })); const loadBundledLLGuidance = jest.fn(async () => ({ createInterpreter })); @@ -22,112 +23,145 @@ describe("LlguidanceConstraint", () => { beforeEach(() => { computeMask.mockReset(); computeMaskInto.mockReset(); - supportsComputeMaskInto = false; commitToken.mockReset(); + disposeInterpreter.mockReset(); createInterpreter.mockClear(); loadBundledLLGuidance.mockClear(); - }); - // Verifies the core integration path: llguidance creates an interpreter and its mask is applied across batches. - it("loads llguidance and applies masks", async () => { - computeMask.mockReturnValue({ mask: [true, false, true, false], vocabSize: 4 }); + computeMaskInto.mockImplementation((target) => { + target[0] = 0b0101; + return { mask: target, vocabSize: 4 }; + }); + commitToken.mockReturnValue({ stop: false, backtrack: 0, ffTokens: [] }); + }); + it("loads llguidance and applies a packed mask", async () => { const tokenizer = { name: "tokenizer" }; const response_format = { type: "json_schema", json_schema: { type: "object" } }; const { logits_processor } = await LlguidanceConstraint.fromResponseFormat(tokenizer, response_format); - const logits = new Tensor("float32", new Float32Array([1, 2, 3, 4, 5, 6, 7, 8]), [2, 4]); + const logits = new Tensor("float32", new Float32Array([1, 2, 3, 4]), [1, 4]); - logits_processor([[0n], [0n]], logits); + logits_processor([[0n]], logits); expect(loadBundledLLGuidance).toHaveBeenCalledTimes(1); - expect(loadBundledLLGuidance).toHaveBeenCalledWith(); expect(createInterpreter).toHaveBeenCalledWith({ tokenizer, response_format }); - expect(Array.from(logits.data)).toEqual([1, -Infinity, 3, -Infinity, 5, -Infinity, 7, -Infinity]); + expect(Array.from(logits.data)).toEqual([1, -Infinity, 3, -Infinity]); }); - it("reuses a packed mask buffer when the runtime supports computeMaskInto", async () => { - supportsComputeMaskInto = true; - computeMaskInto.mockImplementation((target) => { - target[0] = 0b0101; - return { mask: target, vocabSize: 4 }; - }); - + it("reuses the packed mask buffer", async () => { const { logits_processor } = await LlguidanceConstraint.fromResponseFormat({}, { type: "json_object" }); const first = new Tensor("float32", new Float32Array([1, 2, 3, 4]), [1, 4]); const second = new Tensor("float32", new Float32Array([5, 6, 7, 8]), [1, 4]); + logits_processor([[0n]], first); logits_processor([[0n]], second); expect(computeMask).not.toHaveBeenCalled(); expect(computeMaskInto).toHaveBeenCalledTimes(2); expect(computeMaskInto.mock.calls[0][0]).toBe(computeMaskInto.mock.calls[1][0]); - expect(Array.from(second.data)).toEqual([5, -Infinity, 7, -Infinity]); }); - // Confirms sampled tokens are committed back to llguidance so stopping criteria can end generation. - it("commits sampled tokens and stops when llguidance stops", async () => { - computeMask.mockReturnValue({ mask: [true, true], vocabSize: 2 }); - commitToken.mockReturnValueOnce(undefined).mockReturnValueOnce({ stop: true }); + it("rejects batched generation before computing a mask", async () => { + const { logits_processor } = await LlguidanceConstraint.fromResponseFormat({}, { type: "json_object" }); + const logits = new Tensor("float32", new Float32Array(8), [2, 4]); - const { logits_processor, stopping_criteria } = await LlguidanceConstraint.fromResponseFormat({}, { type: "json_object" }); + expect(() => logits_processor([[0n], [0n]], logits)).toThrow("currently supports batch size 1"); + expect(computeMaskInto).not.toHaveBeenCalled(); + expect(disposeInterpreter).toHaveBeenCalledTimes(1); + }); - expect(stopping_criteria([[0n]])).toEqual([false]); + it("commits one sampled token and stops on acceptance", async () => { + commitToken.mockReturnValue({ stop: true, backtrack: 0, ffTokens: [] }); + const { logits_processor, stopping_criteria, stats } = await LlguidanceConstraint.fromResponseFormat({}, { type: "json_object" }); - logits_processor.onTokensSampled( - [0, 1], - [ - [0n, 0n], - [0n, 1n], - ], - ); + expect(stopping_criteria([[0n]])).toEqual([false]); + logits_processor.onTokensSampled([1], [[0n, 1n]]); - expect(commitToken).toHaveBeenCalledWith(0); expect(commitToken).toHaveBeenCalledWith(1); - expect( - stopping_criteria([ - [0n, 0n], - [0n, 1n], - ]), - ).toEqual([true, true]); + expect(stopping_criteria([[0n, 1n]])).toEqual([true]); + expect(stats.stopReason).toBe("accepted"); + expect(disposeInterpreter).toHaveBeenCalledTimes(1); }); - // Covers llguidance's compact bitset mask format, not just boolean arrays. - it("applies packed uint32 masks", async () => { - computeMask.mockReturnValue({ mask: new Uint32Array([0b0101]), vocabSize: 4 }); - + it("throws when llguidance requests backtracking", async () => { + commitToken.mockReturnValue({ stop: false, backtrack: 1, ffTokens: [] }); const { logits_processor } = await LlguidanceConstraint.fromResponseFormat({}, { type: "json_object" }); - const logits = new Tensor("float32", new Float32Array([1, 2, 3, 4]), [1, 4]); - - logits_processor([[0n]], logits); - expect(Array.from(logits.data)).toEqual([1, -Infinity, 3, -Infinity]); + expect(() => logits_processor.onTokensSampled([1], [[0n, 1n]])).toThrow("requested backtracking by 1 token"); + expect(disposeInterpreter).toHaveBeenCalledTimes(1); }); - // Ensures a stop response from computeMask marks the shared state complete without committing more tokens. - it("stops generation when computeMask returns stop", async () => { - computeMask.mockReturnValue({ stop: true }); + it("records and throws on a dead end", async () => { + computeMaskInto.mockReturnValue({ stop: true, reason: "dead_end" }); + const { logits_processor, stats } = await LlguidanceConstraint.fromResponseFormat({}, { type: "json_object" }); + const logits = new Tensor("float32", new Float32Array(4), [1, 4]); + + expect(() => logits_processor([[0n]], logits)).toThrow("reached a dead end"); + expect(stats.stopReason).toBe("dead_end"); + expect(disposeInterpreter).toHaveBeenCalledTimes(1); + }); - const { logits_processor, stopping_criteria } = await LlguidanceConstraint.fromResponseFormat({}, { type: "json_object" }); - const logits = new Tensor("float32", new Float32Array([1, 2]), [1, 2]); + it("stops and disposes when computeMask reports acceptance", async () => { + computeMaskInto.mockReturnValue({ stop: true, reason: "accepted" }); + const { logits_processor, stopping_criteria, stats } = await LlguidanceConstraint.fromResponseFormat({ eos_token_id: 1 }, { type: "json_object" }); + const logits = new Tensor("float32", new Float32Array([1, 2, 3, 4]), [1, 4]); logits_processor([[0n]], logits); logits_processor.onTokensSampled([0], [[0n]]); - expect(Array.from(logits.data)).toEqual([1, 2]); + expect(Array.from(logits.data)).toEqual([-Infinity, 2, -Infinity, -Infinity]); expect(commitToken).not.toHaveBeenCalled(); expect(stopping_criteria([[0n]])).toEqual([true]); + expect(stats.stopReason).toBe("accepted"); + expect(disposeInterpreter).toHaveBeenCalledTimes(1); + }); + + it("fails closed on mask acceptance when the tokenizer has no EOS token ID", async () => { + computeMaskInto.mockReturnValue({ stop: true, reason: "accepted" }); + const { logits_processor } = await LlguidanceConstraint.fromResponseFormat({}, { type: "json_object" }); + const logits = new Tensor("float32", new Float32Array([1, 2, 3, 4]), [1, 4]); + + expect(() => logits_processor([[0n]], logits)).toThrow("tokenizer has no EOS token ID"); + expect(disposeInterpreter).toHaveBeenCalledTimes(1); + }); + + it("fails closed when logits have no vocabulary dimension", async () => { + const { logits_processor } = await LlguidanceConstraint.fromResponseFormat({}, { type: "json_object" }); + const logits = { dims: [], data: new Float32Array(4) }; + + expect(() => logits_processor([[0n]], logits)).toThrow("requires logits with a vocabulary dimension"); + expect(computeMaskInto).not.toHaveBeenCalled(); + expect(disposeInterpreter).toHaveBeenCalledTimes(1); + }); + + it("fails closed when mask vocabulary metadata is invalid", async () => { + computeMaskInto.mockReturnValue({ mask: new Uint32Array([0b0101]), vocabSize: undefined }); + const { logits_processor } = await LlguidanceConstraint.fromResponseFormat({}, { type: "json_object" }); + const logits = new Tensor("float32", new Float32Array(4), [1, 4]); + + expect(() => logits_processor([[0n]], logits)).toThrow("invalid vocabulary size"); + expect(disposeInterpreter).toHaveBeenCalledTimes(1); }); - // Keeps unexpected interpreter failures visible instead of swallowing real bugs. - it("rethrows unexpected computeMask errors", async () => { + it("exposes idempotent disposal", async () => { + const { dispose, logits_processor } = await LlguidanceConstraint.fromResponseFormat({}, { type: "json_object" }); + + dispose(); + dispose(); + + expect(disposeInterpreter).toHaveBeenCalledTimes(1); + expect(() => logits_processor.onTokensSampled([1], [[0n, 1n]])).toThrow("has been disposed"); + }); + + it("disposes and rethrows interpreter errors", async () => { const error = new Error("unexpected"); - computeMask.mockImplementation(() => { + computeMaskInto.mockImplementation(() => { throw error; }); - const { logits_processor } = await LlguidanceConstraint.fromResponseFormat({}, { type: "json_object" }); - const logits = new Tensor("float32", new Float32Array([1, 2]), [1, 2]); + const logits = new Tensor("float32", new Float32Array(4), [1, 4]); expect(() => logits_processor([[0n]], logits)).toThrow(error); + expect(disposeInterpreter).toHaveBeenCalledTimes(1); }); }); diff --git a/packages/transformers/src/transformers.js b/packages/transformers/src/transformers.js index 7d6168020..2b7a035c2 100644 --- a/packages/transformers/src/transformers.js +++ b/packages/transformers/src/transformers.js @@ -53,7 +53,6 @@ export * from './utils/tensor.js'; export { softmax, log_softmax, dot, cos_sim } from './utils/maths.js'; export { random } from './utils/random.js'; export { logger } from './utils/logger.js'; -export { loadWasmBinary, loadWasmFactory } from './backends/utils/cacheWasm.js'; export { DynamicCache } from './cache_utils.js'; From 6331dd4e621c4c3ce849d750595527c6bb2811df Mon Sep 17 00:00:00 2001 From: Nico Martin Date: Fri, 31 Jul 2026 11:37:16 +0200 Subject: [PATCH 13/17] copilot review --- packages/transformers-llguidance/package.json | 1 + packages/transformers/src/models/modeling_utils.js | 11 +++++------ 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/packages/transformers-llguidance/package.json b/packages/transformers-llguidance/package.json index 66fe660c9..f284a8e96 100644 --- a/packages/transformers-llguidance/package.json +++ b/packages/transformers-llguidance/package.json @@ -21,6 +21,7 @@ "typegen": "tsc --build", "dev": "node scripts/dev.mjs", "build": "node scripts/build.mjs && pnpm typegen", + "pretest": "pnpm build", "test": "node --experimental-vm-modules node_modules/jest/bin/jest.js --verbose" }, "repository": { diff --git a/packages/transformers/src/models/modeling_utils.js b/packages/transformers/src/models/modeling_utils.js index f5caec12c..59c07768f 100644 --- a/packages/transformers/src/models/modeling_utils.js +++ b/packages/transformers/src/models/modeling_utils.js @@ -992,6 +992,8 @@ export class PreTrainedModel extends Callable { /** @type {[bigint][]} */ const generated_input_ids = []; + /** @type {number[]} */ + const sampled_token_ids = []; // const new_kv_cache = [];// NOTE: Only used for beam search when concatenating new kv // Loop over each batch for (let batch_idx = 0; batch_idx < next_tokens_scores.dims.at(0); ++batch_idx) { @@ -999,11 +1001,11 @@ export class PreTrainedModel extends Callable { const sampledTokens = await sampler(logs); for (const [newTokenId, logProb] of sampledTokens) { - const bigint = BigInt(newTokenId); // TODO: If branching, use previous beam as a starting point // update generated ids, model inputs, and length for next step scores[batch_idx] += logProb; - generated_input_ids.push([bigint]); + generated_input_ids.push([newTokenId]); + sampled_token_ids.push(Number(newTokenId)); // TODO: Support beam search break; @@ -1012,10 +1014,7 @@ export class PreTrainedModel extends Callable { 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, - ); + prepared_logits_processor.onTokensSampled(sampled_token_ids, all_input_ids); if (streamer) { streamer.put(generated_input_ids); } From b9881afa6c45cae876d3a522100316ee062f5e4e Mon Sep 17 00:00:00 2001 From: Nico Martin Date: Wed, 19 Aug 2026 15:08:24 +0200 Subject: [PATCH 14/17] removed llguidance dependency and created a much faster variant of response-constraint --- packages/transformers-llguidance/README.md | 60 - .../src/LlguidanceConstraint.ts | 315 --- packages/transformers-llguidance/src/index.ts | 2 - .../transformers-llguidance/src/utils/mask.ts | 125 -- .../src/utils/types.ts | 22 - .../tests/llguidance-constraint.test.js | 167 -- .../README.md | 58 + .../jest.config.mjs | 0 .../package.json | 19 +- .../performance/bench-real.mjs | 150 ++ .../performance/complex-json.mjs | 65 + .../performance/run.mjs | 98 + .../performance/simple-json.mjs | 16 + .../performance/simple-regex.mjs | 8 + .../scripts/build.mjs | 3 +- .../scripts/dev.mjs | 10 +- .../src/ResponseConstraint.ts | 92 + .../src/engine/constraint.ts | 312 +++ .../src/engine/index.ts | 2 + .../src/engine/json.ts | 1819 +++++++++++++++++ .../src/engine/regex.ts | 419 ++++ .../src/engine/tokenizer.ts | 183 ++ .../src/engine/types.ts | 12 + .../src/index.ts | 2 + .../src/utils/mask.ts | 30 + .../tests/corpus/helpers.mjs | 43 + .../tests/corpus/loader.mjs | 14 + .../tests/corpus/run.mjs | 12 + .../tests/corpus/runtime.mjs | 113 + .../tests/response-constraint.test.js | 506 +++++ .../tsconfig.json | 0 pnpm-lock.yaml | 11 +- 32 files changed, 3972 insertions(+), 716 deletions(-) delete mode 100644 packages/transformers-llguidance/README.md delete mode 100644 packages/transformers-llguidance/src/LlguidanceConstraint.ts delete mode 100644 packages/transformers-llguidance/src/index.ts delete mode 100644 packages/transformers-llguidance/src/utils/mask.ts delete mode 100644 packages/transformers-llguidance/src/utils/types.ts delete mode 100644 packages/transformers-llguidance/tests/llguidance-constraint.test.js create mode 100644 packages/transformers-response-constraint/README.md rename packages/{transformers-llguidance => transformers-response-constraint}/jest.config.mjs (100%) rename packages/{transformers-llguidance => transformers-response-constraint}/package.json (77%) create mode 100644 packages/transformers-response-constraint/performance/bench-real.mjs create mode 100644 packages/transformers-response-constraint/performance/complex-json.mjs create mode 100644 packages/transformers-response-constraint/performance/run.mjs create mode 100644 packages/transformers-response-constraint/performance/simple-json.mjs create mode 100644 packages/transformers-response-constraint/performance/simple-regex.mjs rename packages/{transformers-llguidance => transformers-response-constraint}/scripts/build.mjs (82%) rename packages/{transformers-llguidance => transformers-response-constraint}/scripts/dev.mjs (74%) create mode 100644 packages/transformers-response-constraint/src/ResponseConstraint.ts create mode 100644 packages/transformers-response-constraint/src/engine/constraint.ts create mode 100644 packages/transformers-response-constraint/src/engine/index.ts create mode 100644 packages/transformers-response-constraint/src/engine/json.ts create mode 100644 packages/transformers-response-constraint/src/engine/regex.ts create mode 100644 packages/transformers-response-constraint/src/engine/tokenizer.ts create mode 100644 packages/transformers-response-constraint/src/engine/types.ts create mode 100644 packages/transformers-response-constraint/src/index.ts create mode 100644 packages/transformers-response-constraint/src/utils/mask.ts create mode 100644 packages/transformers-response-constraint/tests/corpus/helpers.mjs create mode 100644 packages/transformers-response-constraint/tests/corpus/loader.mjs create mode 100644 packages/transformers-response-constraint/tests/corpus/run.mjs create mode 100644 packages/transformers-response-constraint/tests/corpus/runtime.mjs create mode 100644 packages/transformers-response-constraint/tests/response-constraint.test.js rename packages/{transformers-llguidance => transformers-response-constraint}/tsconfig.json (100%) diff --git a/packages/transformers-llguidance/README.md b/packages/transformers-llguidance/README.md deleted file mode 100644 index 55b1e2fee..000000000 --- a/packages/transformers-llguidance/README.md +++ /dev/null @@ -1,60 +0,0 @@ -# @huggingface/transformers-llguidance - -Experimental constrained-generation helpers for Transformers.js. - -This package exports `LlguidanceConstraint`, which turns an llguidance response format into the `logits_processor` and `stopping_criteria` objects accepted by Transformers.js generation. - -```js -import { LlguidanceConstraint } from "@huggingface/transformers-llguidance"; - -const constraint = await LlguidanceConstraint.fromResponseFormat(tokenizer, { - type: "json_schema", - json_schema: { - type: "object", - properties: { - answer: { type: "string" }, - }, - required: ["answer"], - additionalProperties: false, - }, -}); - -try { - await model.generate({ - ...inputs, - logits_processor: constraint.logits_processor, - stopping_criteria: constraint.stopping_criteria, - }); -} finally { - constraint.dispose(); -} -``` - -Constraints currently support a single generated sequence at a time. Generation throws when the logical batch size is not `1` rather than sharing mutable grammar state across sequences. - -The interpreter is automatically disposed when llguidance reaches a terminal state. Call `dispose()` in a `finally` block as shown above to also release resources when generation ends for another reason, such as `max_new_tokens` or cancellation. - -If llguidance reports acceptance before sampling, the constraint forces the tokenizer's EOS token so no unconstrained token is appended. Compatible tokenizer objects must expose an EOS ID such as `eos_token_id` or `eosTokenId`; generation fails closed if acceptance occurs without one. - -## Regex constraints - -Use `type: "regex"` to constrain generation to a regular expression. For example, this only allows ISO-like dates in `YYYY-MM-DD` format: - -```js -import { LlguidanceConstraint } from "@huggingface/transformers-llguidance"; - -const constraint = await LlguidanceConstraint.fromResponseFormat(tokenizer, { - type: "regex", - regex: "\\d{4}-\\d{2}-\\d{2}", -}); - -try { - const output = await model.generate({ - ...inputs, - logits_processor: constraint.logits_processor, - stopping_criteria: constraint.stopping_criteria, - }); -} finally { - constraint.dispose(); -} -``` diff --git a/packages/transformers-llguidance/src/LlguidanceConstraint.ts b/packages/transformers-llguidance/src/LlguidanceConstraint.ts deleted file mode 100644 index 9fb0e560d..000000000 --- a/packages/transformers-llguidance/src/LlguidanceConstraint.ts +++ /dev/null @@ -1,315 +0,0 @@ -import { - env, - LogitsProcessor, - LogitsProcessorList, - StoppingCriteria, - logger, - type Tensor, -} from '@huggingface/transformers'; -import { - loadBundledLLGuidance, - type JSONSchema, - type LLGuidanceMaskResult, - type LLGuidanceTokenizerSource, -} from 'llguidance'; - -import { applyMask, forceTokens, summarizeMaskResult } from './utils/mask'; -import { type LlguidanceState, type LlguidanceStats } from './utils/types'; - -export type ResponseFormat = - | { type: 'json_object' } - | { type: 'json_schema'; json_schema: JSONSchema } - | { type: 'regex'; regex: string }; -export type { LlguidanceStats }; - -export class LlguidanceConstraint { - static async fromResponseFormat(tokenizer: LLGuidanceTokenizerSource, response_format: ResponseFormat) { - logger.debug('[LlguidanceConstraint] loading llguidance', { - response_format, - }); - - const runtime = await loadBundledLLGuidance(); - const eosTokenIds = getEosTokenIds(tokenizer); - const interpreter = runtime.createInterpreter({ - tokenizer, - response_format, - }); - - logger.debug('[LlguidanceConstraint] interpreter created'); - - const state: LlguidanceState = { - completed: false, - disposed: false, - eosTokenIds, - interpreter, - step: 0, - stats: { - steps: 0, - computeMaskMs: 0, - applyMaskMs: 0, - commitTokenMs: 0, - trieNodesVisited: 0, - sharedJsonMaskCacheHits: 0, - sharedJsonMaskCacheMisses: 0, - }, - }; - - const logits_processor = new LogitsProcessorList(); - logits_processor.push(new LlguidanceLogitsProcessor(state)); - - return { - logits_processor, - stopping_criteria: new LlguidanceStoppingCriteria(state), - stats: state.stats, - dispose: () => disposeState(state), - }; - } -} - -class LlguidanceLogitsProcessor extends LogitsProcessor { - private state: LlguidanceState; - - constructor(state: LlguidanceState) { - super(); - this.state = state; - } - - _call(_inputIds: bigint[][], logits: Tensor) { - assertSingleSequence(_inputIds.length, this.state); - - if (this.state.completed) { - if (isDebugEnabled()) { - logger.debug('[LlguidanceLogitsProcessor] skip completed', { - step: this.state.step, - }); - } - return logits; - } - if (this.state.disposed) { - throw new Error('LlguidanceConstraint has been disposed.'); - } - - this.state.step++; - this.state.stats.steps = this.state.step; - const vocabSize = logits.dims.at(-1); - if (vocabSize === undefined || !Number.isInteger(vocabSize) || vocabSize <= 0) { - throw failState(this.state, 'LlguidanceConstraint requires logits with a vocabulary dimension.'); - } - if (isDebugEnabled()) { - logger.debug('[LlguidanceLogitsProcessor] compute mask', { - step: this.state.step, - logitsDims: logits.dims, - vocabSize, - }); - } - - let result: LLGuidanceMaskResult; - const maskStart = performance.now(); - const maskWords = Math.ceil(vocabSize / 32); - if (this.state.maskBuffer?.length !== maskWords) { - this.state.maskBuffer = new Uint32Array(maskWords); - } - try { - result = this.state.interpreter.computeMaskInto(this.state.maskBuffer); - } catch (error) { - disposeState(this.state); - throw error; - } - this.state.stats.computeMaskMs += performance.now() - maskStart; - const maskInstrumentation = this.state.interpreter.maskInstrumentation?.(); - if (maskInstrumentation) { - this.state.stats.trieNodesVisited = maskInstrumentation.trieNodesVisited; - this.state.stats.sharedJsonMaskCacheHits = maskInstrumentation.sharedJsonMaskCacheHits; - this.state.stats.sharedJsonMaskCacheMisses = maskInstrumentation.sharedJsonMaskCacheMisses; - } - - if (isDebugEnabled()) { - logger.debug('[LlguidanceLogitsProcessor] mask result', { - step: this.state.step, - result: summarizeMaskResult(result), - }); - } - - if ('stop' in result && result.stop) { - this.state.stats.stopReason = result.reason; - if (result.reason === 'dead_end') { - throw failState(this.state, 'llguidance reached a dead end before satisfying the constraint.'); - } - this.state.completed = true; - try { - forceTokens(logits, this.state.eosTokenIds); - } catch (error) { - disposeState(this.state); - throw error; - } - disposeState(this.state); - if (isDebugEnabled()) { - logger.debug('[LlguidanceLogitsProcessor] stopped by mask', { - step: this.state.step, - }); - } - return logits; - } - - if ('mask' in result) { - const applyStart = performance.now(); - try { - if (isDebugEnabled()) { - const applied = applyMask(logits, result.mask, result.vocabSize, true); - logger.debug('[LlguidanceLogitsProcessor] mask applied', { - step: this.state.step, - ...applied, - }); - } else { - applyMask(logits, result.mask, result.vocabSize); - } - } catch (error) { - disposeState(this.state); - throw error; - } - this.state.stats.applyMaskMs += performance.now() - applyStart; - } - - return logits; - } - - onTokenSampled(tokenId: number, batchIdx: number, inputIds: bigint[][]) { - if (isDebugEnabled()) { - logger.debug('[LlguidanceLogitsProcessor] token sampled', { - step: this.state.step, - tokenId, - batchIdx, - inputLength: inputIds[batchIdx]?.length, - completed: this.state.completed, - }); - } - - if (this.state.completed) return; - if (this.state.disposed) { - throw new Error('LlguidanceConstraint has been disposed.'); - } - - const commitStart = performance.now(); - let result; - try { - result = this.state.interpreter.commitToken(tokenId); - } catch (error) { - disposeState(this.state); - throw error; - } - this.state.stats.commitTokenMs += performance.now() - commitStart; - if (isDebugEnabled()) { - logger.debug('[LlguidanceLogitsProcessor] token committed', { - step: this.state.step, - tokenId, - result, - }); - } - - if (result.backtrack > 0) { - throw failState( - this.state, - `llguidance requested backtracking by ${result.backtrack} token(s), which Transformers.js does not support.`, - ); - } - - if (result.stop) { - this.state.completed = true; - this.state.stats.stopReason = 'accepted'; - disposeState(this.state); - if (isDebugEnabled()) { - logger.debug('[LlguidanceLogitsProcessor] stopped by commit', { - step: this.state.step, - tokenId, - }); - } - } - } - - onTokensSampled(tokenIds: number[], inputIds: bigint[][]) { - assertSingleSequence(tokenIds.length, this.state); - assertSingleSequence(inputIds.length, this.state); - if (isDebugEnabled()) { - logger.debug('[LlguidanceLogitsProcessor] tokens sampled', { - step: this.state.step, - tokenIds, - inputLengths: inputIds.map((ids) => ids.length), - completed: this.state.completed, - }); - } - - this.onTokenSampled(tokenIds[0], 0, inputIds); - } -} - -class LlguidanceStoppingCriteria extends StoppingCriteria { - private state: LlguidanceState; - - constructor(state: LlguidanceState) { - super(); - this.state = state; - } - - _call(inputIds: ArrayLike[]) { - assertSingleSequence(inputIds.length, this.state); - const result = new Array(inputIds.length).fill(this.state.completed); - if (isDebugEnabled()) { - logger.debug('[LlguidanceStoppingCriteria] call', { - step: this.state.step, - completed: this.state.completed, - result, - inputLengths: inputIds.map((ids) => ids.length), - }); - } - return result; - } -} - -function isDebugEnabled() { - // Keep compatibility with Transformers.js releases that predate the LogLevel export. - return env.logLevel <= DEBUG_LOG_LEVEL; -} - -const DEBUG_LOG_LEVEL = 10; - -function assertSingleSequence(batchSize: number, state: LlguidanceState) { - if (batchSize !== 1) { - throw failState(state, `LlguidanceConstraint currently supports batch size 1; received ${batchSize}.`); - } -} - -function failState(state: LlguidanceState, message: string) { - disposeState(state); - return new Error(message); -} - -function disposeState(state: LlguidanceState) { - if (state.disposed) return; - state.disposed = true; - state.maskBuffer = undefined; - state.interpreter.dispose(); -} - -function getEosTokenIds(tokenizer: LLGuidanceTokenizerSource) { - const source = tokenizer as Record; - const values = [ - source.eos_token_id, - source.eosTokenId, - source.eos_token_ids, - source.eosTokenIds, - source.eos_token, - source.eos_tokens, - ]; - - if (typeof source.eosToken === 'function') { - values.push(source.eosToken.call(tokenizer)); - } - - return [ - ...new Set( - values - .flat() - .filter((value): value is number => typeof value === 'number' && Number.isInteger(value) && value >= 0), - ), - ]; -} diff --git a/packages/transformers-llguidance/src/index.ts b/packages/transformers-llguidance/src/index.ts deleted file mode 100644 index 5f709e319..000000000 --- a/packages/transformers-llguidance/src/index.ts +++ /dev/null @@ -1,2 +0,0 @@ -export { LlguidanceConstraint } from './LlguidanceConstraint'; -export type { LlguidanceStats, ResponseFormat } from './LlguidanceConstraint'; diff --git a/packages/transformers-llguidance/src/utils/mask.ts b/packages/transformers-llguidance/src/utils/mask.ts deleted file mode 100644 index 66706999a..000000000 --- a/packages/transformers-llguidance/src/utils/mask.ts +++ /dev/null @@ -1,125 +0,0 @@ -import { type Tensor } from '@huggingface/transformers'; -import type { LLGuidanceMaskResult } from 'llguidance'; - -type LogitsData = Float32Array | Float64Array | number[]; - -export function summarizeMaskResult(result: LLGuidanceMaskResult): - | { stop: true; reason: 'accepted' | 'dead_end' } - | { - maskLength: number; - vocabSize: number; - allowed: number; - sampleAllowedTokenIds: number[]; - } { - if (!('mask' in result)) { - return { stop: true, reason: result.reason }; - } - - return { - maskLength: result.mask.length, - vocabSize: result.vocabSize, - allowed: countAllowed(result.mask, result.vocabSize), - sampleAllowedTokenIds: sampleAllowedTokenIds(result.mask, result.vocabSize), - }; -} - -export function applyMask(logits: Tensor, mask: Uint32Array, vocabSize: number, includeSummary = false) { - if (!Number.isInteger(vocabSize) || vocabSize <= 0) { - throw new Error(`llguidance returned an invalid vocabulary size: ${vocabSize}.`); - } - - const data = logits.data as LogitsData; - const stride = logits.dims.at(-1); - if (stride === undefined || !Number.isInteger(stride) || stride <= 0) { - throw new Error('LlguidanceConstraint requires logits with a vocabulary dimension.'); - } - if (vocabSize > stride) { - throw new Error(`llguidance vocabulary size ${vocabSize} exceeds logits vocabulary size ${stride}.`); - } - if (mask.length < Math.ceil(vocabSize / 32)) { - throw new Error(`llguidance returned a mask that is too short for vocabulary size ${vocabSize}.`); - } - - const batchSize = Math.floor(data.length / stride); - for (let batch = 0; batch < batchSize; ++batch) { - const offset = batch * stride; - applyPackedMask(data, mask, offset, vocabSize); - if (stride > vocabSize) { - // Logits padded beyond the grammar vocab can never be committed. - data.fill(-Infinity, offset + vocabSize, offset + stride); - } - } - - if (!includeSummary) return undefined; - - const allowed = countAllowed(mask, vocabSize); - return { vocabSize, batchSize, masked: (vocabSize - allowed) * batchSize, allowed }; -} - -export function forceTokens(logits: Tensor, tokenIds: number[]) { - const data = logits.data as LogitsData; - const vocabSize = logits.dims.at(-1); - if (vocabSize === undefined || !Number.isInteger(vocabSize) || vocabSize <= 0) { - throw new Error('LlguidanceConstraint requires logits with a vocabulary dimension.'); - } - if (tokenIds.length === 0) { - throw new Error('LlguidanceConstraint cannot stop on acceptance because the tokenizer has no EOS token ID.'); - } - - const kept = tokenIds.map((tokenId) => { - if (!Number.isInteger(tokenId) || tokenId < 0 || tokenId >= vocabSize) { - throw new Error(`Tokenizer EOS token ID ${tokenId} is outside logits vocabulary size ${vocabSize}.`); - } - return data[tokenId]; - }); - - data.fill(-Infinity); - for (let i = 0; i < tokenIds.length; ++i) { - data[tokenIds[i]] = Number.isFinite(kept[i]) ? kept[i] : 0; - } -} - -// Hot path: runs once per generated token over the whole vocab. Grammar masks -// are skewed: most 32-token words are either fully allowed (skip) or fully -// banned (memset), so per-bit work only happens on the few mixed words. -function applyPackedMask(data: LogitsData, mask: Uint32Array, offset: number, vocabSize: number) { - const numWords = vocabSize >>> 5; - for (let word = 0; word < numWords; ++word) { - const bits = mask[word] | 0; - if (bits === -1) continue; - const base = offset + (word << 5); - if (bits === 0) { - data.fill(-Infinity, base, base + 32); - continue; - } - for (let bit = 0; bit < 32; ++bit) { - if (!(bits & (1 << bit))) data[base + bit] = -Infinity; - } - } - - for (let tokenId = numWords << 5; tokenId < vocabSize; ++tokenId) { - if (!(mask[tokenId >>> 5] & (1 << (tokenId & 31)))) { - data[offset + tokenId] = -Infinity; - } - } -} - -function countAllowed(mask: Uint32Array, vocabSize: number) { - let allowed = 0; - for (let tokenId = 0; tokenId < vocabSize; ++tokenId) { - if (isAllowed(mask, tokenId)) allowed++; - } - return allowed; -} - -function sampleAllowedTokenIds(mask: Uint32Array, vocabSize: number) { - const tokenIds: number[] = []; - for (let tokenId = 0; tokenId < vocabSize && tokenIds.length < 25; ++tokenId) { - if (isAllowed(mask, tokenId)) tokenIds.push(tokenId); - } - return tokenIds; -} - -function isAllowed(mask: Uint32Array, tokenId: number) { - return Boolean(mask[tokenId >> 5] & (1 << (tokenId & 31))); -} diff --git a/packages/transformers-llguidance/src/utils/types.ts b/packages/transformers-llguidance/src/utils/types.ts deleted file mode 100644 index b5c3c1470..000000000 --- a/packages/transformers-llguidance/src/utils/types.ts +++ /dev/null @@ -1,22 +0,0 @@ -import type { LLGuidanceInterpreterHandle } from 'llguidance'; - -export type LlguidanceStats = { - steps: number; - computeMaskMs: number; - applyMaskMs: number; - commitTokenMs: number; - trieNodesVisited: number; - sharedJsonMaskCacheHits: number; - sharedJsonMaskCacheMisses: number; - stopReason?: 'accepted' | 'dead_end'; -}; - -export type LlguidanceState = { - completed: boolean; - disposed: boolean; - eosTokenIds: number[]; - interpreter: LLGuidanceInterpreterHandle; - maskBuffer?: Uint32Array; - step: number; - stats: LlguidanceStats; -}; diff --git a/packages/transformers-llguidance/tests/llguidance-constraint.test.js b/packages/transformers-llguidance/tests/llguidance-constraint.test.js deleted file mode 100644 index ceda77c84..000000000 --- a/packages/transformers-llguidance/tests/llguidance-constraint.test.js +++ /dev/null @@ -1,167 +0,0 @@ -import { jest } from "@jest/globals"; -import { Tensor } from "@huggingface/transformers"; - -const computeMask = jest.fn(); -const computeMaskInto = jest.fn(); -const commitToken = jest.fn(); -const disposeInterpreter = jest.fn(); -const createInterpreter = jest.fn(() => ({ - computeMask, - computeMaskInto, - commitToken, - dispose: disposeInterpreter, -})); -const loadBundledLLGuidance = jest.fn(async () => ({ createInterpreter })); - -jest.unstable_mockModule("llguidance", () => ({ - loadBundledLLGuidance, -})); - -const { LlguidanceConstraint } = await import("../dist/index.js"); - -describe("LlguidanceConstraint", () => { - beforeEach(() => { - computeMask.mockReset(); - computeMaskInto.mockReset(); - commitToken.mockReset(); - disposeInterpreter.mockReset(); - createInterpreter.mockClear(); - loadBundledLLGuidance.mockClear(); - - computeMaskInto.mockImplementation((target) => { - target[0] = 0b0101; - return { mask: target, vocabSize: 4 }; - }); - commitToken.mockReturnValue({ stop: false, backtrack: 0, ffTokens: [] }); - }); - - it("loads llguidance and applies a packed mask", async () => { - const tokenizer = { name: "tokenizer" }; - const response_format = { type: "json_schema", json_schema: { type: "object" } }; - const { logits_processor } = await LlguidanceConstraint.fromResponseFormat(tokenizer, response_format); - const logits = new Tensor("float32", new Float32Array([1, 2, 3, 4]), [1, 4]); - - logits_processor([[0n]], logits); - - expect(loadBundledLLGuidance).toHaveBeenCalledTimes(1); - expect(createInterpreter).toHaveBeenCalledWith({ tokenizer, response_format }); - expect(Array.from(logits.data)).toEqual([1, -Infinity, 3, -Infinity]); - }); - - it("reuses the packed mask buffer", async () => { - const { logits_processor } = await LlguidanceConstraint.fromResponseFormat({}, { type: "json_object" }); - const first = new Tensor("float32", new Float32Array([1, 2, 3, 4]), [1, 4]); - const second = new Tensor("float32", new Float32Array([5, 6, 7, 8]), [1, 4]); - - logits_processor([[0n]], first); - logits_processor([[0n]], second); - - expect(computeMask).not.toHaveBeenCalled(); - expect(computeMaskInto).toHaveBeenCalledTimes(2); - expect(computeMaskInto.mock.calls[0][0]).toBe(computeMaskInto.mock.calls[1][0]); - }); - - it("rejects batched generation before computing a mask", async () => { - const { logits_processor } = await LlguidanceConstraint.fromResponseFormat({}, { type: "json_object" }); - const logits = new Tensor("float32", new Float32Array(8), [2, 4]); - - expect(() => logits_processor([[0n], [0n]], logits)).toThrow("currently supports batch size 1"); - expect(computeMaskInto).not.toHaveBeenCalled(); - expect(disposeInterpreter).toHaveBeenCalledTimes(1); - }); - - it("commits one sampled token and stops on acceptance", async () => { - commitToken.mockReturnValue({ stop: true, backtrack: 0, ffTokens: [] }); - const { logits_processor, stopping_criteria, stats } = await LlguidanceConstraint.fromResponseFormat({}, { type: "json_object" }); - - expect(stopping_criteria([[0n]])).toEqual([false]); - logits_processor.onTokensSampled([1], [[0n, 1n]]); - - expect(commitToken).toHaveBeenCalledWith(1); - expect(stopping_criteria([[0n, 1n]])).toEqual([true]); - expect(stats.stopReason).toBe("accepted"); - expect(disposeInterpreter).toHaveBeenCalledTimes(1); - }); - - it("throws when llguidance requests backtracking", async () => { - commitToken.mockReturnValue({ stop: false, backtrack: 1, ffTokens: [] }); - const { logits_processor } = await LlguidanceConstraint.fromResponseFormat({}, { type: "json_object" }); - - expect(() => logits_processor.onTokensSampled([1], [[0n, 1n]])).toThrow("requested backtracking by 1 token"); - expect(disposeInterpreter).toHaveBeenCalledTimes(1); - }); - - it("records and throws on a dead end", async () => { - computeMaskInto.mockReturnValue({ stop: true, reason: "dead_end" }); - const { logits_processor, stats } = await LlguidanceConstraint.fromResponseFormat({}, { type: "json_object" }); - const logits = new Tensor("float32", new Float32Array(4), [1, 4]); - - expect(() => logits_processor([[0n]], logits)).toThrow("reached a dead end"); - expect(stats.stopReason).toBe("dead_end"); - expect(disposeInterpreter).toHaveBeenCalledTimes(1); - }); - - it("stops and disposes when computeMask reports acceptance", async () => { - computeMaskInto.mockReturnValue({ stop: true, reason: "accepted" }); - const { logits_processor, stopping_criteria, stats } = await LlguidanceConstraint.fromResponseFormat({ eos_token_id: 1 }, { type: "json_object" }); - const logits = new Tensor("float32", new Float32Array([1, 2, 3, 4]), [1, 4]); - - logits_processor([[0n]], logits); - logits_processor.onTokensSampled([0], [[0n]]); - - expect(Array.from(logits.data)).toEqual([-Infinity, 2, -Infinity, -Infinity]); - expect(commitToken).not.toHaveBeenCalled(); - expect(stopping_criteria([[0n]])).toEqual([true]); - expect(stats.stopReason).toBe("accepted"); - expect(disposeInterpreter).toHaveBeenCalledTimes(1); - }); - - it("fails closed on mask acceptance when the tokenizer has no EOS token ID", async () => { - computeMaskInto.mockReturnValue({ stop: true, reason: "accepted" }); - const { logits_processor } = await LlguidanceConstraint.fromResponseFormat({}, { type: "json_object" }); - const logits = new Tensor("float32", new Float32Array([1, 2, 3, 4]), [1, 4]); - - expect(() => logits_processor([[0n]], logits)).toThrow("tokenizer has no EOS token ID"); - expect(disposeInterpreter).toHaveBeenCalledTimes(1); - }); - - it("fails closed when logits have no vocabulary dimension", async () => { - const { logits_processor } = await LlguidanceConstraint.fromResponseFormat({}, { type: "json_object" }); - const logits = { dims: [], data: new Float32Array(4) }; - - expect(() => logits_processor([[0n]], logits)).toThrow("requires logits with a vocabulary dimension"); - expect(computeMaskInto).not.toHaveBeenCalled(); - expect(disposeInterpreter).toHaveBeenCalledTimes(1); - }); - - it("fails closed when mask vocabulary metadata is invalid", async () => { - computeMaskInto.mockReturnValue({ mask: new Uint32Array([0b0101]), vocabSize: undefined }); - const { logits_processor } = await LlguidanceConstraint.fromResponseFormat({}, { type: "json_object" }); - const logits = new Tensor("float32", new Float32Array(4), [1, 4]); - - expect(() => logits_processor([[0n]], logits)).toThrow("invalid vocabulary size"); - expect(disposeInterpreter).toHaveBeenCalledTimes(1); - }); - - it("exposes idempotent disposal", async () => { - const { dispose, logits_processor } = await LlguidanceConstraint.fromResponseFormat({}, { type: "json_object" }); - - dispose(); - dispose(); - - expect(disposeInterpreter).toHaveBeenCalledTimes(1); - expect(() => logits_processor.onTokensSampled([1], [[0n, 1n]])).toThrow("has been disposed"); - }); - - it("disposes and rethrows interpreter errors", async () => { - const error = new Error("unexpected"); - computeMaskInto.mockImplementation(() => { - throw error; - }); - const { logits_processor } = await LlguidanceConstraint.fromResponseFormat({}, { type: "json_object" }); - const logits = new Tensor("float32", new Float32Array(4), [1, 4]); - - expect(() => logits_processor([[0n]], logits)).toThrow(error); - expect(disposeInterpreter).toHaveBeenCalledTimes(1); - }); -}); diff --git a/packages/transformers-response-constraint/README.md b/packages/transformers-response-constraint/README.md new file mode 100644 index 000000000..bcdc91723 --- /dev/null +++ b/packages/transformers-response-constraint/README.md @@ -0,0 +1,58 @@ +# @huggingface/transformers-response-constraint + +Experimental constrained-generation helpers for Transformers.js. + +This dependency-free package exports `ResponseConstraint`, which turns a JSON-schema, JSON-object, or regex response format into the `logits_processor` and `stopping_criteria` objects accepted by Transformers.js generation. + +The constraint engine is implemented specifically for Transformers.js and has no runtime dependencies. + +```js +import { ResponseConstraint } from "@huggingface/transformers-response-constraint"; + +const constraint = ResponseConstraint.fromResponseFormat(tokenizer, { + type: "json_schema", + json_schema: { + type: "object", + properties: { + answer: { type: "string" }, + }, + required: ["answer"], + additionalProperties: false, + }, +}); + +await model.generate({ + ...inputs, + logits_processor: constraint.logits_processor, + stopping_criteria: constraint.stopping_criteria, +}); +``` + +Constraints currently support a single generated sequence at a time. Generation throws when the logical batch size is not `1` rather than sharing mutable grammar state across sequences. + +## Supported constraints + +The JSON engine implements a practical JSON Schema 2020-12 profile. This includes deep `const` and `enum`, exact decimal bounds and `multipleOf`, recognized string formats, tuple and homogeneous arrays, `contains`, deep `uniqueItems`, object property and dependency assertions, `allOf`/`anyOf`/`oneOf`/`not`, conditionals, local `$ref`, recursive `$defs`, draft-07 compatibility, and root-level `x-guidance` separators. External and dynamic references and unevaluated-property/item assertions remain unsupported. + +The regex engine performs full-string matching and supports literals, UTF-8 literals, alternation, groups, character classes, `.`, `\\d`, `\\s`, `\\w`, anchors, and greedy `*`, `+`, `?`, and `{m,n}` quantifiers. Lookarounds, backreferences, lazy quantifiers, and Unicode character classes are intentionally unsupported. + +When the generated bytes satisfy the constraint, the logits processor exposes only the tokenizer's EOS token as a valid completion. Sampling EOS updates the shared stopping criterion. + +## Regex constraints + +Use `type: "regex"` to constrain generation to a regular expression. For example, this only allows ISO-like dates in `YYYY-MM-DD` format: + +```js +import { ResponseConstraint } from "@huggingface/transformers-response-constraint"; + +const constraint = ResponseConstraint.fromResponseFormat(tokenizer, { + type: "regex", + regex: "\\d{4}-\\d{2}-\\d{2}", +}); + +const output = await model.generate({ + ...inputs, + logits_processor: constraint.logits_processor, + stopping_criteria: constraint.stopping_criteria, +}); +``` diff --git a/packages/transformers-llguidance/jest.config.mjs b/packages/transformers-response-constraint/jest.config.mjs similarity index 100% rename from packages/transformers-llguidance/jest.config.mjs rename to packages/transformers-response-constraint/jest.config.mjs diff --git a/packages/transformers-llguidance/package.json b/packages/transformers-response-constraint/package.json similarity index 77% rename from packages/transformers-llguidance/package.json rename to packages/transformers-response-constraint/package.json index f284a8e96..e825674b1 100644 --- a/packages/transformers-llguidance/package.json +++ b/packages/transformers-response-constraint/package.json @@ -1,7 +1,7 @@ { - "name": "@huggingface/transformers-llguidance", + "name": "@huggingface/transformers-response-constraint", "version": "0.0.0", - "description": "llguidance integration helpers for Transformers.js constrained generation", + "description": "Dependency-free constrained generation for Transformers.js", "main": "./dist/index.cjs", "types": "./types/index.d.ts", "type": "module", @@ -18,11 +18,14 @@ "scripts": { "format": "prettier --write . --ignore-path ../../.prettierignore", "format:check": "prettier --check . --ignore-path ../../.prettierignore", - "typegen": "tsc --build", + "typegen": "tsc --build --force", "dev": "node scripts/dev.mjs", "build": "node scripts/build.mjs && pnpm typegen", + "performance": "node performance/run.mjs", "pretest": "pnpm build", - "test": "node --experimental-vm-modules node_modules/jest/bin/jest.js --verbose" + "test": "node --experimental-vm-modules node_modules/jest/bin/jest.js --verbose", + "test:json-corpus": "node tests/corpus/run.mjs", + "test:json-corpus:original": "node tests/corpus/run.mjs --original" }, "repository": { "type": "git", @@ -32,8 +35,9 @@ "transformers", "transformers.js", "huggingface", - "llguidance", - "constrained-generation" + "constrained-generation", + "structured-output", + "json-schema" ], "author": "Hugging Face", "license": "Apache-2.0", @@ -62,8 +66,5 @@ ], "publishConfig": { "access": "public" - }, - "dependencies": { - "llguidance": "0.2.0" } } diff --git a/packages/transformers-response-constraint/performance/bench-real.mjs b/packages/transformers-response-constraint/performance/bench-real.mjs new file mode 100644 index 000000000..d7b517813 --- /dev/null +++ b/packages/transformers-response-constraint/performance/bench-real.mjs @@ -0,0 +1,150 @@ +import { AutoTokenizer, Tensor } from "@huggingface/transformers"; +import { ResponseConstraint } from "/Users/nico/Documents/Dev/transformers.js/packages/transformers-response-constraint/dist/index.js"; + +const MODEL = process.argv[2] ?? "onnx-community/gemma-4-E2B-it-ONNX"; + +const simpleJsonSchemaFormat = { + type: "json_schema", + json_schema: { + "x-guidance": { + whitespace_flexible: false, + item_separator: ", ", + key_separator: ": ", + }, + type: "object", + properties: { + answer: { type: "string", minLength: 1, maxLength: 120 }, + }, + required: ["answer"], + additionalProperties: false, + }, +}; + +const complexJsonSchemaFormat = { + type: "json_schema", + json_schema: { + "x-guidance": { + whitespace_flexible: false, + item_separator: ", ", + key_separator: ": ", + }, + type: "object", + properties: { + answer: { + type: "object", + properties: { + text: { type: "string", minLength: 1, maxLength: 120 }, + tone: { enum: ["friendly", "formal", "playful"] }, + language: { enum: ["en", "de", "fr", "es"] }, + }, + required: ["text", "tone", "language"], + additionalProperties: false, + }, + alternatives: { + type: "array", + minItems: 1, + maxItems: 3, + items: { type: "string", minLength: 1, maxLength: 120 }, + }, + metadata: { + type: "object", + properties: { + confidence: { type: "integer", minimum: 0, maximum: 100 }, + safe: { type: "boolean" }, + tags: { + type: "array", + minItems: 1, + maxItems: 3, + items: { type: "string", minLength: 1, maxLength: 20 }, + }, + }, + required: ["confidence", "safe", "tags"], + additionalProperties: false, + }, + }, + required: ["answer", "alternatives", "metadata"], + additionalProperties: false, + }, +}; + +const regexFormat = { type: "regex", regex: "(Hello|Hi|Hey)( there)?[!.]" }; + +const cases = [ + { + name: "simple JSON", + format: simpleJsonSchemaFormat, + output: '{"answer": "Hello there! How can I help you today?"}', + }, + { name: "regex", format: regexFormat, output: "Hello there!" }, + { + name: "complex JSON", + format: complexJsonSchemaFormat, + output: + '{"answer": {"text": "Hello there! How can I help you today?", "tone": "friendly", "language": "en"}, "alternatives": ["Hi! What can I do for you?", "Hey, great to see you."], "metadata": {"confidence": 95, "safe": true, "tags": ["greeting", "friendly"]}}', + }, +]; + +const tokenizer = await AutoTokenizer.from_pretrained(MODEL); +const eosId = tokenizer.eos_token_id; +console.log("eos:", tokenizer.eos_token, eosId); + +const VOCAB = 262144; + +function percentile(values, p) { + const sorted = [...values].sort((a, b) => a - b); + return sorted[Math.floor(sorted.length * p)]; +} + +function run(format, tokenIds, label) { + const t0 = performance.now(); + const constraint = ResponseConstraint.fromResponseFormat(tokenizer, format); + const setupMs = performance.now() - t0; + + const scores = new Tensor("float32", new Float32Array(VOCAB), [1, VOCAB]); + const inputIds = [0n]; + const stepMs = []; + for (const tokenId of tokenIds) { + scores.data.fill(0); + const s = performance.now(); + constraint.logits_processor([inputIds], scores); + if (!Number.isFinite(scores.data[tokenId])) { + const txt = tokenizer.decode([tokenId]); + throw new Error(`${label} rejected token ${tokenId} (${JSON.stringify(txt)})`); + } + inputIds.push(BigInt(tokenId)); + constraint.logits_processor.onTokensSampled([tokenId], [inputIds]); + constraint.stopping_criteria([inputIds]); + stepMs.push(performance.now() - s); + } + return { setupMs, stepMs }; +} + +// EOS-terminated token sequence for each case +for (const c of cases) { + const ids = tokenizer.encode(c.output, { add_special_tokens: false }); + c.tokenIds = [...ids, Number(eosId)]; +} + +const results = []; +for (const c of cases) { + // Run 1: cold (first-ever fromResponseFormat pays tokenizer trie build; caches empty) + const cold = run(c.format, c.tokenIds, c.name); + // Run 2: warm-ish (mask caches populated from run 1) + const warm1 = run(c.format, c.tokenIds, c.name); + const warm2 = run(c.format, c.tokenIds, c.name); + results.push({ + name: c.name, + tokens: c.tokenIds.length, + "cold setup ms": cold.setupMs.toFixed(1), + "warm setup ms": warm2.setupMs.toFixed(2), + "cold sum ms": cold.stepMs.reduce((a, b) => a + b, 0).toFixed(1), + "cold med": percentile(cold.stepMs, 0.5).toFixed(2), + "cold p90": percentile(cold.stepMs, 0.9).toFixed(2), + "cold max": Math.max(...cold.stepMs).toFixed(2), + "warm sum ms": warm2.stepMs.reduce((a, b) => a + b, 0).toFixed(1), + "warm med": percentile(warm2.stepMs, 0.5).toFixed(3), + "warm p90": percentile(warm2.stepMs, 0.9).toFixed(2), + "warm max": Math.max(...warm2.stepMs).toFixed(2), + }); +} +console.table(results); diff --git a/packages/transformers-response-constraint/performance/complex-json.mjs b/packages/transformers-response-constraint/performance/complex-json.mjs new file mode 100644 index 000000000..3c49c120b --- /dev/null +++ b/packages/transformers-response-constraint/performance/complex-json.mjs @@ -0,0 +1,65 @@ +export default { + name: "complex JSON schema", + responseFormat: { + type: "json_schema", + json_schema: { + type: "object", + properties: { + request_id: { type: "string", pattern: "^[a-z0-9-]{8,36}$" }, + status: { enum: ["queued", "running", "completed", "failed"] }, + user: { + type: "object", + properties: { + id: { type: "integer", minimum: 1 }, + email: { type: "string", format: "email" }, + roles: { + type: "array", + items: { enum: ["admin", "editor", "viewer"] }, + minItems: 1, + uniqueItems: true, + }, + }, + required: ["id", "email", "roles"], + additionalProperties: false, + }, + results: { + type: "array", + minItems: 2, + maxItems: 8, + items: { + type: "object", + properties: { + label: { type: "string", minLength: 2, maxLength: 32 }, + score: { type: "number", minimum: 0, maximum: 1 }, + tags: { + type: "array", + items: { type: "string", pattern: "^[a-z-]+$" }, + maxItems: 5, + }, + metadata: { + anyOf: [ + { type: "null" }, + { + type: "object", + properties: { + source: { type: "string" }, + cached: { type: "boolean" }, + }, + required: ["source", "cached"], + additionalProperties: false, + }, + ], + }, + }, + required: ["label", "score", "tags", "metadata"], + additionalProperties: false, + }, + }, + }, + required: ["request_id", "status", "user", "results"], + additionalProperties: false, + }, + }, + output: + '{"request_id":"req-2026-a1","status":"completed","user":{"id":42,"email":"user@example.com","roles":["admin","editor"]},"results":[{"label":"primary","score":0.97,"tags":["fast","verified"],"metadata":{"source":"cache-v2","cached":true}},{"label":"fallback","score":0.81,"tags":["review"],"metadata":null}]}', +}; diff --git a/packages/transformers-response-constraint/performance/run.mjs b/packages/transformers-response-constraint/performance/run.mjs new file mode 100644 index 000000000..2f3c3a6d8 --- /dev/null +++ b/packages/transformers-response-constraint/performance/run.mjs @@ -0,0 +1,98 @@ +import { Tensor } from "@huggingface/transformers"; + +import { ResponseConstraint } from "../dist/index.js"; +import complexJson from "./complex-json.mjs"; +import simpleJson from "./simple-json.mjs"; +import simpleRegex from "./simple-regex.mjs"; + +const WARMUP_RUNS = 2; +const MEASURED_RUNS = 10; +const EOS_TOKEN_ID = 256; +const VOCAB_SIZE = Number(process.env.VOCAB_SIZE ?? 8192); +const encoder = new TextEncoder(); +const tokenizer = { + tokens: [ + ...Array.from({ length: EOS_TOKEN_ID }, (_, tokenId) => [tokenId]), + [], + ...Array.from({ length: VOCAB_SIZE - EOS_TOKEN_ID - 1 }, (_, index) => [ + ...encoder.encode(` token-${index.toString(36)}`), + ]), + ], + eos_token_id: EOS_TOKEN_ID, + special_token_ids: [EOS_TOKEN_ID], +}; +const encodedCases = [simpleJson, simpleRegex, complexJson].map((testCase) => ({ + ...testCase, + tokenIds: encoder.encode(testCase.output), +})); +const coldResults = new Map(); + +for (const testCase of encodedCases) { + coldResults.set(testCase, await measure(testCase)); +} + +for (const testCase of encodedCases) { + for (let run = 0; run < WARMUP_RUNS; ++run) await measure(testCase); +} + +const results = []; +for (const testCase of encodedCases) { + let elapsedMs = 0; + let processorMs = 0; + let updateMs = 0; + for (let run = 0; run < MEASURED_RUNS; ++run) { + const result = await measure(testCase); + elapsedMs += result.elapsedMs; + processorMs += result.processorMs; + updateMs += result.updateMs; + } + results.push({ + constraint: testCase.name, + vocabulary: VOCAB_SIZE, + tokens: testCase.tokenIds.length + 1, + "cold ms": format(coldResults.get(testCase).elapsedMs), + "total ms": format(elapsedMs / MEASURED_RUNS), + "ms/token": format( + elapsedMs / MEASURED_RUNS / (testCase.tokenIds.length + 1), + ), + "processor ms": format(processorMs / MEASURED_RUNS), + "update ms": format(updateMs / MEASURED_RUNS), + }); +} + +console.table(results); + +async function measure(testCase) { + const constraint = await ResponseConstraint.fromResponseFormat( + tokenizer, + testCase.responseFormat, + ); + const scores = new Tensor("float32", new Float32Array(VOCAB_SIZE), [ + 1, + VOCAB_SIZE, + ]); + const inputIds = [0n]; + const startedAt = performance.now(); + let processorMs = 0; + let updateMs = 0; + for (const tokenId of [...testCase.tokenIds, EOS_TOKEN_ID]) { + scores.data.fill(0); + let stepStartedAt = performance.now(); + constraint.logits_processor([inputIds], scores); + processorMs += performance.now() - stepStartedAt; + if (!Number.isFinite(scores.data[tokenId])) + throw new Error(`${testCase.name} rejected token ${tokenId}`); + inputIds.push(BigInt(tokenId)); + stepStartedAt = performance.now(); + constraint.logits_processor.onTokensSampled([tokenId], [inputIds]); + constraint.stopping_criteria([inputIds]); + updateMs += performance.now() - stepStartedAt; + } + if (!constraint.stopping_criteria([inputIds])[0]) + throw new Error(`${testCase.name} did not stop after EOS`); + return { elapsedMs: performance.now() - startedAt, processorMs, updateMs }; +} + +function format(value) { + return value.toFixed(3); +} diff --git a/packages/transformers-response-constraint/performance/simple-json.mjs b/packages/transformers-response-constraint/performance/simple-json.mjs new file mode 100644 index 000000000..d3ce4eb17 --- /dev/null +++ b/packages/transformers-response-constraint/performance/simple-json.mjs @@ -0,0 +1,16 @@ +export default { + name: "simple JSON schema", + responseFormat: { + type: "json_schema", + json_schema: { + type: "object", + properties: { + answer: { type: "string" }, + confidence: { type: "number", minimum: 0, maximum: 1 }, + }, + required: ["answer", "confidence"], + additionalProperties: false, + }, + }, + output: '{"answer":"Paris","confidence":0.98}', +}; diff --git a/packages/transformers-response-constraint/performance/simple-regex.mjs b/packages/transformers-response-constraint/performance/simple-regex.mjs new file mode 100644 index 000000000..6a8cc6fe2 --- /dev/null +++ b/packages/transformers-response-constraint/performance/simple-regex.mjs @@ -0,0 +1,8 @@ +export default { + name: "simple regex", + responseFormat: { + type: "regex", + regex: "[A-Z]{3}-\\d{4}", + }, + output: "ABC-2026", +}; diff --git a/packages/transformers-llguidance/scripts/build.mjs b/packages/transformers-response-constraint/scripts/build.mjs similarity index 82% rename from packages/transformers-llguidance/scripts/build.mjs rename to packages/transformers-response-constraint/scripts/build.mjs index 9b5ca9f5a..8138e6ce8 100644 --- a/packages/transformers-llguidance/scripts/build.mjs +++ b/packages/transformers-response-constraint/scripts/build.mjs @@ -2,6 +2,7 @@ import { build } from "esbuild"; import { rmSync } from "node:fs"; rmSync("dist", { recursive: true, force: true }); +rmSync("types", { recursive: true, force: true }); const common = { entryPoints: ["src/index.ts"], @@ -9,7 +10,7 @@ const common = { platform: "neutral", target: "es2022", sourcemap: true, - external: ["@huggingface/transformers", "llguidance"], + external: ["@huggingface/transformers"], }; await Promise.all([ diff --git a/packages/transformers-llguidance/scripts/dev.mjs b/packages/transformers-response-constraint/scripts/dev.mjs similarity index 74% rename from packages/transformers-llguidance/scripts/dev.mjs rename to packages/transformers-response-constraint/scripts/dev.mjs index 4c5186e26..74adbb1cd 100644 --- a/packages/transformers-llguidance/scripts/dev.mjs +++ b/packages/transformers-response-constraint/scripts/dev.mjs @@ -12,15 +12,15 @@ const watchLogger = { build.onStart(() => { startTime = performance.now(); - console.log(`[transformers-llguidance] rebuilding ${build.initialOptions.outfile}...`); + console.log(`[transformers-response-constraint] rebuilding ${build.initialOptions.outfile}...`); }); build.onEnd((result) => { const duration = (performance.now() - startTime).toFixed(2); if (result.errors.length > 0) { - console.log(`[transformers-llguidance] rebuild failed in ${duration}ms`); + console.log(`[transformers-response-constraint] rebuild failed in ${duration}ms`); } else { - console.log(`[transformers-llguidance] rebuilt ${build.initialOptions.outfile} in ${duration}ms`); + console.log(`[transformers-response-constraint] rebuilt ${build.initialOptions.outfile} in ${duration}ms`); } }); }, @@ -32,7 +32,7 @@ const common = { platform: "neutral", target: "es2022", sourcemap: true, - external: ["@huggingface/transformers", "llguidance"], + external: ["@huggingface/transformers"], plugins: [watchLogger], }; @@ -56,7 +56,7 @@ const tscWatch = spawn("tsc", ["--build", "--watch", "--preserveWatchOutput"], { shell: true, }); -console.log("Watching @huggingface/transformers-llguidance..."); +console.log("Watching @huggingface/transformers-response-constraint..."); process.on("SIGINT", async () => { tscWatch.kill(); diff --git a/packages/transformers-response-constraint/src/ResponseConstraint.ts b/packages/transformers-response-constraint/src/ResponseConstraint.ts new file mode 100644 index 000000000..bcf1c0420 --- /dev/null +++ b/packages/transformers-response-constraint/src/ResponseConstraint.ts @@ -0,0 +1,92 @@ +import { LogitsProcessor, LogitsProcessorList, StoppingCriteria, type Tensor } from '@huggingface/transformers'; + +import { + createTokenConstraint, + prepareTokenizer, + type JSONSchema, + type TokenConstraint, + type TokenizerSource, +} from './engine'; +import { applyMask } from './utils/mask'; + +export type ResponseFormat = + | { type: 'json_object' } + | { type: 'json_schema'; json_schema: JSONSchema } + | { type: 'regex'; regex: string }; + +type GenerationState = { + completed: boolean; + constraint: TokenConstraint; + mask?: Uint32Array; +}; + +export class ResponseConstraint { + /** + * Precomputes the tokenizer-derived data structures used by every + * constraint. The first constraint per tokenizer otherwise pays this cost + * (hundreds of milliseconds for large vocabularies) inside + * `fromResponseFormat`; call this once after loading the model to pay it + * early instead. + */ + static warmup(tokenizer: TokenizerSource): void { + prepareTokenizer(tokenizer); + } + + static fromResponseFormat(tokenizer: TokenizerSource, responseFormat: ResponseFormat) { + const state: GenerationState = { + completed: false, + constraint: createTokenConstraint(tokenizer, responseFormat), + }; + const logits_processor = new LogitsProcessorList(); + logits_processor.push(new ConstraintLogitsProcessor(state)); + return { + logits_processor, + stopping_criteria: new ConstraintStoppingCriteria(state), + }; + } +} + +class ConstraintLogitsProcessor extends LogitsProcessor { + constructor(private readonly state: GenerationState) { + super(); + } + + _call(inputIds: bigint[][], logits: Tensor) { + assertSingleSequence(inputIds.length); + if (this.state.completed) return logits; + const logitsVocabSize = logits.dims.at(-1); + if (logitsVocabSize === undefined || !Number.isInteger(logitsVocabSize) || logitsVocabSize <= 0) { + throw new Error('ResponseConstraint requires logits with a vocabulary dimension.'); + } + const words = Math.ceil(logitsVocabSize / 32); + if (this.state.mask?.length !== words) this.state.mask = new Uint32Array(words); + if (!this.state.constraint.fillMask(this.state.mask)) { + throw new Error('The constraint reached a dead end before producing a valid output.'); + } + applyMask(logits, this.state.mask, this.state.constraint.vocabSize); + return logits; + } + + onTokensSampled(tokenIds: number[], inputIds: bigint[][]) { + assertSingleSequence(tokenIds.length); + assertSingleSequence(inputIds.length); + if (!this.state.completed) this.state.completed = this.state.constraint.commit(tokenIds[0]); + } +} + +class ConstraintStoppingCriteria extends StoppingCriteria { + constructor(private readonly state: GenerationState) { + super(); + } + + _call(inputIds: ArrayLike[]) { + assertSingleSequence(inputIds.length); + return [this.state.completed]; + } +} + +function assertSingleSequence(batchSize: number): void { + if (batchSize !== 1) { + throw new Error(`ResponseConstraint currently supports batch size 1; received ${batchSize}.`); + } +} diff --git a/packages/transformers-response-constraint/src/engine/constraint.ts b/packages/transformers-response-constraint/src/engine/constraint.ts new file mode 100644 index 000000000..ef762ecc7 --- /dev/null +++ b/packages/transformers-response-constraint/src/engine/constraint.ts @@ -0,0 +1,312 @@ +import { compileJsonSchema } from './json'; +import { compileRegex } from './regex'; +import { extractTokenizer, type TokenizerData } from './tokenizer'; +import type { ConstraintState, JSONSchema, TokenizerSource } from './types'; + +type TrieNode = { childBytes: number[]; childNodes: TrieNode[]; tokenIds: number[] }; +type CachedTokenizer = { + data: TokenizerData; + trie: TrieNode; + stringExceptionalTrie: TrieNode; + stringSafeMask: Uint32Array; + stringSafeCount: number; + stringSafeLengths: Uint32Array; + maxStringSafeLength: number; + maxTokenByteLength: number; + boundedStringMasks: Map; + schemaMaskCaches: WeakMap; + booleanSchemaMaskCaches: [MaskCache, MaskCache]; + jsonObjectMaskCache: MaskCache; + regexMaskCaches: Map; +}; +type ResponseFormat = + | { type: 'json_object' } + | { type: 'json_schema'; json_schema: JSONSchema } + | { type: 'regex'; regex: string }; + +export type TokenConstraint = { + vocabSize: number; + fillMask(target: Uint32Array): boolean; + commit(tokenId: number): boolean; +}; + +const tokenizerCache = new WeakMap(); +const JSON_OBJECT_SCHEMA: JSONSchema = { type: 'object' }; + +/** + * Builds and caches the tokenizer-derived data structures (token tries, string + * masks) ahead of time. This is the expensive part of creating the first + * constraint for a tokenizer (hundreds of milliseconds for a 256k vocabulary), + * so calling this right after loading a model moves that cost off the first + * generation. Subsequent calls with the same tokenizer are free. + */ +export function prepareTokenizer(tokenizerSource: TokenizerSource): void { + cachedTokenizer(tokenizerSource); +} + +export function createTokenConstraint( + tokenizerSource: TokenizerSource, + responseFormat: ResponseFormat, +): TokenConstraint { + const tokenizer = cachedTokenizer(tokenizerSource); + const machine = createMachine(responseFormat, tokenizer); + const maskCache = cacheFor(tokenizer, responseFormat); + let state = machine.initial; + // Post-transition states discovered during the trie walk, so commit() can + // reuse them. Entries are only valid when their stamp matches the current + // fillMask() generation; bumping the stamp invalidates all of them at once + // without refilling the vocabulary-sized array on every step. + const tokenStates: Array = new Array(tokenizer.data.tokens.length); + const tokenStamps = new Int32Array(tokenizer.data.tokens.length); + let stamp = 0; + + return { + vocabSize: tokenizer.data.tokens.length, + fillMask(target) { + const words = Math.ceil(tokenizer.data.tokens.length / 32); + if (target.length < words) throw new RangeError(`Mask target requires at least ${words} words.`); + target.fill(0); + stamp++; + const cacheKey = machine.maskKey?.(state); + const cachedMask = cacheKey === undefined ? undefined : maskCache?.get(cacheKey); + if (cachedMask !== undefined) { + target.set(cachedMask); + return true; + } + let allowed = 0; + if (machine.accepting(state)) { + setBit(target, tokenizer.data.eosTokenId); + allowed++; + } + const stringCapacity = machine.stringCapacity?.(state); + if (stringCapacity !== undefined) { + const safe = boundedStringMask(tokenizer, stringCapacity); + target.set(safe.mask); + allowed += safe.count; + } + const nodes: TrieNode[] = [stringCapacity === undefined ? tokenizer.trie : tokenizer.stringExceptionalTrie]; + const states: unknown[] = [state]; + while (nodes.length > 0) { + const node = nodes.pop()!; + const current = states.pop()!; + for (const tokenId of node.tokenIds) { + if (tokenizer.data.specialTokenIds.has(tokenId)) continue; + setBit(target, tokenId); + tokenStates[tokenId] = current; + tokenStamps[tokenId] = stamp; + allowed++; + } + for (let index = 0; index < node.childNodes.length; ++index) { + const next = machine.transition(current, node.childBytes[index]); + if (!machine.viable(next)) continue; + nodes.push(node.childNodes[index]); + states.push(next); + } + } + if (allowed > 0 && cacheKey !== undefined) { + maskCache?.set(cacheKey, target.subarray(0, words)); + } + return allowed > 0; + }, + commit(tokenId) { + if (!Number.isInteger(tokenId) || tokenId < 0 || tokenId >= tokenizer.data.tokens.length) { + throw new RangeError(`Token ${tokenId} is outside the tokenizer vocabulary.`); + } + if (tokenId === tokenizer.data.eosTokenId) { + if (!machine.accepting(state)) throw new Error(`Token ${tokenId} does not satisfy the constraint.`); + return true; + } + if (tokenizer.data.specialTokenIds.has(tokenId)) { + throw new Error(`Token ${tokenId} does not satisfy the constraint.`); + } + let next: unknown; + if (tokenStamps[tokenId] === stamp && stamp > 0) { + next = tokenStates[tokenId]; + } else { + next = state; + for (const byte of tokenizer.data.tokens[tokenId]) next = machine.transition(next, byte); + } + stamp++; + if (!machine.viable(next)) throw new Error(`Token ${tokenId} does not satisfy the constraint.`); + state = next; + return false; + }, + }; +} + +function createMachine(responseFormat: ResponseFormat, tokenizer: CachedTokenizer): ConstraintState { + if (responseFormat?.type === 'regex') { + if (typeof responseFormat.regex !== 'string') throw new TypeError('response_format.regex must be a string.'); + return compileRegex(responseFormat.regex) as ConstraintState; + } + if (responseFormat?.type === 'json_schema') { + return compileJsonSchema(responseFormat.json_schema, tokenizer.maxTokenByteLength) as ConstraintState; + } + if (responseFormat?.type === 'json_object') { + return compileJsonSchema(JSON_OBJECT_SCHEMA, tokenizer.maxTokenByteLength) as ConstraintState; + } + throw new TypeError(`Unsupported response format: ${String((responseFormat as { type?: unknown })?.type)}.`); +} + +function cachedTokenizer(source: TokenizerSource): CachedTokenizer { + let cached = tokenizerCache.get(source as object); + if (cached === undefined) { + const data = extractTokenizer(source); + const stringExceptionalTokenIds: number[] = []; + const stringSafeMask = new Uint32Array(Math.ceil(data.tokens.length / 32)); + const stringSafeLengths = new Uint32Array(data.tokens.length); + let stringSafeCount = 0; + let maxStringSafeLength = 0; + let maxTokenByteLength = 0; + for (let tokenId = 0; tokenId < data.tokens.length; ++tokenId) { + const special = data.specialTokenIds.has(tokenId); + if (!special && data.tokens[tokenId].length > maxTokenByteLength) { + maxTokenByteLength = data.tokens[tokenId].length; + } + const length = special ? undefined : safeStringTokenLength(data.tokens[tokenId]); + if (length !== undefined) { + stringSafeCount++; + stringSafeLengths[tokenId] = length; + if (length > maxStringSafeLength) maxStringSafeLength = length; + setBit(stringSafeMask, tokenId); + } else { + stringExceptionalTokenIds.push(tokenId); + } + } + cached = { + data, + trie: createTrie(data.tokens), + stringExceptionalTrie: createTrie(data.tokens, stringExceptionalTokenIds), + stringSafeMask, + stringSafeCount, + stringSafeLengths, + maxStringSafeLength, + maxTokenByteLength, + boundedStringMasks: new Map(), + schemaMaskCaches: new WeakMap(), + booleanSchemaMaskCaches: [new MaskCache(), new MaskCache()], + jsonObjectMaskCache: new MaskCache(), + regexMaskCaches: new Map(), + }; + tokenizerCache.set(source as object, cached); + } + return cached; +} + +function cacheFor(tokenizer: CachedTokenizer, responseFormat: ResponseFormat): MaskCache | undefined { + if (responseFormat.type === 'regex') { + let cache = tokenizer.regexMaskCaches.get(responseFormat.regex); + if (cache === undefined) { + cache = new MaskCache(); + if (tokenizer.regexMaskCaches.size >= 16) { + tokenizer.regexMaskCaches.delete(tokenizer.regexMaskCaches.keys().next().value!); + } + tokenizer.regexMaskCaches.set(responseFormat.regex, cache); + } + return cache; + } + if (responseFormat.type === 'json_object') return tokenizer.jsonObjectMaskCache; + const schema = responseFormat.json_schema; + if (typeof schema === 'boolean') return tokenizer.booleanSchemaMaskCaches[schema ? 1 : 0]; + let cache = tokenizer.schemaMaskCaches.get(schema); + if (cache === undefined) { + cache = new MaskCache(); + tokenizer.schemaMaskCaches.set(schema, cache); + } + return cache; +} + +class MaskCache { + private readonly masks = new Map(); + private words = 0; + + get(key: string): Uint32Array | undefined { + const mask = this.masks.get(key); + if (mask === undefined) return undefined; + this.masks.delete(key); + this.masks.set(key, mask); + return mask; + } + + set(key: string, source: Uint32Array): void { + const mask = source.slice(); + const previous = this.masks.get(key); + if (previous !== undefined) { + this.words -= previous.length; + this.masks.delete(key); + } + this.masks.set(key, mask); + this.words += mask.length; + while (this.masks.size > 256 || this.words > 1_048_576) { + const oldestKey = this.masks.keys().next().value!; + const oldest = this.masks.get(oldestKey)!; + this.masks.delete(oldestKey); + this.words -= oldest.length; + } + } +} + +function createTrie(tokens: Uint8Array[], tokenIds?: number[]): TrieNode { + const root: TrieNode = { childBytes: [], childNodes: [], tokenIds: [] }; + const size = tokenIds === undefined ? tokens.length : tokenIds.length; + for (let index = 0; index < size; ++index) { + const tokenId = tokenIds === undefined ? index : tokenIds[index]; + const bytes = tokens[tokenId]; + let node = root; + for (let position = 0; position < bytes.length; ++position) { + const byte = bytes[position]; + const childIndex = node.childBytes.indexOf(byte); + if (childIndex === -1) { + const child: TrieNode = { childBytes: [], childNodes: [], tokenIds: [] }; + node.childBytes.push(byte); + node.childNodes.push(child); + node = child; + } else { + node = node.childNodes[childIndex]; + } + } + node.tokenIds.push(tokenId); + } + return root; +} + +const safeStringDecoder = new TextDecoder('utf-8', { fatal: true }); + +function safeStringTokenLength(bytes: Uint8Array): number | undefined { + if (bytes.length === 0) return undefined; + let length = 0; + for (const byte of bytes) { + if (byte < 0x20 || byte === 0x22 || byte === 0x5c) return undefined; + // Code points equal non-continuation bytes in valid UTF-8. + if ((byte & 0xc0) !== 0x80) length++; + } + try { + safeStringDecoder.decode(bytes); + } catch { + return undefined; + } + return length; +} + +function boundedStringMask(tokenizer: CachedTokenizer, capacity: number): { mask: Uint32Array; count: number } { + if (capacity >= tokenizer.maxStringSafeLength) { + return { mask: tokenizer.stringSafeMask, count: tokenizer.stringSafeCount }; + } + let cached = tokenizer.boundedStringMasks.get(capacity); + if (cached !== undefined) return cached; + const mask = new Uint32Array(Math.ceil(tokenizer.data.tokens.length / 32)); + let count = 0; + for (let tokenId = 0; tokenId < tokenizer.stringSafeLengths.length; ++tokenId) { + const length = tokenizer.stringSafeLengths[tokenId]; + if (length === 0 || length > capacity) continue; + setBit(mask, tokenId); + count++; + } + cached = { mask, count }; + tokenizer.boundedStringMasks.set(capacity, cached); + return cached; +} + +function setBit(mask: Uint32Array, tokenId: number): void { + mask[tokenId >>> 5] |= 1 << (tokenId & 31); +} diff --git a/packages/transformers-response-constraint/src/engine/index.ts b/packages/transformers-response-constraint/src/engine/index.ts new file mode 100644 index 000000000..7a5b500f0 --- /dev/null +++ b/packages/transformers-response-constraint/src/engine/index.ts @@ -0,0 +1,2 @@ +export { createTokenConstraint, prepareTokenizer, type TokenConstraint } from './constraint'; +export type { JSONSchema, TokenizerSource } from './types'; diff --git a/packages/transformers-response-constraint/src/engine/json.ts b/packages/transformers-response-constraint/src/engine/json.ts new file mode 100644 index 000000000..e7558cc3d --- /dev/null +++ b/packages/transformers-response-constraint/src/engine/json.ts @@ -0,0 +1,1819 @@ +import type { ConstraintState, JSONSchema } from './types'; + +type Schema = boolean | Record; +type Kind = 'null' | 'boolean' | 'number' | 'integer' | 'string' | 'array' | 'object'; +type JsonNode = + | { kind: 'null'; value: null } + | { kind: 'boolean'; value: boolean } + | { kind: 'number'; value: number; raw: string } + | { kind: 'string'; value: string } + | { kind: 'array'; value: unknown[]; items: JsonNode[] } + | { kind: 'object'; value: Record; entries: Array<{ key: string; node: JsonNode }> }; +type Guidance = { + itemSeparator: string; + keySeparator: string; + itemBytes: Uint8Array; + keyBytes: Uint8Array; + whitespaceFlexible: boolean; +}; +type Mode = + | 'value' + | 'array-value' + | 'object-key' + | 'colon' + | 'item-separator' + | 'after-value' + | 'string' + | 'key-string' + | 'escape' + | 'key-escape' + | 'unicode' + | 'key-unicode' + | 'number' + | 'literal' + | 'done' + | 'dead'; + +type ObjectFrame = { + kind: 'object'; + schema: Schema; + seen: ReadonlySet; + entries: ReadonlyArray<{ key: string; node: JsonNode }>; + key?: string; + childSchema?: Schema; +}; + +type ArrayFrame = { + kind: 'array'; + schema: Schema; + length: number; + items: readonly JsonNode[]; +}; + +type Frame = ObjectFrame | ArrayFrame; + +type JsonState = { + mode: Mode; + schema: Schema; + stack: readonly Frame[]; + bytes?: readonly number[]; + // Code points started in the current string and UTF-8 continuation bytes + // still expected (-1 once the byte sequence can no longer decode). Both are + // pure functions of `bytes`, maintained incrementally so that length checks + // do not have to re-decode the prefix on every byte. + stringLength?: number; + stringPending?: number; + text?: string; + literal?: string; + literalValue?: null | boolean; + index?: number; + highSurrogate?: number; + guidance: Guidance; +}; + +const DEFAULT_GUIDANCE: Guidance = { + itemSeparator: ',', + keySeparator: ':', + itemBytes: Uint8Array.of(0x2c), + keyBytes: Uint8Array.of(0x3a), + whitespaceFlexible: true, +}; +const DEAD: JsonState = { mode: 'dead', schema: false, stack: [], guidance: DEFAULT_GUIDANCE }; +const encoder = new TextEncoder(); +const decoder = new TextDecoder('utf-8', { fatal: true }); +const schemaIds = new WeakMap(); +const propertyKeyBytes = new WeakMap | null>(); +const finiteStringCache = new WeakMap | null>(); +let nextSchemaId = 0; +const SIMPLE_ESCAPES: Record = { + 0x22: '"', + 0x2f: '/', + 0x5c: '\\', + 0x62: '\b', + 0x66: '\f', + 0x6e: '\n', + 0x72: '\r', + 0x74: '\t', +}; + +export function compileJsonSchema(schema: JSONSchema, stringKeyClamp = Infinity): ConstraintState { + registerSchemaContext(schema, schema); + checkSchema(schema, '$'); + checkReferences(schema); + const guidance = guidanceFrom(schema); + return { + initial: { mode: 'value', schema, stack: [], guidance }, + transition, + viable: (state) => state !== DEAD, + accepting: isAccepting, + stringCapacity, + maskKey: (state) => stateMaskKey(state, stringKeyClamp), + }; +} + +function stringCapacity(state: JsonState): number | undefined { + if (state.mode !== 'string' || state.highSurrogate !== undefined || finiteStringValues(state.schema) !== null) + return undefined; + if ((state.stringPending ?? 0) !== 0) return undefined; + const maximum = directStringMaxLength(state.schema); + if (maximum === undefined) return Infinity; + return maximum - (state.stringLength ?? 0); +} + +function stateMaskKey(state: JsonState, stringKeyClamp: number): string | undefined { + if (state === DEAD) return undefined; + const frames = state.stack.map((frame) => { + if (frame.kind === 'object') { + const entries = [...frame.entries] + .sort((left, right) => left.key.localeCompare(right.key)) + .map((entry) => `${JSON.stringify(entry.key)}:${nodeKey(entry.node)}`) + .join(','); + return `o${schemaId(frame.schema)}[${entries}]${frame.key === undefined ? '' : `:${JSON.stringify(frame.key)}:${schemaId(frame.childSchema!)}`}`; + } + return `a${schemaId(frame.schema)}[${frame.items.map(nodeKey).join(',')}]`; + }); + const key = [ + state.mode, + schemaId(state.schema), + state.bytes === undefined ? '' : (stringLengthKey(state, stringKeyClamp) ?? bytesKey(state.bytes)), + state.text ?? '', + state.literal ?? '', + state.index ?? '', + state.highSurrogate ?? '', + ...frames, + ].join('|'); + return key.length <= 2048 ? key : undefined; +} + +// While inside a plain length-constrained string, the token mask depends only on +// the remaining capacity (clamped to the longest token) and whether minLength is +// already satisfied — not on the actual bytes typed so far. Keying on that makes +// every string-content step after the first share one cached mask. This is only +// sound when nothing can inspect the string's content: neither the string schema +// itself (pattern/format/const/enum or composition keywords) nor any enclosing +// frame (cross-value keywords such as uniqueItems, contains, or conditionals). +const CONTENT_DEPENDENT_STRING_KEYWORDS = ['pattern', 'format', 'const', 'enum', '$ref', 'allOf', 'anyOf', 'oneOf', 'not', 'if']; +const CONTENT_DEPENDENT_FRAME_KEYWORDS = [ + '$ref', + 'allOf', + 'anyOf', + 'oneOf', + 'not', + 'if', + 'uniqueItems', + 'contains', + 'dependentSchemas', + 'dependencies', +]; +const contentIndependentStrings = new WeakMap(); +const contentNeutralFrames = new WeakMap(); + +function isContentIndependent( + schema: Schema, + keywords: readonly string[], + cache: WeakMap, +): boolean { + if (schema === true) return true; + if (schema === false) return false; + let result = cache.get(schema); + if (result === undefined) { + result = keywords.every((keyword) => schema[keyword] === undefined); + cache.set(schema, result); + } + return result; +} + +function stringLengthKey(state: JsonState, clamp: number): string | undefined { + if (state.mode !== 'string' || state.highSurrogate !== undefined) return undefined; + if (!isContentIndependent(state.schema, CONTENT_DEPENDENT_STRING_KEYWORDS, contentIndependentStrings)) { + return undefined; + } + for (const frame of state.stack) { + if (!isContentIndependent(frame.schema, CONTENT_DEPENDENT_FRAME_KEYWORDS, contentNeutralFrames)) { + return undefined; + } + } + if ((state.stringPending ?? 0) !== 0) return undefined; + const length = state.stringLength ?? 0; + const maximum = directStringMaxLength(state.schema); + const remaining = Math.min(maximum === undefined ? Infinity : maximum - length, clamp); + const minimum = + state.schema !== true && state.schema !== false && typeof state.schema.minLength === 'number' + ? state.schema.minLength + : 0; + return `#${remaining === Infinity ? 'inf' : remaining}:${length >= minimum ? '' : length}`; +} + +function nodeKey(node: JsonNode): string { + if (node.kind === 'number') return `n${node.raw}`; + if (node.kind === 'string') return `s${JSON.stringify(node.value)}`; + if (node.kind === 'boolean') return node.value ? 't' : 'f'; + if (node.kind === 'null') return 'z'; + if (node.kind === 'array') return `[${node.items.map(nodeKey).join(',')}]`; + return `{${[...node.entries] + .sort((left, right) => left.key.localeCompare(right.key)) + .map((entry) => `${JSON.stringify(entry.key)}:${nodeKey(entry.node)}`) + .join(',')}}`; +} + +function bytesKey(bytes: readonly number[]): string { + let result = ''; + for (const byte of bytes) result += byte.toString(16).padStart(2, '0'); + return result; +} + +function schemaId(schema: Schema): string | number { + if (schema === true) return 't'; + if (schema === false) return 'f'; + let id = schemaIds.get(schema); + if (id === undefined) { + id = nextSchemaId++; + schemaIds.set(schema, id); + } + return id; +} + +function transition(state: JsonState, byte: number): JsonState { + if (state === DEAD) return DEAD; + if ( + state.guidance.whitespaceFlexible && + isWhitespace(byte) && + allowsWhitespace(state.mode) && + !separatorExpects(state, byte) + ) + return state; + + switch (state.mode) { + case 'value': + case 'array-value': + if (state.mode === 'array-value' && byte === 0x5d) return closeArray(state); + if (state.mode === 'array-value' && exceedsMaxItems(state.stack.at(-1) as ArrayFrame)) return DEAD; + return startValue(state, byte); + case 'object-key': + if (byte === 0x7d) return closeObject(state); + return byte === 0x22 + ? { ...state, mode: 'key-string', bytes: [], stringLength: 0, stringPending: 0 } + : DEAD; + case 'colon': + return separatorByte(state, byte, state.guidance.keyBytes, 'value'); + case 'item-separator': + return separatorByte( + state, + byte, + state.guidance.itemBytes, + state.stack.at(-1)?.kind === 'object' ? 'object-key' : 'array-value', + ); + case 'after-value': + return afterValue(state, byte); + case 'string': + case 'key-string': + return stringByte(state, byte); + case 'escape': + case 'key-escape': + return escapeByte(state, byte); + case 'unicode': + case 'key-unicode': + return unicodeByte(state, byte); + case 'number': + return numberByte(state, byte); + case 'literal': + return literalByte(state, byte); + case 'done': + case 'dead': + return DEAD; + } +} + +function startValue(state: JsonState, byte: number): JsonState { + const kind = byteKind(byte); + if (kind === undefined) return DEAD; + const schema = selectSchema(state.schema, kind); + if (schema === null) return DEAD; + + if (byte === 0x22) return { ...state, mode: 'string', schema, bytes: [], stringLength: 0, stringPending: 0 }; + if (byte === 0x7b) { + return { + mode: 'object-key', + schema, + stack: [...state.stack, { kind: 'object', schema, seen: new Set(), entries: [] }], + guidance: state.guidance, + }; + } + if (byte === 0x5b) { + return { + mode: 'array-value', + schema: itemSchema(schema, 0), + stack: [ + ...state.stack, + { + kind: 'array', + schema, + length: 0, + items: [], + }, + ], + guidance: state.guidance, + }; + } + if (byte === 0x74) return { ...state, mode: 'literal', schema, literal: 'true', literalValue: true, index: 1 }; + if (byte === 0x66) return { ...state, mode: 'literal', schema, literal: 'false', literalValue: false, index: 1 }; + if (byte === 0x6e) return { ...state, mode: 'literal', schema, literal: 'null', literalValue: null, index: 1 }; + const text = String.fromCharCode(byte); + if (!integerPrefixViable(schema, text)) return DEAD; + return { ...state, mode: 'number', schema, text }; +} + +function isAccepting(state: JsonState): boolean { + if (state.mode === 'done') return true; + return ( + state.mode === 'number' && + state.stack.length === 0 && + numberComplete(state.text!) && + validateNode(state.schema, { kind: 'number', value: Number(state.text), raw: state.text! }) + ); +} + +function stringByte(state: JsonState, byte: number): JsonState { + const isKey = state.mode === 'key-string'; + if (state.highSurrogate !== undefined && byte !== 0x5c) return DEAD; + if (byte === 0x22) { + const value = decodeString(state.bytes ?? []); + if (value === null) return DEAD; + if (isKey) return completeKey(state, value); + return completeValue(state, { kind: 'string', value }); + } + if (byte === 0x5c) { + if (isKey ? !keyEscapeAllowed(state) : !stringEscapeAllowed(state)) return DEAD; + return { ...state, mode: isKey ? 'key-escape' : 'escape' }; + } + if (byte < 0x20) return DEAD; + let length = state.stringLength ?? 0; + let pending = state.stringPending ?? 0; + if ((byte & 0xc0) === 0x80) { + pending = pending > 0 ? pending - 1 : -1; + } else { + pending = pending === 0 ? utf8Continuations(byte) : -1; + length++; + } + const bytes = [...(state.bytes ?? []), byte]; + if (isKey && !keyPrefixAllowed(state, bytes)) return DEAD; + if (!isKey && !stringContentAllowed(state.schema, bytes, length)) return DEAD; + return { ...state, bytes, stringLength: length, stringPending: pending }; +} + +function utf8Continuations(lead: number): number { + if (lead < 0x80) return 0; + if (lead < 0xc2) return -1; + if (lead < 0xe0) return 1; + if (lead < 0xf0) return 2; + if (lead < 0xf5) return 3; + return -1; +} + +function escapeByte(state: JsonState, byte: number): JsonState { + const isKey = state.mode === 'key-escape'; + if (state.highSurrogate !== undefined && byte !== 0x75) return DEAD; + if (byte === 0x75) { + return { ...state, mode: isKey ? 'key-unicode' : 'unicode', text: '' }; + } + const escaped = SIMPLE_ESCAPES[byte]; + if (escaped === undefined) return DEAD; + const bytes = [...(state.bytes ?? []), ...encoder.encode(escaped)]; + const length = (state.stringLength ?? 0) + 1; + const pending = (state.stringPending ?? 0) === 0 ? 0 : -1; + if (isKey && !keyPrefixAllowed(state, bytes)) return DEAD; + if (!isKey && !stringContentAllowed(state.schema, bytes, length)) return DEAD; + return { ...state, mode: isKey ? 'key-string' : 'string', bytes, stringLength: length, stringPending: pending }; +} + +function unicodeByte(state: JsonState, byte: number): JsonState { + if (!isHex(byte)) return DEAD; + const text = `${state.text ?? ''}${String.fromCharCode(byte)}`; + if (text.length < 4) { + if (state.mode === 'key-unicode' && !unicodeKeyPrefixAllowed(state, text)) return DEAD; + return { ...state, text }; + } + const codeUnit = Number.parseInt(text, 16); + const isKey = state.mode === 'key-unicode'; + const mode = isKey ? 'key-string' : 'string'; + if (codeUnit >= 0xd800 && codeUnit <= 0xdbff) { + if (state.highSurrogate !== undefined) return DEAD; + return { ...state, mode, text: undefined, highSurrogate: codeUnit }; + } + let value: string; + if (codeUnit >= 0xdc00 && codeUnit <= 0xdfff) { + if (state.highSurrogate === undefined) return DEAD; + const codePoint = 0x10000 + ((state.highSurrogate - 0xd800) << 10) + (codeUnit - 0xdc00); + value = String.fromCodePoint(codePoint); + } else { + if (state.highSurrogate !== undefined) return DEAD; + value = String.fromCharCode(codeUnit); + } + const bytes = [...(state.bytes ?? []), ...encoder.encode(value)]; + const length = (state.stringLength ?? 0) + 1; + const pending = (state.stringPending ?? 0) === 0 ? 0 : -1; + if (isKey && !keyPrefixAllowed(state, bytes)) return DEAD; + if (!isKey && !stringContentAllowed(state.schema, bytes, length)) return DEAD; + return { + ...state, + mode, + bytes, + text: undefined, + highSurrogate: undefined, + stringLength: length, + stringPending: pending, + }; +} + +function completeKey(state: JsonState, key: string): JsonState { + const frame = topObject(state); + if (frame.seen.has(key)) return DEAD; + const childSchema = propertySchema(frame.schema, key, frame.entries); + if (childSchema === null) return DEAD; + return { + ...state, + mode: 'colon', + bytes: undefined, + index: 0, + stack: replaceTop(state.stack, { ...frame, key, childSchema }), + }; +} + +function literalByte(state: JsonState, byte: number): JsonState { + const index = state.index ?? 0; + const literal = state.literal!; + if (byte !== literal.charCodeAt(index)) return DEAD; + if (index + 1 < literal.length) return { ...state, index: index + 1 }; + const node: JsonNode = + state.literalValue === null ? { kind: 'null', value: null } : { kind: 'boolean', value: state.literalValue! }; + return completeValue(state, node); +} + +function numberByte(state: JsonState, byte: number): JsonState { + if (isNumberByte(byte)) { + const text = `${state.text}${String.fromCharCode(byte)}`; + if (numberPrefixValid(text) && integerPrefixViable(state.schema, text)) return { ...state, text }; + } + const value = finishNumber(state); + return value === DEAD ? DEAD : transition(value, byte); +} + +// JSON Schema accepts any number encoding for integer fields ("0.9e1" is 9), +// but during generation that freedom creates traps a model cannot escape: +// after "0.9" under {type: "integer"} only digits and exponents stay viable, +// continuations such as "0.9e-" can never close the value, and zero-padding +// like "0.0e0000…" stays legal forever, so a stuck model streams digits until +// it hits the token limit. Integer-only fields therefore follow llguidance's +// integer shape -?(0|[1-9]\d*)(\.0+)?([eE]+?\d+)? — the value stays +// ±digits×10^exponent, so every prefix can still close — with fraction zeros +// and exponent digits capped (they carry no information beyond two digits), and +// bytes from which no in-range integer remains reachable are pruned. +const SCALE_DIGIT_LIMIT = 3; +const integerRestrictions = new WeakMap(); + +function integerPrefixViable(schema: Schema, text: string): boolean { + const restriction = integerRestriction(schema); + if (restriction === null) return true; + const { lo, hi } = restriction; + if (text === '-') return lo <= Math.min(hi, 0); + const match = /^(-?)(0|[1-9]\d*)(\.0{0,3})?(?:[eE]\+?(\d{0,3}))?$/.exec(text); + if (match === null) return false; + // "1.e" can never complete (the fraction needs a digit before the exponent) + if (match[3] === '.' && match[4] !== undefined) return false; + const negative = match[1] === '-'; + if (match[3] === undefined && match[4] === undefined) { + return integerDigitsReachable(lo, hi, negative, match[2]); + } + return integerScaleReachable(lo, hi, negative, match[2], match[4] ?? ''); +} + +function integerRestriction(schema: Schema): { lo: number; hi: number } | null { + if (schema === true || schema === false) return null; + let restriction = integerRestrictions.get(schema); + if (restriction === undefined) { + restriction = computeIntegerRestriction(schema); + integerRestrictions.set(schema, restriction); + } + return restriction; +} + +function computeIntegerRestriction(schema: Record): { lo: number; hi: number } | null { + let integerOnly = false; + if (schema.type !== undefined) { + const types = Array.isArray(schema.type) ? schema.type.map(String) : [String(schema.type)]; + integerOnly = types.includes('integer') && !types.includes('number'); + } else if (typeof schema.const === 'number') { + integerOnly = Number.isInteger(schema.const); + } else if (Array.isArray(schema.enum)) { + const numbers = schema.enum.filter((value): value is number => typeof value === 'number'); + integerOnly = numbers.length > 0 && numbers.every((value) => Number.isInteger(value)); + } + if (!integerOnly) return null; + let lo = -Infinity; + let hi = Infinity; + if (typeof schema.minimum === 'number') lo = Math.ceil(schema.minimum); + if (typeof schema.exclusiveMinimum === 'number') lo = Math.max(lo, Math.floor(schema.exclusiveMinimum) + 1); + if (typeof schema.maximum === 'number') hi = Math.floor(schema.maximum); + if (typeof schema.exclusiveMaximum === 'number') hi = Math.min(hi, Math.ceil(schema.exclusiveMaximum) - 1); + // Pruning must never reject a value the validator would accept; when a + // bound is beyond exact integer arithmetic, drop it instead of guessing. + if (!Number.isSafeInteger(lo)) lo = -Infinity; + if (!Number.isSafeInteger(hi)) hi = Infinity; + return { lo, hi }; +} + +// Whether some integer whose decimal representation starts with `digits` +// (extended by zero or more digits, or scaled by a later ".0…"/exponent) lies +// in [lo, hi]. The reachable values are ±[|P|·10^k, (|P|+1)·10^k - 1] for +// k ≥ 0, except that "0"/"-0" cannot take further digits. +function integerDigitsReachable(lo: number, hi: number, negative: boolean, digits: string): boolean { + if (digits.length > 16) return negative ? lo === -Infinity : hi === Infinity; + let low = Number(digits); + let high = low + 1; + const extensible = low !== 0; + for (;;) { + if (negative) { + if (-(high - 1) <= hi && -low >= lo) return true; + if (-low < lo) return false; + } else { + if (low <= hi && high - 1 >= lo) return true; + if (low > hi) return false; + } + if (!extensible) return false; + low *= 10; + high *= 10; + } +} + +// Once a fraction or exponent starts, the digits are frozen and the value is +// exactly ±digits×10^E for a still-typeable exponent E ≥ 0. When exponent +// digits have been typed, reachable exponents are those obtainable by extending +// the typed prefix within the digit limit (leading zeros allowed). +function integerScaleReachable( + lo: number, + hi: number, + negative: boolean, + digits: string, + exponentPrefix: string, +): boolean { + if (digits.length > 16) return negative ? lo === -Infinity : hi === Infinity; + const base = Number(digits); + if (base === 0) return lo <= 0 && 0 <= hi; + let magnitude = base; + // A typeable exponent has at most SCALE_DIGIT_LIMIT digits, so 999 bounds + // the search; finite bounds exit early via the overshoot check. + for (let exponent = 0; exponent <= 999; ++exponent) { + const value = negative ? -magnitude : magnitude; + const inRange = (lo === -Infinity || value >= lo) && (hi === Infinity || value <= hi); + if (inRange && exponentTypeable(exponent, exponentPrefix)) return true; + if (negative ? lo !== -Infinity && value < lo : hi !== Infinity && value > hi) return false; + magnitude *= 10; + } + return false; +} + +// Whether exponent value `exponent` can be written by extending the typed +// digits `prefix` to a full exponent of at most SCALE_DIGIT_LIMIT digits +// (at least one digit total; leading zeros allowed). +function exponentTypeable(exponent: number, prefix: string): boolean { + const typed = prefix === '' ? 0 : Number(prefix); + for (let extra = Math.max(0, 1 - prefix.length); extra <= SCALE_DIGIT_LIMIT - prefix.length; ++extra) { + const scale = 10 ** extra; + if (exponent >= typed * scale && exponent < (typed + 1) * scale) return true; + } + return false; +} + +function finishNumber(state: JsonState): JsonState { + const text = state.text!; + if (!numberComplete(text)) return DEAD; + return completeValue(state, { kind: 'number', value: Number(text), raw: text }); +} + +function completeValue(state: JsonState, node: JsonNode): JsonState { + if (!validateNode(state.schema, node)) return DEAD; + if (state.stack.length === 0) return { mode: 'done', schema: state.schema, stack: [], guidance: state.guidance }; + const frame = state.stack.at(-1)!; + if (frame.kind === 'object') { + if (frame.key === undefined) return DEAD; + const seen = new Set(frame.seen); + seen.add(frame.key); + const entries = [...frame.entries, { key: frame.key, node }]; + if (!partialObjectValid(frame.schema, entries, new Set())) return DEAD; + return { + mode: 'after-value', + schema: frame.schema, + stack: replaceTop(state.stack, { kind: 'object', schema: frame.schema, seen, entries }), + guidance: state.guidance, + }; + } + return { + mode: 'after-value', + schema: frame.schema, + stack: replaceTop(state.stack, { ...frame, length: frame.length + 1, items: [...frame.items, node] }), + guidance: state.guidance, + }; +} + +function afterValue(state: JsonState, byte: number): JsonState { + const frame = state.stack.at(-1); + if (frame === undefined) return DEAD; + if (frame.kind === 'object') { + if (byte === 0x7d) return closeObject(state); + if (!objectCanAddProperty(frame)) return DEAD; + return separatorByte( + { ...state, mode: 'item-separator', index: 0 }, + byte, + state.guidance.itemBytes, + 'object-key', + ); + } + if (byte === 0x5d) return closeArray(state); + if (exceedsMaxItems(frame)) return DEAD; + return separatorByte({ ...state, mode: 'item-separator', index: 0 }, byte, state.guidance.itemBytes, 'array-value'); +} + +function separatorByte(state: JsonState, byte: number, separator: Uint8Array, completedMode: Mode): JsonState { + const index = state.index ?? 0; + if (byte !== separator[index]) return DEAD; + if (index + 1 < separator.length) return { ...state, index: index + 1 }; + const frame = state.stack.at(-1); + return { + ...state, + mode: completedMode, + schema: + completedMode === 'value' + ? topObject(state).childSchema! + : completedMode === 'array-value' && frame?.kind === 'array' + ? itemSchema(frame.schema, frame.length) + : state.schema, + index: undefined, + }; +} + +function separatorExpects(state: JsonState, byte: number): boolean { + if (state.mode === 'colon') return state.guidance.keyBytes[state.index ?? 0] === byte; + if (state.mode === 'item-separator') return state.guidance.itemBytes[state.index ?? 0] === byte; + return false; +} + +function closeObject(state: JsonState): JsonState { + const frame = state.stack.at(-1); + if (frame?.kind !== 'object') return DEAD; + const value: Record = Object.create(null) as Record; + for (const entry of frame.entries) value[entry.key] = entry.node.value; + return completeValue( + { ...state, schema: frame.schema, stack: state.stack.slice(0, -1) }, + { kind: 'object', value, entries: [...frame.entries] }, + ); +} + +function closeArray(state: JsonState): JsonState { + const frame = state.stack.at(-1); + if (frame?.kind !== 'array') return DEAD; + return completeValue( + { ...state, schema: frame.schema, stack: state.stack.slice(0, -1) }, + { kind: 'array', value: frame.items.map((item) => item.value), items: [...frame.items] }, + ); +} + +function selectSchema(schema: Schema, kind: Kind): Schema | null { + return schemaMayAcceptKind(schema, kind, new Set()) ? schema : null; +} + +function propertySchema( + schema: Schema, + key: string, + entries: ReadonlyArray<{ key: string; node: JsonNode }> = [], +): Schema | null { + if (schema === true) return true; + if (schema === false) return null; + const schemas: Schema[] = []; + const direct = directPropertySchema(schema, key); + if (direct === null) return null; + if (direct !== true) schemas.push(direct); + for (const candidate of [schema, ...(Array.isArray(schema.allOf) ? schema.allOf.filter(isSchema) : [])]) { + if (candidate === true || candidate === false) continue; + const branch = selectedConditionalBranch(candidate, entries); + if (branch === undefined || branch === true) continue; + if (branch === false) return null; + const branchSchema = directPropertySchema(branch, key); + if (branchSchema === null) return null; + if (branchSchema !== true) schemas.push(branchSchema); + } + if (schemas.length === 0) return true; + return schemas.length === 1 ? schemas[0] : { allOf: schemas }; +} + +function directPropertySchema(schema: Record, key: string): Schema | null { + const properties = isRecord(schema.properties) ? schema.properties : {}; + const schemas: Schema[] = []; + if (isSchema(properties[key])) schemas.push(properties[key]); + if (isRecord(schema.patternProperties)) { + for (const [pattern, candidate] of Object.entries(schema.patternProperties)) { + if (new RegExp(pattern, 'u').test(key) && isSchema(candidate)) schemas.push(candidate); + } + } + if (schemas.length === 0) { + if (schema.additionalProperties === false) return null; + return isSchema(schema.additionalProperties) ? schema.additionalProperties : true; + } + return schemas.length === 1 ? schemas[0] : { allOf: schemas }; +} + +function selectedConditionalBranch( + schema: Record, + entries: ReadonlyArray<{ key: string; node: JsonNode }>, +): Schema | undefined { + if (!isSchema(schema.if) || schema.if === true || schema.if === false) return undefined; + const required = Array.isArray(schema.if.required) ? schema.if.required : []; + if (required.length === 0 || !required.every((key) => entries.some((entry) => entry.key === key))) return undefined; + const node = objectNode(entries); + const branch = validateNode(schema.if, node) ? schema.then : schema.else; + return isSchema(branch) ? branch : undefined; +} + +function partialObjectValid( + schema: Schema, + entries: ReadonlyArray<{ key: string; node: JsonNode }>, + seen: Set, +): boolean { + if (schema === true) return true; + if (schema === false) return false; + if (seen.has(schema)) return true; + seen.add(schema); + try { + for (const entry of entries) { + const child = directPropertySchema(schema, entry.key); + if (child === null || (child !== true && !validateNode(child, entry.node))) return false; + } + if ( + typeof schema.$ref === 'string' && + !partialObjectValid(resolveReference(schema, schema.$ref), entries, seen) + ) + return false; + if ( + Array.isArray(schema.allOf) && + !schema.allOf.every((child) => !isSchema(child) || partialObjectValid(child, entries, seen)) + ) + return false; + const branch = selectedConditionalBranch(schema, entries); + return branch === undefined || partialObjectValid(branch, entries, seen); + } finally { + seen.delete(schema); + } +} + +function objectNode(entries: ReadonlyArray<{ key: string; node: JsonNode }>): Extract { + const value: Record = Object.create(null) as Record; + for (const entry of entries) value[entry.key] = entry.node.value; + return { kind: 'object', value, entries: [...entries] }; +} + +function itemSchema(schema: Schema, index: number): Schema { + if (schema === true || schema === false || hasDeferredStructure(schema)) return schema === false ? false : true; + if (Array.isArray(schema.prefixItems) && isSchema(schema.prefixItems[index])) return schema.prefixItems[index]; + if (Array.isArray(schema.items)) { + if (isSchema(schema.items[index])) return schema.items[index]; + return isSchema(schema.additionalItems) ? schema.additionalItems : true; + } + return isSchema(schema.items) ? schema.items : true; +} + +function keyPrefixAllowed(state: JsonState, bytes: readonly number[]): boolean { + const frame = topObject(state); + if (frame.schema === true || frame.schema === false) return true; + const properties = encodedPropertyKeys(frame.schema); + if (properties === null) return true; + return properties.some(({ key, bytes: propertyBytes }) => { + if (frame.seen.has(key) || bytes.length > propertyBytes.length) return false; + for (let index = 0; index < bytes.length; ++index) { + if (bytes[index] !== propertyBytes[index]) return false; + } + return true; + }); +} + +function keyEscapeAllowed(state: JsonState): boolean { + const frame = topObject(state); + if (frame.schema === true || frame.schema === false) return true; + const properties = encodedPropertyKeys(frame.schema); + if (properties === null) return true; + if (state.highSurrogate !== undefined) return true; + const bytes = state.bytes ?? []; + return properties.some(({ key, bytes: propertyBytes }) => { + if (frame.seen.has(key) || propertyBytes.length <= bytes.length) return false; + for (let index = 0; index < bytes.length; ++index) { + if (bytes[index] !== propertyBytes[index]) return false; + } + const next = propertyBytes[bytes.length]; + return next < 0x20 || next === 0x22 || next === 0x5c || next >= 0x80; + }); +} + +function objectCanAddProperty(frame: ObjectFrame): boolean { + if (frame.schema === true || frame.schema === false) return true; + const properties = encodedPropertyKeys(frame.schema); + return properties === null || properties.some(({ key }) => !frame.seen.has(key)); +} + +function stringContentAllowed(schema: Schema, bytes: readonly number[], length: number): boolean { + const values = finiteStringValues(schema); + if ( + values !== null && + !values.some(({ bytes: valueBytes }) => { + if (bytes.length > valueBytes.length) return false; + for (let index = 0; index < bytes.length; ++index) { + if (bytes[index] !== valueBytes[index]) return false; + } + return true; + }) + ) + return false; + const maximum = directStringMaxLength(schema); + return maximum === undefined || length <= maximum; +} + +function stringEscapeAllowed(state: JsonState): boolean { + const values = finiteStringValues(state.schema); + const maximum = directStringMaxLength(state.schema); + if (maximum !== undefined && (state.stringLength ?? 0) >= maximum) return false; + if (values === null) return true; + if (state.highSurrogate !== undefined) return true; + const bytes = state.bytes ?? []; + return values.some(({ bytes: valueBytes }) => { + if (valueBytes.length <= bytes.length) return false; + for (let index = 0; index < bytes.length; ++index) { + if (bytes[index] !== valueBytes[index]) return false; + } + const next = valueBytes[bytes.length]; + return next < 0x20 || next === 0x22 || next === 0x5c || next >= 0x80; + }); +} + +function directStringMaxLength(schema: Schema): number | undefined { + return schema !== true && schema !== false && typeof schema.maxLength === 'number' ? schema.maxLength : undefined; +} + +function finiteStringValues(schema: Schema): ReadonlyArray<{ value: string; bytes: Uint8Array }> | null { + if (schema === true || schema === false) return null; + let values = finiteStringCache.get(schema); + if (values !== undefined) return values; + let candidates: string[] | null = null; + if (typeof schema.const === 'string') candidates = [schema.const]; + else if (Array.isArray(schema.enum)) + candidates = schema.enum.filter((value): value is string => typeof value === 'string'); + values = + candidates === null + ? null + : [...new Set(candidates)] + .filter((value) => validateNode(schema, { kind: 'string', value })) + .map((value) => ({ value, bytes: encoder.encode(value) })); + finiteStringCache.set(schema, values); + return values; +} + +function unicodeKeyPrefixAllowed(state: JsonState, hexPrefix: string): boolean { + const frame = topObject(state); + if (frame.schema === true || frame.schema === false) return true; + const properties = encodedPropertyKeys(frame.schema); + if (properties === null) return true; + const prefix = decodeString(state.bytes ?? []); + if (prefix === null) return true; + return properties.some(({ key }) => { + if (frame.seen.has(key) || !key.startsWith(prefix)) return false; + const offset = prefix.length; + let codeUnit = key.charCodeAt(offset); + if (state.highSurrogate !== undefined) { + if (codeUnit !== state.highSurrogate) return false; + codeUnit = key.charCodeAt(offset + 1); + } + return Number.isInteger(codeUnit) && codeUnit.toString(16).padStart(4, '0').startsWith(hexPrefix.toLowerCase()); + }); +} + +function encodedPropertyKeys( + schema: Record, +): ReadonlyArray<{ key: string; bytes: Uint8Array }> | null { + let properties = propertyKeyBytes.get(schema); + if (properties === undefined) { + const keys = constrainedPropertyKeys(schema, new Set()); + properties = keys === null ? null : [...keys].map((key) => ({ key, bytes: encoder.encode(key) })); + propertyKeyBytes.set(schema, properties); + } + return properties; +} + +function constrainedPropertyKeys(schema: Schema, seen: Set): Set | null { + if (schema === true || schema === false || seen.has(schema)) return null; + seen.add(schema); + try { + const properties = isRecord(schema.properties) ? schema.properties : {}; + let result = + schema.additionalProperties === false && !isRecord(schema.patternProperties) + ? new Set(Object.keys(properties)) + : null; + if (typeof schema.$ref === 'string') { + result = intersectKeySets(result, constrainedPropertyKeys(resolveReference(schema, schema.$ref), seen)); + } + if (Array.isArray(schema.allOf)) { + for (const child of schema.allOf) { + if (isSchema(child)) result = intersectKeySets(result, constrainedPropertyKeys(child, seen)); + } + } + for (const keyword of ['anyOf', 'oneOf'] as const) { + if (!Array.isArray(schema[keyword])) continue; + let union: Set | null = new Set(); + for (const child of schema[keyword]) { + if (!isSchema(child) || !schemaMayAcceptKind(child, 'object', new Set())) continue; + const childKeys = constrainedPropertyKeys(child, seen); + if (childKeys === null) { + union = null; + break; + } + for (const key of childKeys) union.add(key); + } + result = intersectKeySets(result, union); + } + return result; + } finally { + seen.delete(schema); + } +} + +function intersectKeySets(left: Set | null, right: Set | null): Set | null { + if (left === null) return right; + if (right === null) return left; + return new Set([...left].filter((key) => right.has(key))); +} + +function exceedsMaxItems(frame: ArrayFrame): boolean { + return frame.schema !== true && + frame.schema !== false && + !hasDeferredStructure(frame.schema) && + typeof frame.schema.maxItems === 'number' + ? frame.length >= frame.schema.maxItems + : false; +} + +const schemaContexts = new WeakMap(); + +function validateNode(schema: Schema, node: JsonNode, active = new Map>()): boolean { + if (schema === true) return true; + if (schema === false) return false; + let nodes = active.get(schema); + if (nodes?.has(node)) return true; + if (nodes === undefined) { + nodes = new Set(); + active.set(schema, nodes); + } + nodes.add(node); + try { + if (typeof schema.$ref === 'string' && !validateNode(resolveReference(schema, schema.$ref), node, active)) + return false; + if (schema.type !== undefined) { + const types = Array.isArray(schema.type) ? schema.type.map(String) : [String(schema.type)]; + if (!types.some((type) => nodeHasType(node, type))) return false; + } + if (schema.const !== undefined && !nodeEqualsValue(node, schema.const)) return false; + if (Array.isArray(schema.enum) && !schema.enum.some((value) => nodeEqualsValue(node, value))) return false; + if ( + Array.isArray(schema.allOf) && + !schema.allOf.every((child) => isSchema(child) && validateNode(child, node, active)) + ) + return false; + if ( + Array.isArray(schema.anyOf) && + !schema.anyOf.some((child) => isSchema(child) && validateNode(child, node, active)) + ) + return false; + if ( + Array.isArray(schema.oneOf) && + schema.oneOf.filter((child) => isSchema(child) && validateNode(child, node, active)).length !== 1 + ) + return false; + if (isSchema(schema.not) && validateNode(schema.not, node, active)) return false; + if (isSchema(schema.if)) { + const branch = validateNode(schema.if, node, active) ? schema.then : schema.else; + if (isSchema(branch) && !validateNode(branch, node, active)) return false; + } + if (node.kind === 'number' && !validateNumericNode(schema, node)) return false; + if (node.kind === 'string' && !validateStringNode(schema, node.value)) return false; + if (node.kind === 'array' && !validateArrayNode(schema, node, active)) return false; + if (node.kind === 'object' && !validateObjectNode(schema, node, active)) return false; + return true; + } finally { + nodes.delete(node); + } +} + +function validateNumericNode(schema: Record, node: Extract): boolean { + const value = decimal(node.raw); + if (schema.minimum !== undefined && compareDecimal(value, decimal(String(schema.minimum))) < 0) return false; + if (schema.maximum !== undefined && compareDecimal(value, decimal(String(schema.maximum))) > 0) return false; + if (schema.exclusiveMinimum !== undefined && compareDecimal(value, decimal(String(schema.exclusiveMinimum))) <= 0) + return false; + if (schema.exclusiveMaximum !== undefined && compareDecimal(value, decimal(String(schema.exclusiveMaximum))) >= 0) + return false; + return schema.multipleOf === undefined || decimalMultiple(value, decimal(String(schema.multipleOf))); +} + +function validateStringNode(schema: Record, value: string): boolean { + const length = [...value].length; + if (typeof schema.minLength === 'number' && length < schema.minLength) return false; + if (typeof schema.maxLength === 'number' && length > schema.maxLength) return false; + if (typeof schema.pattern === 'string' && !new RegExp(schema.pattern, 'u').test(value)) return false; + return typeof schema.format !== 'string' || formatMatches(schema.format, value); +} + +function validateArrayNode( + schema: Record, + node: Extract, + active: Map>, +): boolean { + if (typeof schema.minItems === 'number' && node.items.length < schema.minItems) return false; + if (typeof schema.maxItems === 'number' && node.items.length > schema.maxItems) return false; + if (schema.uniqueItems === true) { + for (let index = 0; index < node.items.length; ++index) { + if (node.items.slice(0, index).some((other) => nodesEqual(other, node.items[index]))) return false; + } + } + const legacyTuple = Array.isArray(schema.items) ? schema.items : undefined; + const prefix = (Array.isArray(schema.prefixItems) ? schema.prefixItems : legacyTuple) ?? []; + for (let index = 0; index < Math.min(prefix.length, node.items.length); ++index) { + if (!isSchema(prefix[index]) || !validateNode(prefix[index], node.items[index], active)) return false; + } + const remaining = legacyTuple ? schema.additionalItems : schema.items; + if (isSchema(remaining)) { + for (let index = prefix.length; index < node.items.length; ++index) { + if (!validateNode(remaining, node.items[index], active)) return false; + } + } + if (isSchema(schema.contains)) { + const matches = node.items.filter((item) => validateNode(schema.contains as Schema, item, active)).length; + const minimum = typeof schema.minContains === 'number' ? schema.minContains : 1; + if (matches < minimum || (typeof schema.maxContains === 'number' && matches > schema.maxContains)) return false; + } + return true; +} + +function validateObjectNode( + schema: Record, + node: Extract, + active: Map>, +): boolean { + const entries = new Map(node.entries.map((entry) => [entry.key, entry.node])); + if (entries.size !== node.entries.length) return false; + if (typeof schema.minProperties === 'number' && entries.size < schema.minProperties) return false; + if (typeof schema.maxProperties === 'number' && entries.size > schema.maxProperties) return false; + const required = Array.isArray(schema.required) ? schema.required : []; + if (!required.every((key) => typeof key === 'string' && entries.has(key))) return false; + const properties = isRecord(schema.properties) ? schema.properties : {}; + const patterns = isRecord(schema.patternProperties) ? schema.patternProperties : {}; + for (const [key, child] of entries) { + if ( + isSchema(schema.propertyNames) && + !validateNode(schema.propertyNames, { kind: 'string', value: key }, active) + ) + return false; + let matched = false; + if (isSchema(properties[key])) { + matched = true; + if (!validateNode(properties[key], child, active)) return false; + } + for (const [pattern, patternSchema] of Object.entries(patterns)) { + if (new RegExp(pattern, 'u').test(key)) { + matched = true; + if (!isSchema(patternSchema) || !validateNode(patternSchema, child, active)) return false; + } + } + if (!matched) { + if (schema.additionalProperties === false) return false; + if (isSchema(schema.additionalProperties) && !validateNode(schema.additionalProperties, child, active)) + return false; + } + } + const dependentRequired = isRecord(schema.dependentRequired) ? schema.dependentRequired : {}; + for (const [key, dependencies] of Object.entries(dependentRequired)) { + if ( + entries.has(key) && + Array.isArray(dependencies) && + !dependencies.every((dependency) => typeof dependency === 'string' && entries.has(dependency)) + ) + return false; + } + const dependentSchemas = isRecord(schema.dependentSchemas) ? schema.dependentSchemas : {}; + for (const [key, dependency] of Object.entries(dependentSchemas)) { + if (entries.has(key) && (!isSchema(dependency) || !validateNode(dependency, node, active))) return false; + } + const dependencies = isRecord(schema.dependencies) ? schema.dependencies : {}; + for (const [key, dependency] of Object.entries(dependencies)) { + if (!entries.has(key)) continue; + if (Array.isArray(dependency)) { + if (!dependency.every((requiredKey) => typeof requiredKey === 'string' && entries.has(requiredKey))) + return false; + } else if (!isSchema(dependency) || !validateNode(dependency, node, active)) return false; + } + return true; +} + +function nodeHasType(node: JsonNode, type: string): boolean { + if (type === 'integer') return node.kind === 'number' && decimalInteger(decimal(node.raw)); + if (type === 'number') return node.kind === 'number'; + return node.kind === type; +} + +function nodeEqualsValue(node: JsonNode, value: unknown): boolean { + if (node.kind === 'number') + return typeof value === 'number' && compareDecimal(decimal(node.raw), decimal(String(value))) === 0; + if (node.kind === 'array') + return ( + Array.isArray(value) && + node.items.length === value.length && + node.items.every((item, index) => nodeEqualsValue(item, value[index])) + ); + if (node.kind === 'object') { + if (!isRecord(value)) return false; + const entries = new Map(node.entries.map((entry) => [entry.key, entry.node])); + return ( + entries.size === Object.keys(value).length && + Object.entries(value).every(([key, child]) => entries.has(key) && nodeEqualsValue(entries.get(key)!, child)) + ); + } + return Object.is(node.value, value); +} + +function nodesEqual(left: JsonNode, right: JsonNode): boolean { + if (left.kind !== right.kind) return false; + if (left.kind === 'number' && right.kind === 'number') + return compareDecimal(decimal(left.raw), decimal(right.raw)) === 0; + if (left.kind === 'array' && right.kind === 'array') + return ( + left.items.length === right.items.length && + left.items.every((item, index) => nodesEqual(item, right.items[index])) + ); + if (left.kind === 'object' && right.kind === 'object') { + const rightEntries = new Map(right.entries.map((entry) => [entry.key, entry.node])); + return ( + left.entries.length === rightEntries.size && + left.entries.every( + (entry) => rightEntries.has(entry.key) && nodesEqual(entry.node, rightEntries.get(entry.key)!), + ) + ); + } + return Object.is(left.value, right.value); +} + +type Decimal = { coefficient: bigint; exponent: bigint }; + +function decimal(raw: string): Decimal { + const match = /^(-?)(\d+)(?:\.(\d+))?(?:[eE]([+-]?\d+))?$/.exec(raw)!; + const fraction = match[3] ?? ''; + let digits = `${match[2]}${fraction}`.replace(/^0+/, ''); + if (digits === '') return { coefficient: 0n, exponent: 0n }; + const trailing = /0+$/.exec(digits)?.[0].length ?? 0; + if (trailing > 0) digits = digits.slice(0, -trailing); + return { + coefficient: BigInt(`${match[1]}${digits}`), + exponent: BigInt(match[4] ?? '0') - BigInt(fraction.length) + BigInt(trailing), + }; +} + +function compareDecimal(left: Decimal, right: Decimal): number { + if (left.coefficient === right.coefficient && left.exponent === right.exponent) return 0; + if (left.coefficient === 0n) return right.coefficient < 0n ? 1 : -1; + if (right.coefficient === 0n) return left.coefficient < 0n ? -1 : 1; + if (left.coefficient < 0n && right.coefficient >= 0n) return -1; + if (left.coefficient >= 0n && right.coefficient < 0n) return 1; + const negative = left.coefficient < 0n; + const leftAbsolute = left.coefficient < 0n ? -left.coefficient : left.coefficient; + const rightAbsolute = right.coefficient < 0n ? -right.coefficient : right.coefficient; + const leftMagnitude = BigInt(leftAbsolute.toString().length) + left.exponent; + const rightMagnitude = BigInt(rightAbsolute.toString().length) + right.exponent; + if (leftMagnitude !== rightMagnitude) { + const comparison = leftMagnitude < rightMagnitude ? -1 : 1; + return negative ? -comparison : comparison; + } + const exponent = left.exponent < right.exponent ? left.exponent : right.exponent; + const scaledLeft = left.coefficient * 10n ** (left.exponent - exponent); + const scaledRight = right.coefficient * 10n ** (right.exponent - exponent); + return scaledLeft < scaledRight ? -1 : scaledLeft > scaledRight ? 1 : 0; +} + +function decimalInteger(value: Decimal): boolean { + return value.coefficient === 0n || value.exponent >= 0n; +} + +function decimalMultiple(value: Decimal, divisor: Decimal): boolean { + if (divisor.coefficient === 0n) return false; + if (value.coefficient === 0n) return true; + const numerator = value.coefficient < 0n ? -value.coefficient : value.coefficient; + const denominator = divisor.coefficient < 0n ? -divisor.coefficient : divisor.coefficient; + const exponent = value.exponent - divisor.exponent; + if (exponent >= 0n) { + return ((numerator % denominator) * modularPower(10n, exponent, denominator)) % denominator === 0n; + } + const places = -exponent; + if (places >= BigInt(numerator.toString().length)) return false; + return numerator % (denominator * 10n ** places) === 0n; +} + +function modularPower(base: bigint, exponent: bigint, modulus: bigint): bigint { + if (modulus === 1n) return 0n; + let result = 1n; + base %= modulus; + while (exponent > 0n) { + if (exponent & 1n) result = (result * base) % modulus; + base = (base * base) % modulus; + exponent >>= 1n; + } + return result; +} + +function hasDeferredStructure(schema: Schema): boolean { + return ( + schema !== true && + schema !== false && + ['$ref', 'allOf', 'anyOf', 'oneOf', 'not', 'if'].some((key) => schema[key] !== undefined) + ); +} + +function schemaMayAcceptKind(schema: Schema, kind: Kind, seen: Set): boolean { + if (schema === true) return true; + if (schema === false) return false; + if (seen.has(schema)) return true; + seen.add(schema); + if (typeof schema.$ref === 'string' && !schemaMayAcceptKind(resolveReference(schema, schema.$ref), kind, seen)) + return false; + if (!allowsKind(schema, kind)) return false; + if ( + Array.isArray(schema.allOf) && + !schema.allOf.every((child) => isSchema(child) && schemaMayAcceptKind(child, kind, seen)) + ) + return false; + if ( + Array.isArray(schema.anyOf) && + !schema.anyOf.some((child) => isSchema(child) && schemaMayAcceptKind(child, kind, seen)) + ) + return false; + if ( + Array.isArray(schema.oneOf) && + !schema.oneOf.some((child) => isSchema(child) && schemaMayAcceptKind(child, kind, seen)) + ) + return false; + return true; +} + +function resolveReference(owner: object, reference: string): Schema { + if (!reference.startsWith('#')) + throw new TypeError(`External JSON Schema reference ${JSON.stringify(reference)} is unsupported.`); + const root = schemaContexts.get(owner); + if (root === undefined) throw new Error('Missing JSON Schema context.'); + if (reference === '#') return root; + if (!reference.startsWith('#/')) { + throw new TypeError(`JSON Schema reference ${JSON.stringify(reference)} must contain a JSON Pointer.`); + } + let current: unknown = root; + for (const encoded of reference.slice(2).split('/')) { + const key = decodeURIComponent(encoded).replace(/~1/g, '/').replace(/~0/g, '~'); + if (!isRecord(current) || !(key in current)) + throw new TypeError(`JSON Schema reference ${JSON.stringify(reference)} does not resolve.`); + current = current[key]; + } + if (!isSchema(current)) + throw new TypeError(`JSON Schema reference ${JSON.stringify(reference)} does not resolve to a schema.`); + return current; +} + +function registerSchemaContext(schema: Schema, root: Schema, seen = new Set()): void { + if (!isRecord(schema) || seen.has(schema)) return; + seen.add(schema); + schemaContexts.set(schema, root); + for (const child of schemaChildren(schema)) registerSchemaContext(child, root, seen); +} + +function schemaChildren(schema: Record): Schema[] { + const result: Schema[] = []; + for (const key of [ + 'not', + 'if', + 'then', + 'else', + 'contains', + 'propertyNames', + 'additionalProperties', + 'additionalItems', + ]) { + if (isSchema(schema[key])) result.push(schema[key] as Schema); + } + if (isSchema(schema.items)) result.push(schema.items); + if (Array.isArray(schema.items)) result.push(...schema.items.filter(isSchema)); + for (const key of ['prefixItems', 'allOf', 'anyOf', 'oneOf']) { + if (Array.isArray(schema[key])) result.push(...(schema[key] as unknown[]).filter(isSchema)); + } + for (const key of ['properties', 'patternProperties', 'dependentSchemas', '$defs', 'definitions']) { + if (isRecord(schema[key])) + result.push(...Object.values(schema[key] as Record).filter(isSchema)); + } + if (isRecord(schema.dependencies)) { + result.push(...Object.values(schema.dependencies).filter(isSchema)); + } + return result; +} + +function checkReferences(schema: Schema, seen = new Set()): void { + if (!isRecord(schema) || seen.has(schema)) return; + seen.add(schema); + if (typeof schema.$ref === 'string') resolveReference(schema, schema.$ref); + for (const child of schemaChildren(schema)) checkReferences(child, seen); +} + +function assertPattern(value: unknown, path: string): void { + if (typeof value !== 'string') throw new TypeError(`${path} must be a string.`); + try { + new RegExp(value, 'u'); + } catch { + throw new TypeError(`${path} must be a valid Unicode RegExp.`); + } +} + +function guidanceFrom(schema: Schema): Guidance { + if (!isRecord(schema) || schema['x-guidance'] === undefined) return DEFAULT_GUIDANCE; + if (!isRecord(schema['x-guidance'])) throw new TypeError('x-guidance must be an object.'); + const source = schema['x-guidance']; + for (const key of Object.keys(source)) { + if (!['item_separator', 'key_separator', 'whitespace_flexible'].includes(key)) { + throw new TypeError(`Unsupported x-guidance option ${JSON.stringify(key)}.`); + } + } + const itemSeparator = source.item_separator ?? ','; + const keySeparator = source.key_separator ?? ':'; + const whitespaceFlexible = source.whitespace_flexible ?? true; + if (typeof itemSeparator !== 'string' || itemSeparator.length === 0) { + throw new TypeError('x-guidance.item_separator must be a non-empty string.'); + } + if (typeof keySeparator !== 'string' || keySeparator.length === 0) { + throw new TypeError('x-guidance.key_separator must be a non-empty string.'); + } + if (typeof whitespaceFlexible !== 'boolean') + throw new TypeError('x-guidance.whitespace_flexible must be a boolean.'); + return { + itemSeparator, + keySeparator, + itemBytes: encoder.encode(itemSeparator), + keyBytes: encoder.encode(keySeparator), + whitespaceFlexible, + }; +} + +function formatMatches(format: string, value: string): boolean { + switch (format) { + case 'date': + return validDate(value); + case 'time': + return validTime(value); + case 'date-time': { + const separator = value.search(/[Tt]/); + return separator > 0 && validDate(value.slice(0, separator)) && validTime(value.slice(separator + 1)); + } + case 'duration': + return /^(?:P\d+W|P(?=\d|T\d)(?:\d+Y)?(?:\d+M)?(?:\d+D)?(?:T(?=\d)(?:\d+H)?(?:\d+M)?(?:\d+S)?)?)$/.test( + value, + ); + case 'email': { + const at = value.lastIndexOf('@'); + if (at <= 0 || at === value.length - 1) return false; + const local = value.slice(0, at); + const domain = value.slice(at + 1); + if ( + local.startsWith('.') || + local.endsWith('.') || + local.includes('..') || + !/^[A-Za-z0-9!#$%&'*+\-/=?^_`{|}~.]+$/.test(local) + ) + return false; + return /^\[(.+)\]$/.test(domain) ? validIpv4(domain.slice(1, -1)) : validHostname(domain); + } + case 'hostname': + return validHostname(value); + case 'ipv4': + return validIpv4(value); + case 'ipv6': + return validIpv6(value); + case 'uuid': + return /^[\da-f]{8}(?:-[\da-f]{4}){3}-[\da-f]{12}$/i.test(value); + case 'uri': + return validUriText(value, true); + case 'uri-reference': + return validUriText(value, false); + case 'regex': + try { + new RegExp(value, 'u'); + return true; + } catch { + return false; + } + case 'json-pointer': + return /^(?:\/(?:[^~/]|~[01])*)*$/u.test(value); + case 'relative-json-pointer': + return /^(?:0|[1-9]\d*)(?:#|(?:\/(?:[^~/]|~[01])*)*)$/u.test(value); + default: + return true; + } +} + +function validDate(value: string): boolean { + const match = /^(\d{4})-(\d{2})-(\d{2})$/.exec(value); + if (!match) return false; + const year = Number(match[1]); + const month = Number(match[2]); + const day = Number(match[3]); + const days = [ + 31, + year % 4 === 0 && (year % 100 !== 0 || year % 400 === 0) ? 29 : 28, + 31, + 30, + 31, + 30, + 31, + 31, + 30, + 31, + 30, + 31, + ]; + return month >= 1 && month <= 12 && day >= 1 && day <= days[month - 1]; +} + +function validTime(value: string): boolean { + const match = /^(\d{2}):(\d{2}):(\d{2})(?:\.\d+)?(?:[Zz]|[+-](\d{2}):(\d{2}))$/.exec(value); + if (!match) return false; + const hour = Number(match[1]); + const minute = Number(match[2]); + const second = Number(match[3]); + if (hour > 23 || minute > 59 || second > 60 || (second === 60 && (hour !== 23 || minute !== 59))) return false; + return match[4] === undefined || (Number(match[4]) <= 23 && Number(match[5]) <= 59); +} + +function validHostname(value: string): boolean { + return ( + value.length > 0 && + value.length <= 253 && + !value.endsWith('.') && + value.split('.').every((label) => /^[A-Za-z0-9](?:[A-Za-z0-9-]{0,61}[A-Za-z0-9])?$/.test(label)) + ); +} + +function validIpv4(value: string): boolean { + const parts = value.split('.'); + return parts.length === 4 && parts.every((part) => /^(?:0|[1-9]\d{0,2})$/.test(part) && Number(part) <= 255); +} + +function validIpv6(value: string): boolean { + if (value.includes(':::') || value.split('::').length > 2) return false; + const compressed = value.includes('::'); + const groups = value.split(':').filter(Boolean); + if ( + !groups.every((group, index) => + group.includes('.') ? index === groups.length - 1 && validIpv4(group) : /^[\da-f]{1,4}$/i.test(group), + ) + ) + return false; + const count = groups.reduce((total, group) => total + (group.includes('.') ? 2 : 1), 0); + return compressed ? count < 8 : count === 8; +} + +function validUriText(value: string, absolute: boolean): boolean { + if (/[\s\\<>^`{|}\0]/.test(value) || /%(?![\da-f]{2})/i.test(value)) return false; + let rest = value; + const scheme = /^[A-Za-z][A-Za-z0-9+.-]*:/.exec(rest); + if (absolute && !scheme) return false; + if (scheme) rest = rest.slice(scheme[0].length); + const fragment = rest.indexOf('#'); + if (fragment !== -1) { + if (!uriQuery.test(rest.slice(fragment + 1))) return false; + rest = rest.slice(0, fragment); + } + const query = rest.indexOf('?'); + if (query !== -1) { + if (!uriQuery.test(rest.slice(query + 1))) return false; + rest = rest.slice(0, query); + } + if (rest.startsWith('//')) { + const slash = rest.indexOf('/', 2); + const authority = rest.slice(2, slash === -1 ? undefined : slash); + const path = slash === -1 ? '' : rest.slice(slash); + return validAuthority(authority) && uriPath.test(path); + } + if (!uriPath.test(rest)) return false; + if (scheme || rest === '' || rest.startsWith('/')) return true; + return !rest.slice(0, rest.indexOf('/') === -1 ? undefined : rest.indexOf('/')).includes(':'); +} + +const uriEncoded = '%[\\da-fA-F]{2}'; +const uriPchar = `(?:[A-Za-z0-9._~!$&'()*+,;=:@-]|${uriEncoded})`; +const uriPath = new RegExp(`^(?:${uriPchar}|/)*$`); +const uriQuery = new RegExp(`^(?:${uriPchar}|[/?])*$`); +const uriRegName = new RegExp(`^(?:[A-Za-z0-9._~!$&'()*+,;=-]|${uriEncoded})*$`); +const uriUserInfo = new RegExp(`^(?:[A-Za-z0-9._~!$&'()*+,;=:-]|${uriEncoded})*$`); + +function validAuthority(authority: string): boolean { + const at = authority.lastIndexOf('@'); + if (at !== -1) { + if (authority.indexOf('@') !== at || !uriUserInfo.test(authority.slice(0, at))) return false; + authority = authority.slice(at + 1); + } + if (authority.startsWith('[')) { + const close = authority.indexOf(']'); + if (close === -1) return false; + const host = authority.slice(1, close); + const suffix = authority.slice(close + 1); + return ( + (validIpv6(host) || /^v[\da-f]+\.[A-Za-z0-9._~!$&'()*+,;=:-]+$/i.test(host)) && + (suffix === '' || /^:\d*$/.test(suffix)) + ); + } + const colon = authority.lastIndexOf(':'); + let host = authority; + if (colon !== -1) { + if (authority.indexOf(':') !== colon || !/^\d*$/.test(authority.slice(colon + 1))) return false; + host = authority.slice(0, colon); + } + return uriRegName.test(host); +} + +function checkSchema(schema: Schema, path: string): void { + if (typeof schema === 'boolean') return; + if (!isRecord(schema)) throw new TypeError(`${path} must be a boolean or object schema.`); + const supported = new Set([ + 'type', + 'const', + 'enum', + 'properties', + 'required', + 'additionalProperties', + 'minProperties', + 'maxProperties', + 'items', + 'prefixItems', + 'minItems', + 'maxItems', + 'uniqueItems', + 'minLength', + 'maxLength', + 'pattern', + 'format', + 'minimum', + 'maximum', + 'exclusiveMinimum', + 'exclusiveMaximum', + 'multipleOf', + 'contains', + 'minContains', + 'maxContains', + 'patternProperties', + 'propertyNames', + 'dependentRequired', + 'dependentSchemas', + 'dependencies', + 'anyOf', + 'allOf', + 'oneOf', + 'not', + 'if', + 'then', + 'else', + '$ref', + '$defs', + 'definitions', + 'additionalItems', + 'x-guidance', + 'title', + 'description', + '$schema', + '$id', + '$comment', + 'default', + 'examples', + 'deprecated', + 'readOnly', + 'writeOnly', + 'contentEncoding', + 'contentMediaType', + 'contentSchema', + ]); + for (const key of Object.keys(schema)) { + if (!supported.has(key)) + throw new TypeError(`${path}: unsupported JSON Schema keyword ${JSON.stringify(key)}.`); + } + const configuredTypes = schema.type === undefined ? [] : Array.isArray(schema.type) ? schema.type : [schema.type]; + if ( + configuredTypes.some( + (type) => !['null', 'boolean', 'number', 'integer', 'string', 'array', 'object'].includes(String(type)), + ) + ) { + throw new TypeError(`${path}.type contains an unsupported JSON type.`); + } + for (const keyword of [ + 'not', + 'if', + 'then', + 'else', + 'contains', + 'propertyNames', + 'additionalProperties', + 'additionalItems', + ]) { + if (schema[keyword] !== undefined && !isSchema(schema[keyword])) { + throw new TypeError(`${path}.${keyword} must be a boolean or object schema.`); + } + } + for (const keyword of ['allOf', 'anyOf', 'oneOf', 'prefixItems']) { + const value = schema[keyword]; + if ( + value !== undefined && + (!Array.isArray(value) || + (keyword !== 'prefixItems' && value.length === 0) || + value.some((child) => !isSchema(child))) + ) { + throw new TypeError(`${path}.${keyword} must be an array of schemas.`); + } + } + for (const keyword of ['properties', 'patternProperties', 'dependentSchemas', '$defs', 'definitions']) { + const value = schema[keyword]; + if (value !== undefined && (!isRecord(value) || Object.values(value).some((child) => !isSchema(child)))) { + throw new TypeError(`${path}.${keyword} must be an object containing schemas.`); + } + } + if ( + schema.items !== undefined && + !isSchema(schema.items) && + !(Array.isArray(schema.items) && schema.items.every(isSchema)) + ) { + throw new TypeError(`${path}.items must be a schema or array of schemas.`); + } + if ( + schema.required !== undefined && + (!Array.isArray(schema.required) || + schema.required.some((key) => typeof key !== 'string') || + new Set(schema.required).size !== schema.required.length) + ) { + throw new TypeError(`${path}.required must be an array of unique strings.`); + } + if (schema.dependentRequired !== undefined) { + if (!isRecord(schema.dependentRequired)) throw new TypeError(`${path}.dependentRequired must be an object.`); + for (const dependency of Object.values(schema.dependentRequired)) { + if ( + !Array.isArray(dependency) || + dependency.some((key) => typeof key !== 'string') || + new Set(dependency).size !== dependency.length + ) { + throw new TypeError(`${path}.dependentRequired values must be arrays of unique strings.`); + } + } + } + if (schema.dependencies !== undefined && !isRecord(schema.dependencies)) { + throw new TypeError(`${path}.dependencies must be an object.`); + } + if (schema.$ref !== undefined && typeof schema.$ref !== 'string') + throw new TypeError(`${path}.$ref must be a string.`); + if (schema.enum !== undefined && (!Array.isArray(schema.enum) || schema.enum.length === 0)) { + throw new TypeError(`${path}.enum must be a non-empty array.`); + } + for (const keyword of ['minimum', 'maximum', 'exclusiveMinimum', 'exclusiveMaximum', 'multipleOf']) { + const value = schema[keyword]; + if ( + value !== undefined && + (typeof value !== 'number' || !Number.isFinite(value) || (keyword === 'multipleOf' && value <= 0)) + ) { + throw new TypeError( + `${path}.${keyword} must be ${keyword === 'multipleOf' ? 'a positive' : 'a finite'} number.`, + ); + } + } + if (schema.pattern !== undefined) assertPattern(schema.pattern, `${path}.pattern`); + if (isRecord(schema.patternProperties)) { + for (const pattern of Object.keys(schema.patternProperties)) + assertPattern(pattern, `${path}.patternProperties`); + } + if (schema.uniqueItems !== undefined && typeof schema.uniqueItems !== 'boolean') { + throw new TypeError(`${path}.uniqueItems must be a boolean.`); + } + for (const keyword of [ + 'minItems', + 'maxItems', + 'minContains', + 'maxContains', + 'minProperties', + 'maxProperties', + 'minLength', + 'maxLength', + ]) { + const value = schema[keyword]; + if (value !== undefined && (!Number.isInteger(value) || (value as number) < 0)) { + throw new TypeError(`${path}.${keyword} must be a non-negative integer.`); + } + } + if (schema['x-guidance'] !== undefined && path !== '$') { + throw new TypeError('x-guidance is only supported on the root schema.'); + } + schemaChildren(schema).forEach((child, index) => checkSchema(child, `${path}.schema[${index}]`)); +} + +function allowsKind(schema: Record, kind: Kind): boolean { + if ( + schema.const !== undefined && + valueKind(schema.const) !== kind && + !(kind === 'number' && valueKind(schema.const) === 'integer') + ) + return false; + if ( + Array.isArray(schema.enum) && + !schema.enum.some((value) => valueKind(value) === kind || (kind === 'number' && valueKind(value) === 'integer')) + ) + return false; + if (schema.type === undefined) return true; + const types = Array.isArray(schema.type) ? schema.type : [schema.type]; + return ( + types.includes(kind) || + (kind === 'number' && types.includes('integer')) || + (kind === 'integer' && types.includes('number')) + ); +} + +function valueKind(value: unknown): Kind | undefined { + if (value === null) return 'null'; + if (Array.isArray(value)) return 'array'; + if (isRecord(value)) return 'object'; + if (typeof value === 'number') return Number.isInteger(value) ? 'integer' : 'number'; + if (['boolean', 'string'].includes(typeof value)) return typeof value as Kind; + return undefined; +} + +function byteKind(byte: number): Kind | undefined { + if (byte === 0x22) return 'string'; + if (byte === 0x7b) return 'object'; + if (byte === 0x5b) return 'array'; + if (byte === 0x74 || byte === 0x66) return 'boolean'; + if (byte === 0x6e) return 'null'; + if (byte === 0x2d || isDigit(byte)) return 'number'; + return undefined; +} + +function numberPrefixValid(value: string): boolean { + // The exponent may only start after at least one fraction digit — "1.e" + // would be a syntactically extendable prefix that can never complete. + return ( + /^-?(?:0|[1-9]\d*)(?:\.\d+)?(?:[eE][+-]?\d*)?$/.test(value) || + /^-?(?:0|[1-9]\d*)\.$/.test(value) || + value === '-' + ); +} + +function numberComplete(value: string): boolean { + return /^-?(?:0|[1-9]\d*)(?:\.\d+)?(?:[eE][+-]?\d+)?$/.test(value); +} + +function isNumberByte(byte: number): boolean { + return isDigit(byte) || byte === 0x2e || byte === 0x45 || byte === 0x65 || byte === 0x2b || byte === 0x2d; +} + +function decodeString(bytes: readonly number[]): string | null { + try { + return decoder.decode(Uint8Array.from(bytes)); + } catch { + return null; + } +} + +function topObject(state: JsonState): ObjectFrame { + return state.stack.at(-1) as ObjectFrame; +} + +function replaceTop(stack: readonly Frame[], frame: Frame): Frame[] { + return [...stack.slice(0, -1), frame]; +} + +function allowsWhitespace(mode: Mode): boolean { + return ['value', 'array-value', 'object-key', 'colon', 'item-separator', 'after-value', 'done'].includes(mode); +} + +function isWhitespace(byte: number): boolean { + return byte === 0x09 || byte === 0x0a || byte === 0x0d || byte === 0x20; +} + +function isDigit(byte: number): boolean { + return byte >= 0x30 && byte <= 0x39; +} + +function isHex(byte: number): boolean { + return isDigit(byte) || (byte >= 0x41 && byte <= 0x46) || (byte >= 0x61 && byte <= 0x66); +} + +function isSchema(value: unknown): value is Schema { + return typeof value === 'boolean' || isRecord(value); +} + +function isRecord(value: unknown): value is Record { + return value !== null && typeof value === 'object' && !Array.isArray(value); +} diff --git a/packages/transformers-response-constraint/src/engine/regex.ts b/packages/transformers-response-constraint/src/engine/regex.ts new file mode 100644 index 000000000..c03725af9 --- /dev/null +++ b/packages/transformers-response-constraint/src/engine/regex.ts @@ -0,0 +1,419 @@ +import type { ConstraintState } from './types'; + +type Expression = + | { kind: 'empty' } + | { kind: 'set'; bytes: Uint32Array } + | { kind: 'choice'; choices: Expression[] } + | { kind: 'sequence'; parts: Expression[] } + | { kind: 'star'; child: Expression }; + +const EMPTY: Expression = { kind: 'empty' }; +const encoder = new TextEncoder(); + +export function compileRegex(source: string): ConstraintState { + const machine = new RegexMachine(new RegexParser(source).parse()); + return { + initial: machine.initial, + transition: (state, byte) => machine.transition(state, byte), + viable: (state) => state >= 0, + accepting: (state) => machine.accepting(state), + maskKey: (state) => machine.stateKey(state), + }; +} + +class RegexParser { + private index = 0; + + constructor(private readonly source: string) {} + + parse(): Expression { + if (this.source.startsWith('^')) this.index++; + const expression = this.alternation(); + if (this.peek() === '$' && this.index === this.source.length - 1) this.index++; + if (this.index !== this.source.length) this.fail(`unexpected ${JSON.stringify(this.peek())}`); + return expression; + } + + private alternation(): Expression { + const choices = [this.concatenation()]; + while (this.peek() === '|') { + this.index++; + choices.push(this.concatenation()); + } + return choice(choices); + } + + private concatenation(): Expression { + const parts: Expression[] = []; + while (this.index < this.source.length && this.peek() !== ')' && this.peek() !== '|') { + if (this.peek() === '$' && this.index === this.source.length - 1) break; + parts.push(this.quantified()); + } + return sequence(parts); + } + + private quantified(): Expression { + const atom = this.atom(); + const quantifier = this.peek(); + if (!['*', '+', '?', '{'].includes(quantifier ?? '')) return atom; + let minimum: number; + let maximum: number | null; + if (quantifier === '*') { + this.index++; + [minimum, maximum] = [0, null]; + } else if (quantifier === '+') { + this.index++; + [minimum, maximum] = [1, null]; + } else if (quantifier === '?') { + this.index++; + [minimum, maximum] = [0, 1]; + } else { + [minimum, maximum] = this.bounds(); + } + if (this.peek() === '?' || this.peek() === '+') this.fail('lazy and possessive quantifiers are unsupported'); + const parts = Array.from({ length: minimum }, () => atom); + if (maximum === null) parts.push(star(atom)); + else for (let count = minimum; count < maximum; ++count) parts.push(choice([EMPTY, atom])); + return sequence(parts); + } + + private atom(): Expression { + const character = this.peek(); + if (character === undefined) this.fail('expected an expression'); + if (character === '(') { + this.index++; + if (this.source.startsWith('?:', this.index)) this.index += 2; + else if (this.peek() === '?') this.fail('only non-capturing special groups are supported'); + const result = this.alternation(); + if (this.peek() !== ')') this.fail('unterminated group'); + this.index++; + return result; + } + if (character === '[') return this.characterClass(); + if (character === '.') { + this.index++; + return byteSet(range(0, 255)); + } + if (character === '\\') return this.escape(false).expression; + if ('*+?{})'.includes(character)) this.fail(`unexpected ${JSON.stringify(character)}`); + this.index += character.length; + return literal(character); + } + + private characterClass(): Expression { + this.index++; + const negated = this.peek() === '^'; + if (negated) this.index++; + const bytes = new Uint32Array(8); + let hasValue = false; + while (this.peek() !== ']') { + if (this.peek() === undefined) this.fail('unterminated character class'); + const first = this.classValue(); + if (this.peek() === '-' && this.source[this.index + 1] !== ']') { + this.index++; + const last = this.classValue(); + if (first.single === undefined || last.single === undefined || first.single > last.single) { + this.fail('invalid character class range'); + } + addRange(bytes, first.single, last.single); + } else { + union(bytes, first.bytes); + } + hasValue = true; + } + this.index++; + if (!hasValue) this.fail('empty character class'); + if (negated) for (let word = 0; word < bytes.length; ++word) bytes[word] = ~bytes[word]; + return byteSet(bytes); + } + + private classValue(): { bytes: Uint32Array; single?: number } { + if (this.peek() === '\\') { + const escaped = this.escape(true); + if (escaped.bytes === undefined) this.fail('multi-byte escapes are unsupported in character classes'); + return { bytes: escaped.bytes, single: escaped.single }; + } + const character = this.peek()!; + this.index += character.length; + const encoded = encoder.encode(character); + if (encoded.length !== 1) this.fail('non-ASCII character classes are unsupported'); + return { bytes: singleton(encoded[0]), single: encoded[0] }; + } + + private escape(inClass: boolean): { expression: Expression; bytes?: Uint32Array; single?: number } { + this.index++; + const code = this.peek(); + if (code === undefined) this.fail('trailing escape'); + this.index++; + if ('dDsSwW'.includes(code)) { + const bytes = shorthand(code.toLowerCase()); + if (code === code.toUpperCase()) for (let word = 0; word < bytes.length; ++word) bytes[word] = ~bytes[word]; + return { expression: byteSet(bytes), bytes }; + } + if (code === 'b' && !inClass) this.fail('word boundaries are unsupported'); + let value: number; + if (code === 'x') value = this.hex(2); + else if (code === 'u') value = this.hex(4); + else + value = + ({ n: 10, r: 13, t: 9, f: 12, v: 11, b: 8 } as Record)[code] ?? code.codePointAt(0)!; + const text = String.fromCodePoint(value); + const encoded = encoder.encode(text); + const bytes = encoded.length === 1 ? singleton(encoded[0]) : undefined; + return { expression: literal(text), bytes, single: encoded.length === 1 ? encoded[0] : undefined }; + } + + private bounds(): [number, number | null] { + this.index++; + const minimum = this.decimal(); + let maximum: number | null = minimum; + if (this.peek() === ',') { + this.index++; + maximum = this.peek() === '}' ? null : this.decimal(); + } + if (this.peek() !== '}') this.fail('unterminated repetition'); + this.index++; + if (minimum > 1000 || (maximum !== null && (maximum < minimum || maximum > 1000))) { + this.fail('invalid or excessive repetition'); + } + return [minimum, maximum]; + } + + private decimal(): number { + const start = this.index; + while (/\d/.test(this.peek() ?? '')) this.index++; + if (start === this.index) this.fail('expected a repetition count'); + return Number(this.source.slice(start, this.index)); + } + + private hex(length: number): number { + const value = this.source.slice(this.index, this.index + length); + if (!new RegExp(`^[\\da-f]{${length}}$`, 'i').test(value)) this.fail('invalid hexadecimal escape'); + this.index += length; + return Number.parseInt(value, 16); + } + + private peek(): string | undefined { + return this.source[this.index]; + } + + private fail(message: string): never { + throw new SyntaxError(`Invalid regex at index ${this.index}: ${message}.`); + } +} + +const OP_SET = 0; +const OP_SPLIT = 1; +const OP_JUMP = 2; +const OP_MATCH = 3; +const UNKNOWN = -2; + +type Fragment = { start: number; outs: number[] }; + +class NfaBuilder { + readonly ops: number[] = []; + readonly out1: number[] = []; + readonly out2: number[] = []; + readonly sets: Uint32Array[] = []; + + compile(expression: Expression): Fragment { + switch (expression.kind) { + case 'empty': { + const state = this.emit(OP_JUMP); + return { start: state, outs: [state << 1] }; + } + case 'set': { + const state = this.emit(OP_SET, -1, -1, expression.bytes); + return { start: state, outs: [state << 1] }; + } + case 'sequence': { + let result = this.compile(expression.parts[0]); + for (let index = 1; index < expression.parts.length; ++index) { + const next = this.compile(expression.parts[index]); + this.patch(result.outs, next.start); + result = { start: result.start, outs: next.outs }; + } + return result; + } + case 'choice': { + let result = this.compile(expression.choices[0]); + for (let index = 1; index < expression.choices.length; ++index) { + const right = this.compile(expression.choices[index]); + result = { + start: this.emit(OP_SPLIT, result.start, right.start), + outs: [...result.outs, ...right.outs], + }; + } + return result; + } + case 'star': { + const child = this.compile(expression.child); + const split = this.emit(OP_SPLIT, child.start); + this.patch(child.outs, split); + return { start: split, outs: [(split << 1) | 1] }; + } + } + } + + emit(op: number, first = -1, second = -1, set?: Uint32Array): number { + const state = this.ops.length; + this.ops.push(op); + this.out1.push(first); + this.out2.push(second); + this.sets.push(set ?? new Uint32Array(0)); + return state; + } + + patch(outs: number[], target: number): void { + for (const output of outs) { + if (output & 1) this.out2[output >>> 1] = target; + else this.out1[output >>> 1] = target; + } + } +} + +class RegexMachine { + readonly initial: number; + private readonly builder = new NfaBuilder(); + private readonly states: number[][] = []; + private readonly acceptingStates: boolean[] = []; + private readonly transitionTables: Int32Array[] = []; + private readonly stateIds = new Map(); + private readonly stateKeys: string[] = []; + + constructor(expression: Expression) { + const fragment = this.builder.compile(expression); + const match = this.builder.emit(OP_MATCH); + this.builder.patch(fragment.outs, match); + this.initial = this.intern(this.closure([fragment.start])); + } + + transition(state: number, byte: number): number { + if (state < 0) return -1; + const table = this.transitionTables[state]; + const cached = table[byte]; + if (cached !== UNKNOWN) return cached; + const seeds: number[] = []; + for (const pc of this.states[state]) { + if (this.builder.ops[pc] !== OP_SET) continue; + const set = this.builder.sets[pc]; + if (set[byte >>> 5] & (1 << (byte & 31))) seeds.push(this.builder.out1[pc]); + } + const next = seeds.length === 0 ? -1 : this.intern(this.closure(seeds)); + table[byte] = next; + return next; + } + + accepting(state: number): boolean { + return state >= 0 && this.acceptingStates[state]; + } + + // The interned NFA state set is intrinsic to the regex (unlike the interned + // ids, which depend on discovery order), so it is a stable mask-cache key + // across constraint instances compiled from the same source. + stateKey(state: number): string | undefined { + return state >= 0 ? this.stateKeys[state] : undefined; + } + + private closure(seeds: number[]): number[] { + const result: number[] = []; + const stack = [...seeds]; + const seen = new Set(); + while (stack.length > 0) { + const state = stack.pop()!; + if (state < 0 || seen.has(state)) continue; + seen.add(state); + const op = this.builder.ops[state]; + if (op === OP_SPLIT) { + stack.push(this.builder.out1[state], this.builder.out2[state]); + } else if (op === OP_JUMP) { + stack.push(this.builder.out1[state]); + } else { + result.push(state); + } + } + result.sort((left, right) => left - right); + return result; + } + + private intern(active: number[]): number { + const key = active.join(','); + const existing = this.stateIds.get(key); + if (existing !== undefined) return existing; + if (this.states.length >= 4096) throw new Error('Regex produced too many runtime states.'); + const id = this.states.length; + const transitions = new Int32Array(256); + transitions.fill(UNKNOWN); + this.states.push(active); + this.acceptingStates.push(active.some((state) => this.builder.ops[state] === OP_MATCH)); + this.transitionTables.push(transitions); + this.stateIds.set(key, id); + this.stateKeys.push(key); + return id; + } +} + +function choice(items: Expression[]): Expression { + const flattened = items.flatMap((item) => (item.kind === 'choice' ? item.choices : [item])); + if (flattened.length === 1) return flattened[0]; + return { kind: 'choice', choices: flattened }; +} + +function sequence(items: Expression[]): Expression { + const flattened = items + .flatMap((item) => (item.kind === 'sequence' ? item.parts : [item])) + .filter((item) => item !== EMPTY); + if (flattened.length === 0) return EMPTY; + if (flattened.length === 1) return flattened[0]; + return { kind: 'sequence', parts: flattened }; +} + +function star(child: Expression): Expression { + if (child === EMPTY) return EMPTY; + if (child.kind === 'star') return child; + return { kind: 'star', child }; +} + +function literal(value: string): Expression { + return sequence([...encoder.encode(value)].map((byte) => byteSet(singleton(byte)))); +} + +function byteSet(bytes: Uint32Array): Expression { + return { kind: 'set', bytes }; +} + +function shorthand(code: string): Uint32Array { + const bytes = new Uint32Array(8); + if (code === 'd' || code === 'w') addRange(bytes, 48, 57); + if (code === 'w') { + addRange(bytes, 65, 90); + addRange(bytes, 97, 122); + add(bytes, 95); + } + if (code === 's') for (const byte of [9, 10, 11, 12, 13, 32]) add(bytes, byte); + return bytes; +} + +function singleton(byte: number): Uint32Array { + const bytes = new Uint32Array(8); + add(bytes, byte); + return bytes; +} + +function range(first: number, last: number): Uint32Array { + const bytes = new Uint32Array(8); + addRange(bytes, first, last); + return bytes; +} + +function addRange(bytes: Uint32Array, first: number, last: number): void { + for (let byte = first; byte <= last; ++byte) add(bytes, byte); +} + +function add(bytes: Uint32Array, byte: number): void { + bytes[byte >>> 5] |= 1 << (byte & 31); +} + +function union(target: Uint32Array, source: Uint32Array): void { + for (let word = 0; word < target.length; ++word) target[word] |= source[word]; +} diff --git a/packages/transformers-response-constraint/src/engine/tokenizer.ts b/packages/transformers-response-constraint/src/engine/tokenizer.ts new file mode 100644 index 000000000..ed191029d --- /dev/null +++ b/packages/transformers-response-constraint/src/engine/tokenizer.ts @@ -0,0 +1,183 @@ +import type { TokenizerSource } from './types'; + +type RecordLike = Record; + +export type TokenizerData = { + tokens: Uint8Array[]; + eosTokenId: number; + specialTokenIds: Set; +}; + +const encoder = new TextEncoder(); +let byteLevelMap: Map | undefined; + +export function extractTokenizer(tokenizer: TokenizerSource): TokenizerData { + const source = asRecord(tokenizer, 'tokenizer'); + if (Array.isArray(source.tokens)) { + return normalizeDirectTokenizer(source); + } + + const tokenizerJson = getTokenizerJson(source); + const vocabulary = getVocabulary(source, tokenizerJson); + if (vocabulary === undefined) { + throw new TypeError('Could not extract the tokenizer vocabulary.'); + } + + const vocabularyTokens = Object.keys(vocabulary); + let size = 0; + for (const token of vocabularyTokens) { + const id = Number(vocabulary[token]); + if (!Number.isInteger(id) || id < 0) throw new TypeError(`Tokenizer has an invalid token ID for ${token}.`); + if (id + 1 > size) size = id + 1; + } + const tokenBytes = tokenBytesConverter(tokenizerJson); + const tokens = new Array(size); + for (const token of vocabularyTokens) { + tokens[Number(vocabulary[token])] = tokenBytes(token); + } + + const addedTokens = field(tokenizerJson, 'added_tokens'); + if (Array.isArray(addedTokens)) { + for (const added of addedTokens) { + if (!isRecord(added) || !Number.isInteger(added.id)) continue; + while (tokens.length <= Number(added.id)) tokens.push(undefined); + tokens[Number(added.id)] = encoder.encode(typeof added.content === 'string' ? added.content : ''); + } + } + for (let id = 0; id < tokens.length; ++id) { + if (tokens[id] === undefined) throw new Error(`Tokenizer vocabulary is missing token ID ${id}.`); + } + + const eosTokenId = tokenId(source, tokenizerJson, ['eos_token_id', 'eosTokenId', 'eos_token', 'eosToken']); + if (eosTokenId === undefined) throw new TypeError('Tokenizer does not expose an EOS token ID.'); + const specialTokenIds = new Set([eosTokenId]); + for (const value of [ + source.special_token_ids, + source.specialTokenIds, + source.all_special_ids, + source.allSpecialIds, + ]) { + if (Array.isArray(value)) for (const id of value) if (Number.isInteger(id)) specialTokenIds.add(Number(id)); + } + if (Array.isArray(addedTokens)) { + for (const added of addedTokens) { + if (isRecord(added) && added.special === true && Number.isInteger(added.id)) + specialTokenIds.add(Number(added.id)); + } + } + return { tokens: tokens as Uint8Array[], eosTokenId, specialTokenIds }; +} + +function normalizeDirectTokenizer(source: RecordLike): TokenizerData { + const configuredTokens = source.tokens as unknown[]; + const tokens = configuredTokens.map((token, id) => { + if (!(token instanceof Uint8Array) && !Array.isArray(token)) { + throw new TypeError(`Tokenizer token ${id} must be a byte array.`); + } + const values = Array.from(token as ArrayLike); + if (values.some((byte) => !Number.isInteger(byte) || byte < 0 || byte > 255)) + throw new TypeError(`Tokenizer token ${id} is invalid.`); + return Uint8Array.from(values); + }); + const eosTokenId = Number(source.eosTokenId ?? source.eos_token_id); + if (!Number.isInteger(eosTokenId) || eosTokenId < 0 || eosTokenId >= tokens.length) { + throw new TypeError('A valid eos_token_id is required with tokenizer tokens.'); + } + const configured = source.specialTokenIds ?? source.special_token_ids; + const specialTokenIds = new Set([eosTokenId]); + if (Array.isArray(configured)) for (const id of configured) specialTokenIds.add(Number(id)); + return { tokens, eosTokenId, specialTokenIds }; +} + +function getTokenizerJson(source: RecordLike): unknown { + const value = source._tokenizerJSON ?? source.tokenizerJSON ?? source.tokenizer_json; + return typeof value === 'string' ? JSON.parse(value) : value; +} + +function getVocabulary(source: RecordLike, tokenizerJson: unknown): RecordLike | undefined { + const modelVocabulary = field(field(tokenizerJson, 'model'), 'vocab'); + if (isRecord(modelVocabulary)) return modelVocabulary; + for (const name of ['get_vocab', 'getVocab']) { + const method = source[name]; + if (typeof method !== 'function') continue; + const value = method.call(source, true); + if (value instanceof Map) return Object.fromEntries(value); + if (isRecord(value)) return value; + } + return isRecord(source.vocab) ? source.vocab : undefined; +} + +function tokenBytesConverter(tokenizerJson: unknown): (token: string) => Uint8Array { + if ( + hasComponent(field(tokenizerJson, 'decoder'), 'ByteLevel') || + hasComponent(field(tokenizerJson, 'pre_tokenizer'), 'ByteLevel') + ) { + const map = getByteLevelMap(); + return (token) => { + const fallback = /^<0x([\da-f]{2})>$/i.exec(token); + if (fallback) return Uint8Array.of(Number.parseInt(fallback[1], 16)); + const bytes: number[] = []; + for (const character of token) { + const byte = map.get(character); + if (byte === undefined) bytes.push(...encoder.encode(character)); + else bytes.push(byte); + } + return Uint8Array.from(bytes); + }; + } + const modelType = field(field(tokenizerJson, 'model'), 'type'); + const sentencePiece = modelType === 'Unigram' || modelType === 'SentencePiece'; + const prefix = field(field(tokenizerJson, 'model'), 'continuing_subword_prefix'); + return (token) => { + const fallback = /^<0x([\da-f]{2})>$/i.exec(token); + if (fallback) return Uint8Array.of(Number.parseInt(fallback[1], 16)); + if (sentencePiece || token.includes('▁')) return encoder.encode(token.replaceAll('▁', ' ')); + return encoder.encode( + typeof prefix === 'string' && token.startsWith(prefix) ? token.slice(prefix.length) : token, + ); + }; +} + +function getByteLevelMap(): Map { + if (byteLevelMap !== undefined) return byteLevelMap; + const visible = new Set(); + for (let code = 33; code <= 126; ++code) visible.add(code); + for (let code = 161; code <= 172; ++code) visible.add(code); + for (let code = 174; code <= 255; ++code) visible.add(code); + let extra = 0; + byteLevelMap = new Map(); + for (let byte = 0; byte < 256; ++byte) { + byteLevelMap.set(String.fromCharCode(visible.has(byte) ? byte : 256 + extra++), byte); + } + return byteLevelMap; +} + +function tokenId(source: RecordLike, tokenizerJson: unknown, keys: string[]): number | undefined { + const vocabulary = getVocabulary(source, tokenizerJson); + for (const key of keys) { + const value = source[key] ?? field(tokenizerJson, key); + if (Number.isInteger(value)) return Number(value); + if (typeof value === 'string' && Number.isInteger(vocabulary?.[value])) return Number(vocabulary![value]); + } + return undefined; +} + +function hasComponent(value: unknown, type: string): boolean { + if (Array.isArray(value)) return value.some((item) => hasComponent(item, type)); + return isRecord(value) && (value.type === type || Object.values(value).some((item) => hasComponent(item, type))); +} + +function field(value: unknown, key: string): unknown { + return isRecord(value) ? value[key] : undefined; +} + +function asRecord(value: unknown, name: string): RecordLike { + if ((typeof value !== 'object' && typeof value !== 'function') || value === null) { + throw new TypeError(`${name} must be an object.`); + } + return value as RecordLike; +} + +function isRecord(value: unknown): value is RecordLike { + return value !== null && typeof value === 'object' && !Array.isArray(value); +} diff --git a/packages/transformers-response-constraint/src/engine/types.ts b/packages/transformers-response-constraint/src/engine/types.ts new file mode 100644 index 000000000..e8a18352a --- /dev/null +++ b/packages/transformers-response-constraint/src/engine/types.ts @@ -0,0 +1,12 @@ +export type JSONValue = null | boolean | number | string | JSONValue[] | { [key: string]: JSONValue }; +export type JSONSchema = boolean | { [key: string]: JSONValue }; +export type TokenizerSource = object | ((...args: unknown[]) => unknown); + +export interface ConstraintState { + readonly initial: State; + transition(state: State, byte: number): State; + viable(state: State): boolean; + accepting(state: State): boolean; + stringCapacity?(state: State): number | undefined; + maskKey?(state: State): string | undefined; +} diff --git a/packages/transformers-response-constraint/src/index.ts b/packages/transformers-response-constraint/src/index.ts new file mode 100644 index 000000000..5910fe9b6 --- /dev/null +++ b/packages/transformers-response-constraint/src/index.ts @@ -0,0 +1,2 @@ +export { ResponseConstraint } from './ResponseConstraint'; +export type { ResponseFormat } from './ResponseConstraint'; diff --git a/packages/transformers-response-constraint/src/utils/mask.ts b/packages/transformers-response-constraint/src/utils/mask.ts new file mode 100644 index 000000000..dcdedce25 --- /dev/null +++ b/packages/transformers-response-constraint/src/utils/mask.ts @@ -0,0 +1,30 @@ +import { type Tensor } from '@huggingface/transformers'; + +type LogitsData = Float32Array | Float64Array | number[]; + +export function applyMask(logits: Tensor, mask: Uint32Array, vocabSize: number): void { + const data = logits.data as LogitsData; + const stride = logits.dims.at(-1)!; + if (vocabSize > stride) { + throw new Error(`Constraint vocabulary size ${vocabSize} exceeds logits vocabulary size ${stride}.`); + } + for (let offset = 0; offset < data.length; offset += stride) { + const fullWords = vocabSize >>> 5; + for (let word = 0; word < fullWords; ++word) { + const bits = mask[word] | 0; + if (bits === -1) continue; + const start = offset + (word << 5); + if (bits === 0) { + data.fill(-Infinity, start, start + 32); + } else { + for (let bit = 0; bit < 32; ++bit) { + if (!(bits & (1 << bit))) data[start + bit] = -Infinity; + } + } + } + for (let tokenId = fullWords << 5; tokenId < vocabSize; ++tokenId) { + if (!(mask[tokenId >>> 5] & (1 << (tokenId & 31)))) data[offset + tokenId] = -Infinity; + } + data.fill(-Infinity, offset + vocabSize, offset + stride); + } +} diff --git a/packages/transformers-response-constraint/tests/corpus/helpers.mjs b/packages/transformers-response-constraint/tests/corpus/helpers.mjs new file mode 100644 index 000000000..d0f33e0e7 --- /dev/null +++ b/packages/transformers-response-constraint/tests/corpus/helpers.mjs @@ -0,0 +1,43 @@ +import assert from "node:assert/strict"; + +import { loadLLGuidance } from "./runtime.mjs"; + +const encoder = new TextEncoder(); + +export const llguidance = await loadLLGuidance(); +export const tokenizer = new llguidance.LLTokenizer({}); + +export function tokenAllowed(mask, token) { + return Boolean(mask?.[token >>> 5] & (1 << (token & 31))); +} + +export function resultField(result, field) { + return result instanceof Map ? result.get(field) : result?.[field]; +} + +export function grammarAccepts(grammar, input) { + const matcher = new llguidance.LLMatcher(tokenizer, grammar); + for (const token of encoder.encode(input)) { + const result = matcher._interpreter.computeMask(); + if (!tokenAllowed(resultField(result, "mask"), token)) return false; + matcher._interpreter.commitToken(token); + } + const result = matcher._interpreter.computeMask(); + return tokenAllowed(resultField(result, "mask"), tokenizer.eos_token) || Boolean(resultField(result, "stop")); +} + +export function jsonSchemaAccepts(schema, value, jsonText) { + return grammarAccepts(llguidance.LLMatcher.grammar_from_json_schema(schema), jsonText ?? JSON.stringify(value)); +} + +export function assertJsonSchema(schema, value, expected, jsonText) { + assert.equal(jsonSchemaAccepts(schema, value, jsonText), expected, `${jsonText ?? JSON.stringify(value)} should ${expected ? "match" : "not match"} ${JSON.stringify(schema)}`); +} + +export function compileJsonSchema(schema) { + return new llguidance.LLMatcher(tokenizer, llguidance.LLMatcher.grammar_from_json_schema(schema)); +} + +export function assertJsonSchemaCompileError(schema, message) { + assert.throws(() => compileJsonSchema(schema), message ? (error) => String(error).includes(message) : undefined); +} diff --git a/packages/transformers-response-constraint/tests/corpus/loader.mjs b/packages/transformers-response-constraint/tests/corpus/loader.mjs new file mode 100644 index 000000000..f1bbd5472 --- /dev/null +++ b/packages/transformers-response-constraint/tests/corpus/loader.mjs @@ -0,0 +1,14 @@ +import { pathToFileURL } from "node:url"; + +const runtime = pathToFileURL(new URL("./runtime.mjs", import.meta.url).pathname).href; +const helpers = pathToFileURL(new URL("./helpers.mjs", import.meta.url).pathname).href; + +export async function resolve(specifier, context, nextResolve) { + if (context.parentURL?.includes("/llguidance-js/test/upstream/") && specifier === "./helpers.mjs") { + return { url: helpers, shortCircuit: true }; + } + if (context.parentURL?.includes("/llguidance-js/test/") && specifier.endsWith("/src/llguidance.node.ts")) { + return { url: runtime, shortCircuit: true }; + } + return nextResolve(specifier, context); +} diff --git a/packages/transformers-response-constraint/tests/corpus/run.mjs b/packages/transformers-response-constraint/tests/corpus/run.mjs new file mode 100644 index 000000000..ceae5244e --- /dev/null +++ b/packages/transformers-response-constraint/tests/corpus/run.mjs @@ -0,0 +1,12 @@ +import { spawnSync } from "node:child_process"; +import { dirname, resolve } from "node:path"; +import { fileURLToPath } from "node:url"; + +const here = dirname(fileURLToPath(import.meta.url)); +const corpusRoot = process.env.LLGUIDANCE_JS_PATH ?? resolve(here, "../../../../../llguidance-js"); +const files = ["test/upstream/json-primitives.test.mjs", "test/upstream/json-structures.test.mjs", "test/upstream/json-combinations.test.mjs", "test/upstream/json-string-format.test.mjs", "test/upstream/json-x-guidance.test.mjs", "test/json-schema-2020-profile.test.mjs"].map((file) => resolve(corpusRoot, file)); +const original = process.argv.includes("--original"); +const unsupported = ["resolves general local JSON Pointers, escaped segments, and nested defs", "rejects malformed shapes for the production profile at compile time", "flexible separators with spaces: pretty_print$", "flexible separators with spaces: pretty_print_2$"].join("|"); +const args = [...(original ? [] : ["--experimental-loader", resolve(here, "loader.mjs"), "--test-skip-pattern", unsupported]), "--test", ...files]; +const result = spawnSync(process.execPath, args, { stdio: "inherit" }); +process.exit(result.status ?? 1); diff --git a/packages/transformers-response-constraint/tests/corpus/runtime.mjs b/packages/transformers-response-constraint/tests/corpus/runtime.mjs new file mode 100644 index 000000000..ad68792da --- /dev/null +++ b/packages/transformers-response-constraint/tests/corpus/runtime.mjs @@ -0,0 +1,113 @@ +import { Tensor } from "@huggingface/transformers"; + +import { ResponseConstraint } from "../../dist/index.js"; + +const DEFAULT_EOS = 261; +const DEFAULT_TOKENIZER = { + tokens: [...Array.from({ length: 256 }, (_, byte) => [byte]), ...Array.from({ length: 6 }, () => [])], + eos_token_id: DEFAULT_EOS, + special_token_ids: [256, 257, 258, 259, 260, 261], +}; + +export async function loadLLGuidance() { + return runtime; +} + +export const loadBundledLLGuidance = loadLLGuidance; + +class LLTokenizer { + constructor(config = {}) { + this._tokenizer = normalizeTokenizer(config); + this.eos_token = this._tokenizer.eos_token_id; + this.eos_tokens = [this.eos_token]; + this.vocab_size = this._tokenizer.tokens.length; + } + + dispose() {} + free() {} +} + +class LLMatcher { + static grammar_from_json_schema(schema) { + return typeof schema === "string" ? JSON.parse(schema) : schema; + } + + constructor(tokenizer, schema) { + this._interpreter = createInterpreter({ + tokenizer: unwrapTokenizer(tokenizer), + response_format: { type: "json_schema", json_schema: schema }, + }); + } + + dispose() {} + free() {} +} + +function createTokenizer(config = {}) { + return normalizeTokenizer(config); +} + +function createInterpreter({ tokenizer, response_format }) { + const source = unwrapTokenizer(tokenizer); + const constraint = ResponseConstraint.fromResponseFormat(source, response_format); + const vocabSize = source.tokens.length; + let inputIds = [0n]; + + return { + computeMask() { + const logits = new Tensor("float32", new Float32Array(vocabSize), [1, vocabSize]); + try { + constraint.logits_processor([inputIds], logits); + } catch { + return { stop: true, reason: "dead_end" }; + } + const mask = new Uint32Array(Math.ceil(vocabSize / 32)); + for (let token = 0; token < vocabSize; ++token) { + if (Number.isFinite(logits.data[token])) mask[token >>> 5] |= 1 << (token & 31); + } + return { mask, vocabSize }; + }, + computeMaskInto(target) { + const result = this.computeMask(); + target.fill(0); + if ("mask" in result) { + target.set(result.mask); + return { mask: target, vocabSize }; + } + return result; + }, + commitToken(token) { + inputIds = [...inputIds, BigInt(token)]; + constraint.logits_processor.onTokensSampled([token], [inputIds]); + return { + stop: constraint.stopping_criteria([inputIds])[0], + backtrack: 0, + ffTokens: [], + }; + }, + dispose() {}, + free() {}, + }; +} + +function normalizeTokenizer(config) { + if (!config.tokens) return DEFAULT_TOKENIZER; + const eos = config.eosTokenId ?? config.eos_token_id; + return { + ...config, + eos_token_id: eos, + special_token_ids: config.specialTokenIds ?? config.special_token_ids ?? [eos], + }; +} + +function unwrapTokenizer(tokenizer) { + return tokenizer?._tokenizer ?? tokenizer ?? DEFAULT_TOKENIZER; +} + +const runtime = { + LLTokenizer, + LLMatcher, + createTokenizer, + createInterpreter, + get_version: () => "llguidance-ts-core", +}; diff --git a/packages/transformers-response-constraint/tests/response-constraint.test.js b/packages/transformers-response-constraint/tests/response-constraint.test.js new file mode 100644 index 000000000..aa132c4a6 --- /dev/null +++ b/packages/transformers-response-constraint/tests/response-constraint.test.js @@ -0,0 +1,506 @@ +import { Tensor } from "@huggingface/transformers"; + +import { ResponseConstraint } from "../dist/index.js"; + +const EOS_TOKEN_ID = 256; +const tokenizer = { + tokens: [...Array.from({ length: EOS_TOKEN_ID }, (_, tokenId) => [tokenId]), []], + eos_token_id: EOS_TOKEN_ID, + special_token_ids: [EOS_TOKEN_ID], +}; + +function logits() { + return new Tensor("float32", new Float32Array(EOS_TOKEN_ID + 1).fill(1), [1, EOS_TOKEN_ID + 1]); +} + +function isAllowed(scores, tokenId) { + return Number.isFinite(scores.data[tokenId]); +} + +async function consume(constraint, text) { + const inputIds = [0n]; + for (const tokenId of new TextEncoder().encode(text)) { + const scores = logits(); + constraint.logits_processor([inputIds], scores); + expect(isAllowed(scores, tokenId)).toBe(true); + inputIds.push(BigInt(tokenId)); + constraint.logits_processor.onTokensSampled([tokenId], [inputIds]); + expect(constraint.stopping_criteria([inputIds])).toEqual([false]); + } + return inputIds; +} + +function schemaAccepts(schema, text) { + const constraint = ResponseConstraint.fromResponseFormat(tokenizer, { + type: "json_schema", + json_schema: schema, + }); + const inputIds = [0n]; + for (const tokenId of new TextEncoder().encode(text)) { + const scores = logits(); + try { + constraint.logits_processor([inputIds], scores); + } catch { + return false; + } + if (!isAllowed(scores, tokenId)) return false; + inputIds.push(BigInt(tokenId)); + constraint.logits_processor.onTokensSampled([tokenId], [inputIds]); + } + const scores = logits(); + try { + constraint.logits_processor([inputIds], scores); + } catch { + return false; + } + return isAllowed(scores, EOS_TOKEN_ID); +} + +describe("ResponseConstraint", () => { + it("applies a regex mask", async () => { + const constraint = await ResponseConstraint.fromResponseFormat(tokenizer, { type: "regex", regex: "[ac]" }); + const scores = logits(); + + constraint.logits_processor([[0n]], scores); + + expect(isAllowed(scores, "a".charCodeAt(0))).toBe(true); + expect(isAllowed(scores, "b".charCodeAt(0))).toBe(false); + expect(isAllowed(scores, "c".charCodeAt(0))).toBe(true); + expect(isAllowed(scores, EOS_TOKEN_ID)).toBe(false); + }); + + it("accepts JSON that satisfies a schema", async () => { + const constraint = await ResponseConstraint.fromResponseFormat(tokenizer, { + type: "json_schema", + json_schema: { + type: "object", + properties: { answer: { type: "string", minLength: 1 } }, + required: ["answer"], + additionalProperties: false, + }, + }); + const inputIds = await consume(constraint, '{"answer":"yes"}'); + const scores = logits(); + + constraint.logits_processor([inputIds], scores); + expect(isAllowed(scores, EOS_TOKEN_ID)).toBe(true); + constraint.logits_processor.onTokensSampled([EOS_TOKEN_ID], [[...inputIds, BigInt(EOS_TOKEN_ID)]]); + + expect(constraint.stopping_criteria([[...inputIds, BigInt(EOS_TOKEN_ID)]])).toEqual([true]); + }); + + it("applies schema structure while producing JSON", async () => { + const constraint = await ResponseConstraint.fromResponseFormat(tokenizer, { + type: "json_schema", + json_schema: { + type: "object", + properties: { answer: { type: "string" } }, + required: ["answer"], + additionalProperties: false, + }, + }); + const initial = logits(); + constraint.logits_processor([[0n]], initial); + expect(isAllowed(initial, "{".charCodeAt(0))).toBe(true); + expect(isAllowed(initial, "[".charCodeAt(0))).toBe(false); + + const inputIds = await consume(constraint, "{"); + const afterOpen = logits(); + constraint.logits_processor([inputIds], afterOpen); + expect(isAllowed(afterOpen, "}".charCodeAt(0))).toBe(false); + }); + + it("rejects impossible property prefixes and invalid scalar endings", async () => { + const constraint = ResponseConstraint.fromResponseFormat(tokenizer, { + type: "json_schema", + json_schema: { + type: "object", + properties: { answer: { type: "string", pattern: "^yes$" } }, + required: ["answer"], + additionalProperties: false, + }, + }); + const propertyInput = await consume(constraint, '{"'); + const propertyScores = logits(); + constraint.logits_processor([propertyInput], propertyScores); + expect(isAllowed(propertyScores, "a".charCodeAt(0))).toBe(true); + expect(isAllowed(propertyScores, "z".charCodeAt(0))).toBe(false); + + const valueInput = await consume(constraint, 'answer":"x'); + const valueScores = logits(); + constraint.logits_processor([valueInput], valueScores); + expect(isAllowed(valueScores, '"'.charCodeAt(0))).toBe(false); + + const unicodeProperty = { + type: "object", + properties: { "😀": { type: "string" } }, + required: ["😀"], + additionalProperties: false, + }; + expect(schemaAccepts(unicodeProperty, '{"😀":"yes"}')).toBe(true); + expect(schemaAccepts(unicodeProperty, '{"😁":"yes"}')).toBe(false); + + const escapedProperty = { + type: "object", + properties: { confidence: { type: "number" } }, + required: ["confidence"], + additionalProperties: false, + }; + expect(schemaAccepts(escapedProperty, '{"c\\u006fnfidence":1}')).toBe(false); + expect(schemaAccepts(escapedProperty, '{"c\\n\\n\\n":1}')).toBe(false); + const canonicalKeyConstraint = ResponseConstraint.fromResponseFormat(tokenizer, { + type: "json_schema", + json_schema: escapedProperty, + }); + const canonicalKeyInput = await consume(canonicalKeyConstraint, '{"confidence'); + const canonicalKeyScores = logits(); + canonicalKeyConstraint.logits_processor([canonicalKeyInput], canonicalKeyScores); + expect(isAllowed(canonicalKeyScores, '"'.charCodeAt(0))).toBe(true); + expect(isAllowed(canonicalKeyScores, "\\".charCodeAt(0))).toBe(false); + const completedPropertyInput = await consume(canonicalKeyConstraint, '":1'); + const completedPropertyScores = logits(); + canonicalKeyConstraint.logits_processor([completedPropertyInput], completedPropertyScores); + expect(isAllowed(completedPropertyScores, "}".charCodeAt(0))).toBe(true); + expect(isAllowed(completedPropertyScores, ",".charCodeAt(0))).toBe(false); + expect(schemaAccepts({ type: "object", properties: { "a\nb": true }, required: ["a\nb"], additionalProperties: false }, '{"a\\nb":1}')).toBe(true); + + const languageConstraint = ResponseConstraint.fromResponseFormat(tokenizer, { + type: "json_schema", + json_schema: { enum: ["en", "de", "fr", "es"] }, + }); + const languageInput = await consume(languageConstraint, '"'); + const languageScores = logits(); + languageConstraint.logits_processor([languageInput], languageScores); + expect(isAllowed(languageScores, "e".charCodeAt(0))).toBe(true); + expect(isAllowed(languageScores, "d".charCodeAt(0))).toBe(true); + expect(isAllowed(languageScores, "x".charCodeAt(0))).toBe(false); + expect(isAllowed(languageScores, "\\".charCodeAt(0))).toBe(false); + expect(schemaAccepts({ const: "\n" }, '"\\n"')).toBe(true); + + const boundedString = ResponseConstraint.fromResponseFormat(tokenizer, { + type: "json_schema", + json_schema: { type: "string", maxLength: 2 }, + }); + const boundedStringInput = await consume(boundedString, '"ab'); + const boundedStringScores = logits(); + boundedString.logits_processor([boundedStringInput], boundedStringScores); + expect(isAllowed(boundedStringScores, '"'.charCodeAt(0))).toBe(true); + expect(isAllowed(boundedStringScores, "c".charCodeAt(0))).toBe(false); + expect(isAllowed(boundedStringScores, "\\".charCodeAt(0))).toBe(false); + expect(schemaAccepts({ type: "string", maxLength: 1 }, '"😀"')).toBe(true); + expect(schemaAccepts({ type: "string", maxLength: 1 }, '"😀x"')).toBe(false); + + const composed = ResponseConstraint.fromResponseFormat(tokenizer, { + type: "json_schema", + json_schema: { + oneOf: [{ const: "general" }, { type: "object", properties: { role: { type: "string" } }, required: ["role"], additionalProperties: false }], + }, + }); + const composedInput = await consume(composed, '{"'); + const composedScores = logits(); + composed.logits_processor([composedInput], composedScores); + expect(isAllowed(composedScores, "r".charCodeAt(0))).toBe(true); + expect(isAllowed(composedScores, "z".charCodeAt(0))).toBe(false); + }); + + it("restricts integer fields to reachable canonical syntax", async () => { + const confidence = { + type: "object", + properties: { confidence: { type: "integer", minimum: 0, maximum: 100 } }, + required: ["confidence"], + additionalProperties: false, + }; + + // Integer fields follow llguidance's shape: fractions may only contain + // zeros and exponents may not be negative, so "0.9" (which would strand the + // model in states like "0.9e-" that can never close) is cut off at the "9". + const constraint = ResponseConstraint.fromResponseFormat(tokenizer, { + type: "json_schema", + json_schema: confidence, + }); + const input = await consume(constraint, '{"confidence":0.'); + const scores = logits(); + constraint.logits_processor([input], scores); + expect(isAllowed(scores, "0".charCodeAt(0))).toBe(true); + expect(isAllowed(scores, "9".charCodeAt(0))).toBe(false); + + const exponent = ResponseConstraint.fromResponseFormat(tokenizer, { + type: "json_schema", + json_schema: confidence, + }); + const exponentInput = await consume(exponent, '{"confidence":9e'); + const exponentScores = logits(); + exponent.logits_processor([exponentInput], exponentScores); + expect(isAllowed(exponentScores, "-".charCodeAt(0))).toBe(false); + // 9e1 = 90 fits [0, 100]; every exponent starting with 3 puts 9e3+ out of range + expect(isAllowed(exponentScores, "1".charCodeAt(0))).toBe(true); + expect(isAllowed(exponentScores, "3".charCodeAt(0))).toBe(false); + + expect(schemaAccepts(confidence, '{"confidence":15}')).toBe(true); + expect(schemaAccepts(confidence, '{"confidence":1.0}')).toBe(true); + expect(schemaAccepts(confidence, '{"confidence":9e1}')).toBe(true); + expect(schemaAccepts(confidence, '{"confidence":0.9}')).toBe(false); + expect(schemaAccepts(confidence, '{"confidence":9e-1}')).toBe(false); + + // Zero padding carries no information, so it is capped: a model stuck on + // "0" is eventually forced to close instead of streaming digits forever. + const padded = ResponseConstraint.fromResponseFormat(tokenizer, { + type: "json_schema", + json_schema: confidence, + }); + const paddedInput = await consume(padded, '{"confidence":95e000'); + const paddedScores = logits(); + padded.logits_processor([paddedInput], paddedScores); + expect(isAllowed(paddedScores, "0".charCodeAt(0))).toBe(false); + expect(isAllowed(paddedScores, "}".charCodeAt(0))).toBe(true); + + // Digits that could never get back into [0, 100] are pruned: after "15", + // any further digit forces 150+. + const bounded = ResponseConstraint.fromResponseFormat(tokenizer, { + type: "json_schema", + json_schema: confidence, + }); + const boundedInput = await consume(bounded, '{"confidence":15'); + const boundedScores = logits(); + bounded.logits_processor([boundedInput], boundedScores); + expect(isAllowed(boundedScores, "0".charCodeAt(0))).toBe(false); + expect(isAllowed(boundedScores, "}".charCodeAt(0))).toBe(true); + + // A first digit that cannot start any in-range integer is masked: 4, 40-49, + // 400+ all miss [50, 100], while 1 can still reach 100. + const range = ResponseConstraint.fromResponseFormat(tokenizer, { + type: "json_schema", + json_schema: { type: "integer", minimum: 50, maximum: 100 }, + }); + const rangeScores = logits(); + range.logits_processor([[0n]], rangeScores); + expect(isAllowed(rangeScores, "5".charCodeAt(0))).toBe(true); + expect(isAllowed(rangeScores, "1".charCodeAt(0))).toBe(true); + expect(isAllowed(rangeScores, "4".charCodeAt(0))).toBe(false); + expect(isAllowed(rangeScores, "-".charCodeAt(0))).toBe(false); + + const negative = { type: "integer", minimum: -50, maximum: -10 }; + expect(schemaAccepts(negative, "-25")).toBe(true); + const negativeConstraint = ResponseConstraint.fromResponseFormat(tokenizer, { + type: "json_schema", + json_schema: negative, + }); + const negativeInput = await consume(negativeConstraint, "-"); + const negativeScores = logits(); + negativeConstraint.logits_processor([negativeInput], negativeScores); + expect(isAllowed(negativeScores, "2".charCodeAt(0))).toBe(true); + expect(isAllowed(negativeScores, "6".charCodeAt(0))).toBe(false); + + // Integer-valued enums get the same protection. + expect(schemaAccepts({ enum: [1, 2, 30] }, "30")).toBe(true); + expect(schemaAccepts({ enum: [1, 2, 30] }, "1.0")).toBe(true); + expect(schemaAccepts({ enum: [1, 2, 30] }, "1.5")).toBe(false); + + // Plain number fields keep full JSON syntax. + const ratio = { type: "number", minimum: 0, maximum: 1 }; + expect(schemaAccepts(ratio, "0.9")).toBe(true); + expect(schemaAccepts(ratio, "9e-1")).toBe(true); + }); + + it("supports composition, conditionals, and local references", async () => { + expect(schemaAccepts({ not: { type: "string" } }, "42")).toBe(true); + expect(schemaAccepts({ not: { type: "string" } }, '"no"')).toBe(false); + expect(schemaAccepts({ allOf: [{ type: "integer", minimum: 2 }, { multipleOf: 2 }] }, "4")).toBe(true); + expect(schemaAccepts({ oneOf: [{ type: "string" }, { type: "integer" }] }, "2")).toBe(true); + + const conditional = { + type: "object", + properties: { kind: { enum: ["text", "count"] }, value: true }, + required: ["kind", "value"], + if: { properties: { kind: { const: "text" } }, required: ["kind"] }, + then: { properties: { value: { type: "string" } } }, + else: { properties: { value: { type: "integer" } } }, + }; + expect(schemaAccepts(conditional, '{"value":"ok","kind":"text"}')).toBe(true); + expect(schemaAccepts(conditional, '{"kind":"text","value":2}')).toBe(false); + + const followUp = { + type: "object", + properties: { needed: { type: "boolean" }, question: { type: ["string", "null"] } }, + required: ["needed", "question"], + additionalProperties: false, + allOf: [ + { + if: { properties: { needed: { const: true } }, required: ["needed"] }, + then: { properties: { question: { type: "string", minLength: 1 } } }, + else: { properties: { question: { type: "null" } } }, + }, + ], + }; + const followUpConstraint = ResponseConstraint.fromResponseFormat(tokenizer, { + type: "json_schema", + json_schema: followUp, + }); + const followUpInput = await consume(followUpConstraint, '{"needed":true,"question":'); + const followUpScores = logits(); + followUpConstraint.logits_processor([followUpInput], followUpScores); + expect(isAllowed(followUpScores, '"'.charCodeAt(0))).toBe(true); + expect(isAllowed(followUpScores, "n".charCodeAt(0))).toBe(false); + expect(schemaAccepts(followUp, '{"question":null,"needed":true}')).toBe(false); + + const referenced = { + $defs: { answer: { type: "integer", minimum: 2 } }, + $ref: "#/$defs/answer", + }; + expect(schemaAccepts(referenced, "3")).toBe(true); + expect(schemaAccepts(referenced, "1")).toBe(false); + + const objectUnion = { + anyOf: [ + { type: "object", properties: { a: { type: "string" }, b: { type: "integer" } }, required: ["a"], additionalProperties: false }, + { type: "object", properties: { a: { type: "string" }, c: { type: "number" } }, required: ["a"], additionalProperties: false }, + ], + }; + expect(schemaAccepts(objectUnion, '{"a":"x","b":2}')).toBe(true); + expect(schemaAccepts(objectUnion, '{"a":"x","b":2,"c":3}')).toBe(false); + }); + + it("supports deep equality and structural assertions", () => { + expect(schemaAccepts({ const: { name: "John", values: [1] } }, '{"values":[1.0],"name":"John"}')).toBe(true); + expect(schemaAccepts({ const: "😀" }, '"\\ud83d\\ude00"')).toBe(true); + expect(schemaAccepts({ const: "😀" }, '"\\ud83dx"')).toBe(false); + expect(schemaAccepts({ type: "array", uniqueItems: true }, '[{"a":[1]},{"a":[1.0]}]')).toBe(false); + expect(schemaAccepts({ type: "array", contains: { type: "integer", minimum: 2 }, minContains: 2, maxContains: 2 }, '["x",2,3]')).toBe(true); + expect(schemaAccepts({ type: "array", contains: { const: 1 } }, "[0]")).toBe(false); + }); + + it("supports property patterns and dependencies", () => { + const schema = { + type: "object", + properties: { code: { type: "integer" }, card: { type: "string" }, billing: { type: "string" } }, + patternProperties: { "^code$": { minimum: 1 }, "^x-": { type: "string" } }, + dependentRequired: { card: ["billing"] }, + additionalProperties: false, + }; + expect(schemaAccepts(schema, '{"x-note":"ok","code":2,"billing":"x","card":"1"}')).toBe(true); + expect(schemaAccepts(schema, '{"code":0}')).toBe(false); + expect(schemaAccepts(schema, '{"card":"1"}')).toBe(false); + }); + + it("supports recursive references and draft-07 compatibility", () => { + const linkedList = { + $defs: { + node: { + type: "object", + properties: { + value: { type: "string" }, + next: { anyOf: [{ $ref: "#/$defs/node" }, { type: "null" }] }, + }, + required: ["value", "next"], + additionalProperties: false, + }, + }, + $ref: "#/$defs/node", + }; + expect(schemaAccepts(linkedList, '{"value":"a","next":{"value":"b","next":null}}')).toBe(true); + expect(schemaAccepts(linkedList, '{"value":"a","next":2}')).toBe(false); + + const tuple = { + type: "array", + items: [{ type: "string" }, { type: "integer" }], + additionalItems: false, + }; + expect(schemaAccepts(tuple, '["x",1]')).toBe(true); + expect(schemaAccepts(tuple, '["x",1,true]')).toBe(false); + }); + + it("supports recognized formats and x-guidance separators", () => { + const formats = { + date: ["2024-02-29", "2023-02-29"], + time: ["23:59:60Z", "22:59:60Z"], + "date-time": ["2024-02-29T12:00:00Z", "2023-02-29T12:00:00Z"], + duration: ["P1Y2M3DT4H5M6S", "P1YT"], + email: ["user+tag@example.com", "user@@example.com"], + hostname: ["www.example.com", "-example.com"], + ipv4: ["255.255.255.255", "256.0.0.1"], + ipv6: ["2001:db8::1", "1::d6::42"], + uuid: ["98d80576-482e-427f-8434-7f86890ab222", "98d80576-482e-427f"], + uri: ["https://example.com/a%20b?x=1#part", "http://example.com/%GG"], + "uri-reference": ["../a/b?x=1#part", "a:b c"], + regex: ["^(?:😀|[a-z]+)$", "["], + "json-pointer": ["/a~1b/~0key", "/bad~escape"], + "relative-json-pointer": ["2/items/0", "01/value"], + }; + for (const [format, [valid, invalid]] of Object.entries(formats)) { + if (!schemaAccepts({ type: "string", format }, JSON.stringify(valid))) throw new Error(`Rejected valid ${format}`); + if (schemaAccepts({ type: "string", format }, JSON.stringify(invalid))) throw new Error(`Accepted invalid ${format}`); + } + + const guided = { + type: "object", + properties: { a: { type: "integer" }, b: { type: "integer" } }, + required: ["a", "b"], + additionalProperties: false, + "x-guidance": { item_separator: "-", key_separator: "_ ", whitespace_flexible: false }, + }; + expect(schemaAccepts(guided, '{"a"_ 1-"b"_ 2}')).toBe(true); + expect(schemaAccepts(guided, '{"a":1,"b":2}')).toBe(false); + }); + + it("rejects malformed and unsupported schema shapes", () => { + for (const schema of [{ not: [] }, { patternProperties: { "[": true } }, { dependentRequired: { a: ["b", "b"] } }, { minContains: -1 }, { uniqueItems: "yes" }, { unevaluatedProperties: false }]) { + expect(() => ResponseConstraint.fromResponseFormat(tokenizer, { type: "json_schema", json_schema: schema })).toThrow(); + } + }); + + it("supports unconstrained JSON objects", async () => { + const constraint = await ResponseConstraint.fromResponseFormat(tokenizer, { type: "json_object" }); + const inputIds = await consume(constraint, '{"nested":{"enabled":true},"count":2}'); + const scores = logits(); + + constraint.logits_processor([inputIds], scores); + + expect(isAllowed(scores, EOS_TOKEN_ID)).toBe(true); + }); + + it("reuses cached masks without sharing generation state", async () => { + const schema = { + type: "object", + properties: { answer: { enum: ["yes", "no"] } }, + required: ["answer"], + additionalProperties: false, + }; + const first = ResponseConstraint.fromResponseFormat(tokenizer, { + type: "json_schema", + json_schema: schema, + }); + const second = ResponseConstraint.fromResponseFormat(tokenizer, { + type: "json_schema", + json_schema: schema, + }); + + const firstIds = await consume(first, '{"answer":"yes"}'); + const secondIds = await consume(second, '{"answer":"no"}'); + const firstScores = logits(); + const secondScores = logits(); + first.logits_processor([firstIds], firstScores); + second.logits_processor([secondIds], secondScores); + + expect(isAllowed(firstScores, EOS_TOKEN_ID)).toBe(true); + expect(isAllowed(secondScores, EOS_TOKEN_ID)).toBe(true); + }); + + it("rejects batched generation", async () => { + const constraint = await ResponseConstraint.fromResponseFormat(tokenizer, { type: "json_object" }); + + expect(() => constraint.logits_processor([[0n], [0n]], new Tensor("float32", new Float32Array((EOS_TOKEN_ID + 1) * 2), [2, EOS_TOKEN_ID + 1]))).toThrow("currently supports batch size 1"); + }); + + it("rejects a sampled token outside the constraint", async () => { + const constraint = await ResponseConstraint.fromResponseFormat(tokenizer, { type: "regex", regex: "a" }); + constraint.logits_processor([[0n]], logits()); + + expect(() => constraint.logits_processor.onTokensSampled(["b".charCodeAt(0)], [[0n, 98n]])).toThrow("does not satisfy the constraint"); + }); + + it("returns only the generation hooks", async () => { + const constraint = await ResponseConstraint.fromResponseFormat(tokenizer, { type: "regex", regex: "a" }); + + expect(Object.keys(constraint).sort()).toEqual(["logits_processor", "stopping_criteria"]); + }); +}); diff --git a/packages/transformers-llguidance/tsconfig.json b/packages/transformers-response-constraint/tsconfig.json similarity index 100% rename from packages/transformers-llguidance/tsconfig.json rename to packages/transformers-response-constraint/tsconfig.json diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 5bb74e840..56ac2d997 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -58,11 +58,7 @@ importers: specifier: 5.9.3 version: 5.9.3 - packages/transformers-llguidance: - dependencies: - llguidance: - specifier: 0.2.0 - version: 0.2.0 + packages/transformers-response-constraint: devDependencies: '@huggingface/transformers': specifier: workspace:* @@ -1585,9 +1581,6 @@ packages: linkify-it@5.0.0: resolution: {integrity: sha512-5aHCbzQRADcdP+ATqnDuhhJ/MRIqDkZX5pyjFHRRysS8vZ5AbqGEoFIb6pYHPZ+L/OC2Lc+xT8uHVVR5CAK/wQ==} - llguidance@0.2.0: - resolution: {integrity: sha512-BvN0WlQLLZC9kkRpz7k4DoAVpplcaJKcKrRVI68tCwPZlWtrGv/kriFyjKwP7y0qBct4M7vmOWEqXJTDFsMB6w==} - locate-path@5.0.0: resolution: {integrity: sha512-t7hw9pI+WvuwNJXwk5zVHpyhIqzg2qTlklJOf0mVxGSbe3Fp2VieZcduNYjaLDoy6p9uGpQEGWG87WpMKlNq8g==} engines: {node: '>=8'} @@ -3750,8 +3743,6 @@ snapshots: dependencies: uc.micro: 2.1.0 - llguidance@0.2.0: {} - locate-path@5.0.0: dependencies: p-locate: 4.1.0 From 7c4593c6b43ad5975865184b3f79e17535b99eeb Mon Sep 17 00:00:00 2001 From: Nico Martin Date: Wed, 19 Aug 2026 15:25:58 +0200 Subject: [PATCH 15/17] clean up --- .../performance/bench-real.mjs | 150 ------------------ .../performance/complex-json.mjs | 65 -------- .../performance/run.mjs | 98 ------------ .../performance/simple-json.mjs | 16 -- .../performance/simple-regex.mjs | 8 - 5 files changed, 337 deletions(-) delete mode 100644 packages/transformers-response-constraint/performance/bench-real.mjs delete mode 100644 packages/transformers-response-constraint/performance/complex-json.mjs delete mode 100644 packages/transformers-response-constraint/performance/run.mjs delete mode 100644 packages/transformers-response-constraint/performance/simple-json.mjs delete mode 100644 packages/transformers-response-constraint/performance/simple-regex.mjs diff --git a/packages/transformers-response-constraint/performance/bench-real.mjs b/packages/transformers-response-constraint/performance/bench-real.mjs deleted file mode 100644 index d7b517813..000000000 --- a/packages/transformers-response-constraint/performance/bench-real.mjs +++ /dev/null @@ -1,150 +0,0 @@ -import { AutoTokenizer, Tensor } from "@huggingface/transformers"; -import { ResponseConstraint } from "/Users/nico/Documents/Dev/transformers.js/packages/transformers-response-constraint/dist/index.js"; - -const MODEL = process.argv[2] ?? "onnx-community/gemma-4-E2B-it-ONNX"; - -const simpleJsonSchemaFormat = { - type: "json_schema", - json_schema: { - "x-guidance": { - whitespace_flexible: false, - item_separator: ", ", - key_separator: ": ", - }, - type: "object", - properties: { - answer: { type: "string", minLength: 1, maxLength: 120 }, - }, - required: ["answer"], - additionalProperties: false, - }, -}; - -const complexJsonSchemaFormat = { - type: "json_schema", - json_schema: { - "x-guidance": { - whitespace_flexible: false, - item_separator: ", ", - key_separator: ": ", - }, - type: "object", - properties: { - answer: { - type: "object", - properties: { - text: { type: "string", minLength: 1, maxLength: 120 }, - tone: { enum: ["friendly", "formal", "playful"] }, - language: { enum: ["en", "de", "fr", "es"] }, - }, - required: ["text", "tone", "language"], - additionalProperties: false, - }, - alternatives: { - type: "array", - minItems: 1, - maxItems: 3, - items: { type: "string", minLength: 1, maxLength: 120 }, - }, - metadata: { - type: "object", - properties: { - confidence: { type: "integer", minimum: 0, maximum: 100 }, - safe: { type: "boolean" }, - tags: { - type: "array", - minItems: 1, - maxItems: 3, - items: { type: "string", minLength: 1, maxLength: 20 }, - }, - }, - required: ["confidence", "safe", "tags"], - additionalProperties: false, - }, - }, - required: ["answer", "alternatives", "metadata"], - additionalProperties: false, - }, -}; - -const regexFormat = { type: "regex", regex: "(Hello|Hi|Hey)( there)?[!.]" }; - -const cases = [ - { - name: "simple JSON", - format: simpleJsonSchemaFormat, - output: '{"answer": "Hello there! How can I help you today?"}', - }, - { name: "regex", format: regexFormat, output: "Hello there!" }, - { - name: "complex JSON", - format: complexJsonSchemaFormat, - output: - '{"answer": {"text": "Hello there! How can I help you today?", "tone": "friendly", "language": "en"}, "alternatives": ["Hi! What can I do for you?", "Hey, great to see you."], "metadata": {"confidence": 95, "safe": true, "tags": ["greeting", "friendly"]}}', - }, -]; - -const tokenizer = await AutoTokenizer.from_pretrained(MODEL); -const eosId = tokenizer.eos_token_id; -console.log("eos:", tokenizer.eos_token, eosId); - -const VOCAB = 262144; - -function percentile(values, p) { - const sorted = [...values].sort((a, b) => a - b); - return sorted[Math.floor(sorted.length * p)]; -} - -function run(format, tokenIds, label) { - const t0 = performance.now(); - const constraint = ResponseConstraint.fromResponseFormat(tokenizer, format); - const setupMs = performance.now() - t0; - - const scores = new Tensor("float32", new Float32Array(VOCAB), [1, VOCAB]); - const inputIds = [0n]; - const stepMs = []; - for (const tokenId of tokenIds) { - scores.data.fill(0); - const s = performance.now(); - constraint.logits_processor([inputIds], scores); - if (!Number.isFinite(scores.data[tokenId])) { - const txt = tokenizer.decode([tokenId]); - throw new Error(`${label} rejected token ${tokenId} (${JSON.stringify(txt)})`); - } - inputIds.push(BigInt(tokenId)); - constraint.logits_processor.onTokensSampled([tokenId], [inputIds]); - constraint.stopping_criteria([inputIds]); - stepMs.push(performance.now() - s); - } - return { setupMs, stepMs }; -} - -// EOS-terminated token sequence for each case -for (const c of cases) { - const ids = tokenizer.encode(c.output, { add_special_tokens: false }); - c.tokenIds = [...ids, Number(eosId)]; -} - -const results = []; -for (const c of cases) { - // Run 1: cold (first-ever fromResponseFormat pays tokenizer trie build; caches empty) - const cold = run(c.format, c.tokenIds, c.name); - // Run 2: warm-ish (mask caches populated from run 1) - const warm1 = run(c.format, c.tokenIds, c.name); - const warm2 = run(c.format, c.tokenIds, c.name); - results.push({ - name: c.name, - tokens: c.tokenIds.length, - "cold setup ms": cold.setupMs.toFixed(1), - "warm setup ms": warm2.setupMs.toFixed(2), - "cold sum ms": cold.stepMs.reduce((a, b) => a + b, 0).toFixed(1), - "cold med": percentile(cold.stepMs, 0.5).toFixed(2), - "cold p90": percentile(cold.stepMs, 0.9).toFixed(2), - "cold max": Math.max(...cold.stepMs).toFixed(2), - "warm sum ms": warm2.stepMs.reduce((a, b) => a + b, 0).toFixed(1), - "warm med": percentile(warm2.stepMs, 0.5).toFixed(3), - "warm p90": percentile(warm2.stepMs, 0.9).toFixed(2), - "warm max": Math.max(...warm2.stepMs).toFixed(2), - }); -} -console.table(results); diff --git a/packages/transformers-response-constraint/performance/complex-json.mjs b/packages/transformers-response-constraint/performance/complex-json.mjs deleted file mode 100644 index 3c49c120b..000000000 --- a/packages/transformers-response-constraint/performance/complex-json.mjs +++ /dev/null @@ -1,65 +0,0 @@ -export default { - name: "complex JSON schema", - responseFormat: { - type: "json_schema", - json_schema: { - type: "object", - properties: { - request_id: { type: "string", pattern: "^[a-z0-9-]{8,36}$" }, - status: { enum: ["queued", "running", "completed", "failed"] }, - user: { - type: "object", - properties: { - id: { type: "integer", minimum: 1 }, - email: { type: "string", format: "email" }, - roles: { - type: "array", - items: { enum: ["admin", "editor", "viewer"] }, - minItems: 1, - uniqueItems: true, - }, - }, - required: ["id", "email", "roles"], - additionalProperties: false, - }, - results: { - type: "array", - minItems: 2, - maxItems: 8, - items: { - type: "object", - properties: { - label: { type: "string", minLength: 2, maxLength: 32 }, - score: { type: "number", minimum: 0, maximum: 1 }, - tags: { - type: "array", - items: { type: "string", pattern: "^[a-z-]+$" }, - maxItems: 5, - }, - metadata: { - anyOf: [ - { type: "null" }, - { - type: "object", - properties: { - source: { type: "string" }, - cached: { type: "boolean" }, - }, - required: ["source", "cached"], - additionalProperties: false, - }, - ], - }, - }, - required: ["label", "score", "tags", "metadata"], - additionalProperties: false, - }, - }, - }, - required: ["request_id", "status", "user", "results"], - additionalProperties: false, - }, - }, - output: - '{"request_id":"req-2026-a1","status":"completed","user":{"id":42,"email":"user@example.com","roles":["admin","editor"]},"results":[{"label":"primary","score":0.97,"tags":["fast","verified"],"metadata":{"source":"cache-v2","cached":true}},{"label":"fallback","score":0.81,"tags":["review"],"metadata":null}]}', -}; diff --git a/packages/transformers-response-constraint/performance/run.mjs b/packages/transformers-response-constraint/performance/run.mjs deleted file mode 100644 index 2f3c3a6d8..000000000 --- a/packages/transformers-response-constraint/performance/run.mjs +++ /dev/null @@ -1,98 +0,0 @@ -import { Tensor } from "@huggingface/transformers"; - -import { ResponseConstraint } from "../dist/index.js"; -import complexJson from "./complex-json.mjs"; -import simpleJson from "./simple-json.mjs"; -import simpleRegex from "./simple-regex.mjs"; - -const WARMUP_RUNS = 2; -const MEASURED_RUNS = 10; -const EOS_TOKEN_ID = 256; -const VOCAB_SIZE = Number(process.env.VOCAB_SIZE ?? 8192); -const encoder = new TextEncoder(); -const tokenizer = { - tokens: [ - ...Array.from({ length: EOS_TOKEN_ID }, (_, tokenId) => [tokenId]), - [], - ...Array.from({ length: VOCAB_SIZE - EOS_TOKEN_ID - 1 }, (_, index) => [ - ...encoder.encode(` token-${index.toString(36)}`), - ]), - ], - eos_token_id: EOS_TOKEN_ID, - special_token_ids: [EOS_TOKEN_ID], -}; -const encodedCases = [simpleJson, simpleRegex, complexJson].map((testCase) => ({ - ...testCase, - tokenIds: encoder.encode(testCase.output), -})); -const coldResults = new Map(); - -for (const testCase of encodedCases) { - coldResults.set(testCase, await measure(testCase)); -} - -for (const testCase of encodedCases) { - for (let run = 0; run < WARMUP_RUNS; ++run) await measure(testCase); -} - -const results = []; -for (const testCase of encodedCases) { - let elapsedMs = 0; - let processorMs = 0; - let updateMs = 0; - for (let run = 0; run < MEASURED_RUNS; ++run) { - const result = await measure(testCase); - elapsedMs += result.elapsedMs; - processorMs += result.processorMs; - updateMs += result.updateMs; - } - results.push({ - constraint: testCase.name, - vocabulary: VOCAB_SIZE, - tokens: testCase.tokenIds.length + 1, - "cold ms": format(coldResults.get(testCase).elapsedMs), - "total ms": format(elapsedMs / MEASURED_RUNS), - "ms/token": format( - elapsedMs / MEASURED_RUNS / (testCase.tokenIds.length + 1), - ), - "processor ms": format(processorMs / MEASURED_RUNS), - "update ms": format(updateMs / MEASURED_RUNS), - }); -} - -console.table(results); - -async function measure(testCase) { - const constraint = await ResponseConstraint.fromResponseFormat( - tokenizer, - testCase.responseFormat, - ); - const scores = new Tensor("float32", new Float32Array(VOCAB_SIZE), [ - 1, - VOCAB_SIZE, - ]); - const inputIds = [0n]; - const startedAt = performance.now(); - let processorMs = 0; - let updateMs = 0; - for (const tokenId of [...testCase.tokenIds, EOS_TOKEN_ID]) { - scores.data.fill(0); - let stepStartedAt = performance.now(); - constraint.logits_processor([inputIds], scores); - processorMs += performance.now() - stepStartedAt; - if (!Number.isFinite(scores.data[tokenId])) - throw new Error(`${testCase.name} rejected token ${tokenId}`); - inputIds.push(BigInt(tokenId)); - stepStartedAt = performance.now(); - constraint.logits_processor.onTokensSampled([tokenId], [inputIds]); - constraint.stopping_criteria([inputIds]); - updateMs += performance.now() - stepStartedAt; - } - if (!constraint.stopping_criteria([inputIds])[0]) - throw new Error(`${testCase.name} did not stop after EOS`); - return { elapsedMs: performance.now() - startedAt, processorMs, updateMs }; -} - -function format(value) { - return value.toFixed(3); -} diff --git a/packages/transformers-response-constraint/performance/simple-json.mjs b/packages/transformers-response-constraint/performance/simple-json.mjs deleted file mode 100644 index d3ce4eb17..000000000 --- a/packages/transformers-response-constraint/performance/simple-json.mjs +++ /dev/null @@ -1,16 +0,0 @@ -export default { - name: "simple JSON schema", - responseFormat: { - type: "json_schema", - json_schema: { - type: "object", - properties: { - answer: { type: "string" }, - confidence: { type: "number", minimum: 0, maximum: 1 }, - }, - required: ["answer", "confidence"], - additionalProperties: false, - }, - }, - output: '{"answer":"Paris","confidence":0.98}', -}; diff --git a/packages/transformers-response-constraint/performance/simple-regex.mjs b/packages/transformers-response-constraint/performance/simple-regex.mjs deleted file mode 100644 index 6a8cc6fe2..000000000 --- a/packages/transformers-response-constraint/performance/simple-regex.mjs +++ /dev/null @@ -1,8 +0,0 @@ -export default { - name: "simple regex", - responseFormat: { - type: "regex", - regex: "[A-Z]{3}-\\d{4}", - }, - output: "ABC-2026", -}; From 7fd417914e4d7e4c6c5e96e096a7b8c675656e8d Mon Sep 17 00:00:00 2001 From: Nico Martin Date: Fri, 21 Aug 2026 08:54:07 +0200 Subject: [PATCH 16/17] removed onTokensSampled --- .../src/ResponseConstraint.ts | 16 ++++--- .../tests/response-constraint.test.js | 25 ++++++++--- .../src/generation/logits_process.js | 13 ------ .../transformers/src/models/modeling_utils.js | 4 -- .../tests/utils/generation.test.js | 44 ------------------- 5 files changed, 29 insertions(+), 73 deletions(-) diff --git a/packages/transformers-response-constraint/src/ResponseConstraint.ts b/packages/transformers-response-constraint/src/ResponseConstraint.ts index bcf1c0420..35dc67912 100644 --- a/packages/transformers-response-constraint/src/ResponseConstraint.ts +++ b/packages/transformers-response-constraint/src/ResponseConstraint.ts @@ -17,6 +17,7 @@ export type ResponseFormat = type GenerationState = { completed: boolean; constraint: TokenConstraint; + processedInputLength?: number; mask?: Uint32Array; }; @@ -53,6 +54,7 @@ class ConstraintLogitsProcessor extends LogitsProcessor { _call(inputIds: bigint[][], logits: Tensor) { assertSingleSequence(inputIds.length); + this.state.processedInputLength ??= inputIds[0].length; if (this.state.completed) return logits; const logitsVocabSize = logits.dims.at(-1); if (logitsVocabSize === undefined || !Number.isInteger(logitsVocabSize) || logitsVocabSize <= 0) { @@ -66,12 +68,6 @@ class ConstraintLogitsProcessor extends LogitsProcessor { applyMask(logits, this.state.mask, this.state.constraint.vocabSize); return logits; } - - onTokensSampled(tokenIds: number[], inputIds: bigint[][]) { - assertSingleSequence(tokenIds.length); - assertSingleSequence(inputIds.length); - if (!this.state.completed) this.state.completed = this.state.constraint.commit(tokenIds[0]); - } } class ConstraintStoppingCriteria extends StoppingCriteria { @@ -79,8 +75,14 @@ class ConstraintStoppingCriteria extends StoppingCriteria { super(); } - _call(inputIds: ArrayLike[]) { + _call(inputIds: ArrayLike[]) { assertSingleSequence(inputIds.length); + const input = inputIds[0]; + const start = this.state.processedInputLength ?? input.length; + for (let i = start; i < input.length && !this.state.completed; ++i) { + this.state.completed = this.state.constraint.commit(Number(input[i])); + } + this.state.processedInputLength = input.length; return [this.state.completed]; } } diff --git a/packages/transformers-response-constraint/tests/response-constraint.test.js b/packages/transformers-response-constraint/tests/response-constraint.test.js index aa132c4a6..8372aef6c 100644 --- a/packages/transformers-response-constraint/tests/response-constraint.test.js +++ b/packages/transformers-response-constraint/tests/response-constraint.test.js @@ -17,14 +17,16 @@ function isAllowed(scores, tokenId) { return Number.isFinite(scores.data[tokenId]); } +const inputIdsByConstraint = new WeakMap(); + async function consume(constraint, text) { - const inputIds = [0n]; + const inputIds = inputIdsByConstraint.get(constraint) ?? [0n]; + inputIdsByConstraint.set(constraint, inputIds); for (const tokenId of new TextEncoder().encode(text)) { const scores = logits(); constraint.logits_processor([inputIds], scores); expect(isAllowed(scores, tokenId)).toBe(true); inputIds.push(BigInt(tokenId)); - constraint.logits_processor.onTokensSampled([tokenId], [inputIds]); expect(constraint.stopping_criteria([inputIds])).toEqual([false]); } return inputIds; @@ -45,7 +47,7 @@ function schemaAccepts(schema, text) { } if (!isAllowed(scores, tokenId)) return false; inputIds.push(BigInt(tokenId)); - constraint.logits_processor.onTokensSampled([tokenId], [inputIds]); + constraint.stopping_criteria([inputIds]); } const scores = logits(); try { @@ -69,6 +71,20 @@ describe("ResponseConstraint", () => { expect(isAllowed(scores, EOS_TOKEN_ID)).toBe(false); }); + it("does not process the same sampled token twice", async () => { + const constraint = ResponseConstraint.fromResponseFormat(tokenizer, { type: "regex", regex: "ab" }); + const inputIds = [0n]; + constraint.logits_processor([inputIds], logits()); + inputIds.push(BigInt("a".charCodeAt(0))); + + expect(constraint.stopping_criteria([inputIds])).toEqual([false]); + expect(constraint.stopping_criteria([inputIds])).toEqual([false]); + + const scores = logits(); + constraint.logits_processor([inputIds], scores); + expect(isAllowed(scores, "b".charCodeAt(0))).toBe(true); + }); + it("accepts JSON that satisfies a schema", async () => { const constraint = await ResponseConstraint.fromResponseFormat(tokenizer, { type: "json_schema", @@ -84,7 +100,6 @@ describe("ResponseConstraint", () => { constraint.logits_processor([inputIds], scores); expect(isAllowed(scores, EOS_TOKEN_ID)).toBe(true); - constraint.logits_processor.onTokensSampled([EOS_TOKEN_ID], [[...inputIds, BigInt(EOS_TOKEN_ID)]]); expect(constraint.stopping_criteria([[...inputIds, BigInt(EOS_TOKEN_ID)]])).toEqual([true]); }); @@ -495,7 +510,7 @@ describe("ResponseConstraint", () => { const constraint = await ResponseConstraint.fromResponseFormat(tokenizer, { type: "regex", regex: "a" }); constraint.logits_processor([[0n]], logits()); - expect(() => constraint.logits_processor.onTokensSampled(["b".charCodeAt(0)], [[0n, 98n]])).toThrow("does not satisfy the constraint"); + expect(() => constraint.stopping_criteria([[0n, 98n]])).toThrow("does not satisfy the constraint"); }); it("returns only the generation hooks", async () => { diff --git a/packages/transformers/src/generation/logits_process.js b/packages/transformers/src/generation/logits_process.js index 7bb9545eb..647a30806 100644 --- a/packages/transformers/src/generation/logits_process.js +++ b/packages/transformers/src/generation/logits_process.js @@ -88,19 +88,6 @@ 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 59c07768f..68387468d 100644 --- a/packages/transformers/src/models/modeling_utils.js +++ b/packages/transformers/src/models/modeling_utils.js @@ -992,8 +992,6 @@ export class PreTrainedModel extends Callable { /** @type {[bigint][]} */ const generated_input_ids = []; - /** @type {number[]} */ - const sampled_token_ids = []; // const new_kv_cache = [];// NOTE: Only used for beam search when concatenating new kv // Loop over each batch for (let batch_idx = 0; batch_idx < next_tokens_scores.dims.at(0); ++batch_idx) { @@ -1005,7 +1003,6 @@ export class PreTrainedModel extends Callable { // update generated ids, model inputs, and length for next step scores[batch_idx] += logProb; generated_input_ids.push([newTokenId]); - sampled_token_ids.push(Number(newTokenId)); // TODO: Support beam search break; @@ -1014,7 +1011,6 @@ export class PreTrainedModel extends Callable { 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(sampled_token_ids, 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 db991affc..d827ba3e3 100644 --- a/packages/transformers/tests/utils/generation.test.js +++ b/packages/transformers/tests/utils/generation.test.js @@ -11,8 +11,6 @@ import { // Other TextStreamer, DynamicCache, - LogitsProcessor, - LogitsProcessorList, StoppingCriteria, random, full, @@ -211,48 +209,6 @@ 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 () => { From 47e71193286d0c668c96a729ec05a8e1a95b3259 Mon Sep 17 00:00:00 2001 From: Nico Martin Date: Fri, 21 Aug 2026 08:58:51 +0200 Subject: [PATCH 17/17] added whitespace penalty --- .../src/ResponseConstraint.ts | 25 +++++++++++++++++ .../src/engine/constraint.ts | 19 +++++++++++++ .../tests/response-constraint.test.js | 28 +++++++++++++++++++ 3 files changed, 72 insertions(+) diff --git a/packages/transformers-response-constraint/src/ResponseConstraint.ts b/packages/transformers-response-constraint/src/ResponseConstraint.ts index 35dc67912..1e58cd9c4 100644 --- a/packages/transformers-response-constraint/src/ResponseConstraint.ts +++ b/packages/transformers-response-constraint/src/ResponseConstraint.ts @@ -21,6 +21,9 @@ type GenerationState = { mask?: Uint32Array; }; +const WHITESPACE_REPETITION_PENALTY = 1.2; +const MAX_CONSECUTIVE_WHITESPACE_TOKENS = 4; + export class ResponseConstraint { /** * Precomputes the tokenizer-derived data structures used by every @@ -66,6 +69,10 @@ class ConstraintLogitsProcessor extends LogitsProcessor { throw new Error('The constraint reached a dead end before producing a valid output.'); } applyMask(logits, this.state.mask, this.state.constraint.vocabSize); + const repeatedWhitespace = this.state.constraint.repeatedWhitespace(); + if (repeatedWhitespace !== undefined) { + discourageRepeatedWhitespace(logits, repeatedWhitespace.tokenIds, repeatedWhitespace.count); + } return logits; } } @@ -92,3 +99,21 @@ function assertSingleSequence(batchSize: number): void { throw new Error(`ResponseConstraint currently supports batch size 1; received ${batchSize}.`); } } + +function discourageRepeatedWhitespace(logits: Tensor, tokenIds: readonly number[], count: number): void { + const data = logits.data as Float32Array | Float64Array | number[]; + const stride = logits.dims.at(-1)!; + const penalty = WHITESPACE_REPETITION_PENALTY ** count; + for (let offset = 0; offset < data.length; offset += stride) { + for (const tokenId of tokenIds) { + const index = offset + tokenId; + if (count >= MAX_CONSECUTIVE_WHITESPACE_TOKENS) { + data[index] = -Infinity; + } else if (data[index] < 0) { + data[index] *= penalty; + } else { + data[index] /= penalty; + } + } + } +} diff --git a/packages/transformers-response-constraint/src/engine/constraint.ts b/packages/transformers-response-constraint/src/engine/constraint.ts index ef762ecc7..8a62de7de 100644 --- a/packages/transformers-response-constraint/src/engine/constraint.ts +++ b/packages/transformers-response-constraint/src/engine/constraint.ts @@ -7,6 +7,7 @@ type TrieNode = { childBytes: number[]; childNodes: TrieNode[]; tokenIds: number type CachedTokenizer = { data: TokenizerData; trie: TrieNode; + whitespaceTokenIds: number[]; stringExceptionalTrie: TrieNode; stringSafeMask: Uint32Array; stringSafeCount: number; @@ -28,6 +29,7 @@ export type TokenConstraint = { vocabSize: number; fillMask(target: Uint32Array): boolean; commit(tokenId: number): boolean; + repeatedWhitespace(): { tokenIds: readonly number[]; count: number } | undefined; }; const tokenizerCache = new WeakMap(); @@ -59,6 +61,8 @@ export function createTokenConstraint( const tokenStates: Array = new Array(tokenizer.data.tokens.length); const tokenStamps = new Int32Array(tokenizer.data.tokens.length); let stamp = 0; + let consecutiveWhitespace = 0; + const tracksJsonWhitespace = responseFormat.type !== 'regex'; return { vocabSize: tokenizer.data.tokens.length, @@ -128,9 +132,17 @@ export function createTokenConstraint( } stamp++; if (!machine.viable(next)) throw new Error(`Token ${tokenId} does not satisfy the constraint.`); + consecutiveWhitespace = + tracksJsonWhitespace && next === state && isJsonWhitespace(tokenizer.data.tokens[tokenId]) + ? consecutiveWhitespace + 1 + : 0; state = next; return false; }, + repeatedWhitespace() { + if (consecutiveWhitespace === 0) return undefined; + return { tokenIds: tokenizer.whitespaceTokenIds, count: consecutiveWhitespace }; + }, }; } @@ -153,6 +165,7 @@ function cachedTokenizer(source: TokenizerSource): CachedTokenizer { if (cached === undefined) { const data = extractTokenizer(source); const stringExceptionalTokenIds: number[] = []; + const whitespaceTokenIds: number[] = []; const stringSafeMask = new Uint32Array(Math.ceil(data.tokens.length / 32)); const stringSafeLengths = new Uint32Array(data.tokens.length); let stringSafeCount = 0; @@ -160,6 +173,7 @@ function cachedTokenizer(source: TokenizerSource): CachedTokenizer { let maxTokenByteLength = 0; for (let tokenId = 0; tokenId < data.tokens.length; ++tokenId) { const special = data.specialTokenIds.has(tokenId); + if (!special && isJsonWhitespace(data.tokens[tokenId])) whitespaceTokenIds.push(tokenId); if (!special && data.tokens[tokenId].length > maxTokenByteLength) { maxTokenByteLength = data.tokens[tokenId].length; } @@ -176,6 +190,7 @@ function cachedTokenizer(source: TokenizerSource): CachedTokenizer { cached = { data, trie: createTrie(data.tokens), + whitespaceTokenIds, stringExceptionalTrie: createTrie(data.tokens, stringExceptionalTokenIds), stringSafeMask, stringSafeCount, @@ -307,6 +322,10 @@ function boundedStringMask(tokenizer: CachedTokenizer, capacity: number): { mask return cached; } +function isJsonWhitespace(bytes: Uint8Array): boolean { + return bytes.length > 0 && bytes.every((byte) => byte === 0x09 || byte === 0x0a || byte === 0x0d || byte === 0x20); +} + function setBit(mask: Uint32Array, tokenId: number): void { mask[tokenId >>> 5] |= 1 << (tokenId & 31); } diff --git a/packages/transformers-response-constraint/tests/response-constraint.test.js b/packages/transformers-response-constraint/tests/response-constraint.test.js index 8372aef6c..f23535521 100644 --- a/packages/transformers-response-constraint/tests/response-constraint.test.js +++ b/packages/transformers-response-constraint/tests/response-constraint.test.js @@ -85,6 +85,34 @@ describe("ResponseConstraint", () => { expect(isAllowed(scores, "b".charCodeAt(0))).toBe(true); }); + it("discourages repeated non-progressing JSON whitespace", () => { + const constraint = ResponseConstraint.fromResponseFormat(tokenizer, { type: "json_object" }); + const inputIds = [0n]; + constraint.logits_processor([inputIds], logits()); + + for (let count = 1; count <= 4; ++count) { + inputIds.push(10n); + expect(constraint.stopping_criteria([inputIds])).toEqual([false]); + + const scores = logits(); + scores.data[10] = 12; + scores.data[32] = 12; + scores.data[13] = -12; + constraint.logits_processor([inputIds], scores); + + if (count < 4) { + expect(scores.data[10]).toBeCloseTo(12 / 1.2 ** count); + expect(scores.data[32]).toBeCloseTo(12 / 1.2 ** count); + expect(scores.data[13]).toBeCloseTo(-12 * 1.2 ** count); + } else { + expect(scores.data[10]).toBe(-Infinity); + expect(scores.data[32]).toBe(-Infinity); + expect(scores.data[13]).toBe(-Infinity); + } + expect(scores.data["{".charCodeAt(0)]).toBe(1); + } + }); + it("accepts JSON that satisfies a schema", async () => { const constraint = await ResponseConstraint.fromResponseFormat(tokenizer, { type: "json_schema",