Hugging Face's logo

Diffusers documentation

BlockRefinementScheduler

Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

BlockRefinementScheduler

O BlockRefinementScheduler gerencia o refinamento iterativo por bloco para difusão de tokens discretos. A cada etapa, ele confirma os tokens mais confiantes e, opcionalmente, edita tokens já confirmados quando o modelo prevê um token diferente com alta confiança.

Este scheduler é usado por LLaDA2Pipeline.

BlockRefinementScheduler

classe diffusers.BlockRefinementScheduler

< >

( block_length: int = 32num_inference_steps: int = 32threshold: float = 0.95editing_threshold: float | None = Noneminimal_topk: int = 1 )

Agendador para refinamento iterativo por blocos (commit por confiança).

A cada etapa, o agendador amostra tokens candidatos dos logits do modelo e confirma aqueles com a maior confiança. O número de tokens a serem confirmados por etapa é determinado distribuindo uniformemente o comprimento do bloco entre o número de etapas de refinamento.

Suporta opcionalmente edição: após todos os tokens de máscara serem resolvidos, os tokens podem ser substituídos se o modelo prever um token diferente com confiança acima de um editing_threshold positivo (None, 0.0 ou negativo desativa a edição).

add_noise

< >

( original_samples: torch.LongTensorattention_mask: torch.LongTensorprompt_length: intblock_length: intmask_token_id: intgenerator: torch.Generator | None = None ) → tuple[torch.LongTensor, torch.LongTensor, torch.BoolTensor, torch.BoolTensor]

Parâmetros

  • original_samples (torch.LongTensor de forma (batch_size, seq_len)) — IDs de tokens limpos.
  • attention_mask (torch.LongTensor de forma (batch_size, seq_len)) — Máscara de preenchimento (1 para válido, 0 para preenchimento).
  • prompt_length (int) — Número de tokens de prompt iniciais a manter sem máscara.
  • block_length (int) — Tamanho do bloco para mascaramento.
  • mask_token_id (int) — ID do token a ser usado para posições mascaradas.
  • generator (torch.Generator, opcional) — RNG para reprodutibilidade.

Retorna

tuple[torch.LongTensor, torch.LongTensor, torch.BoolTensor, torch.BoolTensor]

(noisy, noisy_rev, masked, masked_rev) — as duas sequências ruidosas complementares e suas máscaras booleanas correspondentes.

Aplique o processo de encaminhamento (adição de ruído) para mascaramento de blocos semi-autorregressivos.

Para cada bloco após o prompt, uma fração aleatória de tokens válidos (não de preenchimento) é substituída por mask_token_id. Duas visões complementares são retornadas: noisy e noisy_rev, nas quais as posições mascaradas em uma correspondem às posições não mascaradas na outra.

check_block_should_continue

< >

( step_idx: intmasks_remaining: boolediting_enabled: boolediting_transfer_index: torch.BoolTensorpost_steps: intmax_post_steps: intfinished: torch.BoolTensor ) → bool

Parâmetros

  • step_idx (int) — Índice atual da etapa de refinamento dentro deste bloco.
  • masks_remaining (bool) — Se ainda há tokens de máscara no bloco.
  • editing_enabled (bool) — Se o modo de edição está ativo.
  • editing_transfer_index (torch.BoolTensor) — Quais tokens foram editados nesta etapa.
  • post_steps (int) — Número de etapas de edição pós-máscara realizadas até o momento.
  • max_post_steps (int) — Número máximo de etapas de edição pós-máscara permitidas.
  • finished (torch.BoolTensor) — Flags de finalização por lote (da detecção de EOS).

Retorna

bool

True se o refinamento deve continuar, False para interromper.

Determine se o loop de refinamento interno deve continuar para o bloco atual.

check_eos_finished

< >

( cur_x: torch.LongTensorsampled_tokens: torch.LongTensorfinal_transfer: torch.BoolTensorfinished: torch.BoolTensoreos_token_id: intmask_token_id: intprompt_length: int ) → torch.BoolTensor

Parâmetros

  • cur_x (torch.LongTensor de forma (batch_size, seq_len)) — Sequência completa atual, incluindo todos os blocos até a janela atual.
  • sampled_tokens (torch.LongTensor de forma (batch_size, block_length)) — Tokens amostrados pelo agendador nesta etapa.
  • final_transfer (torch.BoolTensor de forma (batch_size, block_length)) — Máscara combinada das posições confirmadas e editadas.
  • finished (torch.BoolTensor de forma (batch_size,)) — Flags de finalização atuais por lote.
  • eos_token_id (int) — ID do token EOS.
  • mask_token_id (int) — ID do token de máscara.
  • prompt_length (int) — Número de tokens de prompt no início da sequência.

Retorna

torch.BoolTensor

Atualizados os flags de finalização.

Atualize as flags de conclusão por lote quando os tokens EOS forem confirmados.

get_num_transfer_tokens

< >

( block_length: intnum_inference_steps: int )

Distribua uniformemente os block_length tokens entre as num_inference_steps etapas.

etapa

< >

( model_output: torch.Tensortimestep: int | torch.Tensorsample: torch.LongTensormask_token_id: int | None = Nonetemperature: float = 0.0top_p: float | None = Nonetop_k: int | None = Nonesampling_method: str = 'auto'threshold: float | None = Noneediting_threshold: float | None = Noneminimal_topk: int | None = Noneprompt_mask: torch.BoolTensor | None = Nonegenerator: torch.Generator | None = Nonereturn_dict: bool = True )

Parâmetros

  • model_output (torch.Tensor de forma (batch_size, block_length, vocab_size)) — Logits brutos do modelo para o bloco atual.
  • timestep (int ou torch.Tensor) — Índice atual da etapa dentro do cronograma de refinamento do bloco.
  • sample (torch.LongTensor de forma (batch_size, block_length)) — IDs dos tokens do bloco atual (contém tokens de máscara para posições não confirmadas no modo baseado em máscara).
  • mask_token_id (int, opcional) — ID do token usado para posições mascaradas. Quando None, o agendador opera no modo de corrupção uniforme: ele rastreia internamente as posições comprometidas (reiniciando em timestep == 0) e renasce as não comprometidas com tokens aleatórios uniformes, correspondendo ao amostrador de refinamento de blocos do DiffusionGemma.
  • temperature (float) — Temperatura de amostragem.
  • top_p (float, opcional) — Corte de amostragem de núcleo.
  • top_k (int, opcional) — Corte de amostragem top-k.
  • sampling_method (str) — Método de amostragem (auto, greedy, multinomial).
  • threshold (float, opcional) — Limiar de confiança para confirmar tokens. O valor padrão é o definido na configuração.
  • editing_threshold (float, opcional) — Limiar de confiança para a edição de tokens não mascarados; deve ser positivo para habilitar a edição. O valor padrão é o valor da configuração.
  • minimal_topk (int, opcional) — Número mínimo de tokens a serem confirmados por etapa. O valor padrão é o definido na configuração.
  • prompt_mask (torch.BoolTensor, opcional) — Máscara booleana de forma (block_length,) onde True marca as posições do prompt (não editáveis).
  • generator (torch.Generator, opcional) — Gera números aleatórios para amostragem.
  • return_dict (bool) — Se deve retornar um BlockRefinementSchedulerOutput ou uma tupla.

Execute uma etapa única de refinamento: amostrar dos logits, confirmar tokens confiantes e, opcionalmente, editar os existentes.

BlockRefinementSchedulerOutput

classe diffusers.BlockRefinementSchedulerOutput

< >

( prev_sample: torch.LongTensortransfer_index: torch.BoolTensorediting_transfer_index: torch.BoolTensorsampled_tokens: torch.LongTensorsampled_probs: torch.Tensorpred_logits: torch.Tensor )

Parâmetros

  • prev_sample (torch.LongTensor de forma (batch_size, block_length)) — Tokens do bloco atualizados após a etapa de refinamento atual.
  • transfer_index (torch.BoolTensor de forma (batch_size, block_length)) — Máscara booleana que indica quais tokens foram confirmados (preenchimento de máscara).
  • editing_transfer_index (torch.BoolTensor de forma (batch_size, block_length)) — Máscara booleana indicando quais tokens foram editados (substituição não-máscara).
  • sampled_tokens (torch.LongTensor de forma (batch_size, block_length)) — IDs dos tokens amostrados a partir dos logits do modelo.
  • sampled_probs (torch.Tensor de forma (batch_size, block_length)) — Probabilidades dos tokens amostrados.
  • pred_logits (torch.Tensor de forma (batch_size, block_length, vocab_size)) — Os logits do desruidor, transmitidos para autocondicionar o próximo passo.

Classe de saída para agendamento de refinamento de blocos.

Update on GitHub

←Visão geral CMStochasticIterativeScheduler→