Skip to content

Add Generation Extension Points for Constrained Decoding#1720

Open
nico-martin wants to merge 4 commits into
mainfrom
feat/response_format
Open

Add Generation Extension Points for Constrained Decoding#1720
nico-martin wants to merge 4 commits into
mainfrom
feat/response_format

Conversation

@nico-martin

Copy link
Copy Markdown
Collaborator

In this PR I keep structured/constrained generation out of core and expose the same kind of generation extension points that Python transformers already uses:

  • logits_processor for transforming token scores before sampling.
  • stopping_criteria for deciding when generation should stop.

The main idea is that JSON Schema / llguidance-style decoding should be possible from an external package, without adding WASM, tokenizer bridge code, or schema validation to Transformers.js itself.

API Shape

The core API stays small: callers pass custom generation pieces through the existing generate() path.

const constraint = await createLlguidanceConstraint(
  tokenizer,
  response_format,
);

const output = await model.generate({
  ...inputs,
  max_new_tokens: 128,
  logits_processor: constraint.logits_processor,
  stopping_criteria: constraint.stopping_criteria,
});

I keep logits_processor list-like, matching Python transformers usage. A single custom processor can be wrapped in LogitsProcessorList, or an external helper can return the list directly.

I also keep stopping_criteria separate from logits processing. That preserves the upstream split: processors modify scores, criteria stops generation.

createLlguidanceConstraint() is only an example external-package helper. The Transformers.js API surface is the generate() arguments, not the helper name.

Python transformers Compatibility

Python transformers already models constrained generation as composable generation hooks:

outputs = model.generate(
    **inputs,
    logits_processor=logits_processor_list,
    stopping_criteria=stopping_criteria_list,
)

This PR follows that shape instead of adding a pipeline-specific response_format option to core. The important portable part stays the same: LogitsProcessor.__call__(input_ids, scores) -> scores.

The only Transformers.js-specific piece is an optional batch-level onTokensSampled(tokenIds, inputIds) hook on processors inside a LogitsProcessorList. I added that for stateful decoders that need to commit the sampled token after sampling. It runs once per generation step after the full batch step has been appended, so processors do not see a partially updated batch.

Minimal External Constraint Example

This is not real llguidance. It is a small, visible toy constraint: generated tokens may not decode to emoji. The behavior is easy to verify, but it uses the same shape an external llguidance package would use.

import {
  LogitsProcessor,
  LogitsProcessorList,
  StoppingCriteria,
} from "@huggingface/transformers";

const emojiPattern = /[\p{Emoji_Presentation}\p{Emoji}\uFE0F]/u;

class NoEmojiProcessor extends LogitsProcessor {
  constructor(state) {
    super();
    this.state = state;
  }

  _call(_inputIds, logits) {
    const vocabSize = logits.dims.at(-1);
    if (!vocabSize) return logits;

    const data = logits.data;
    const batchSize = Math.max(1, data.length / vocabSize);

    for (let batch = 0; batch < batchSize; ++batch) {
      const offset = batch * vocabSize;
      for (let tokenId = 0; tokenId < vocabSize; ++tokenId) {
        if (!this.state.allowedTokenIds.has(tokenId)) {
          data[offset + tokenId] = -Infinity;
        }
      }
    }

    return logits;
  }

  onTokensSampled(tokenIds) {
    this.state.generated += tokenIds.length;
    if (this.state.generated >= this.state.maxTokens) {
      this.state.done = true;
    }
  }
}

class NoEmojiStoppingCriteria extends StoppingCriteria {
  constructor(state) {
    super();
    this.state = state;
  }

  _call(inputIds) {
    return new Array(inputIds.length).fill(this.state.done);
  }
}

export async function createNoEmojiConstraint(tokenizer, options = {}) {
  const allowedTokenIds = new Set();
  const vocabSize = getVocabSize(tokenizer);

  for (let tokenId = 0; tokenId < vocabSize; ++tokenId) {
    const text = tokenizer.decode([tokenId]);
    if (!emojiPattern.test(text)) {
      allowedTokenIds.add(tokenId);
    }
  }

  const state = {
    allowedTokenIds,
    generated: 0,
    maxTokens: options.max_tokens ?? 32,
    done: false,
  };

  const logits_processor = new LogitsProcessorList();
  logits_processor.push(new NoEmojiProcessor(state));

  return {
    logits_processor,
    stopping_criteria: new NoEmojiStoppingCriteria(state),
  };
}

function getVocabSize(tokenizer) {
  const vocab = getVocab(tokenizer);
  const vocabSize =
    tokenizer.vocab_size ??
    tokenizer.vocabSize ??
    tokenizer.get_vocab_size?.() ??
    tokenizer.getVocabSize?.() ??
    getVocabSizeFromVocab(vocab);

  if (!Number.isInteger(vocabSize)) {
    throw new Error("Unable to determine tokenizer vocabulary size");
  }

  return Number(vocabSize);
}

function getVocab(tokenizer) {
  return (
    tokenizer.vocab ??
    tokenizer.get_vocab?.() ??
    tokenizer.getVocab?.() ??
    tokenizer._tokenizer?.get_vocab?.() ??
    tokenizer._tokenizer?.getVocab?.()
  );
}

function getVocabSizeFromVocab(vocab) {
  if (!vocab) return undefined;

  const ids = vocab instanceof Map ? vocab.values() : Object.values(vocab);
  let maxTokenId = -1;
  for (const id of ids) {
    maxTokenId = Math.max(maxTokenId, Number(id));
  }

  return maxTokenId + 1;
}

Usage:

const { logits_processor, stopping_criteria } = await createNoEmojiConstraint(tokenizer, {
  max_tokens: 32,
});

await model.generate({
  ...inputs,
  logits_processor,
  stopping_criteria,
});

Benefits of this direction:

  • It keeps Transformers.js core lightweight.
  • It avoids bundling llguidance/WASM for users who do not need constrained decoding.
  • It matches the Python transformers split between score processing and stopping.
  • It gives external packages enough hooks to implement schema-constrained decoding.
  • It keeps future constrained decoding backends independent from core release cycles.

@nico-martin nico-martin requested a review from xenova July 10, 2026 08:56
@HuggingFaceDocBuilderDev

Copy link
Copy Markdown

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants