diffusers/main/pt/api/schedulers/block_refinement.md

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[[diffusers.BlockRefinementScheduler]]

diffusers.BlockRefinementScheduler[[diffusers.BlockRefinementScheduler]]

diffusers.BlockRefinementScheduler(block_length: int = 32, num_inference_steps: int = 32, threshold: float = 0.95, editing_threshold: float | None = None, minimal_topk: int = 1)

Fonte

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 pelo 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[[diffusers.BlockRefinementScheduler.add_noise]]

add_noise(original_samples: torch.LongTensor, attention_mask: torch.LongTensor, prompt_length: int, block_length: int, mask_token_id: int, generator: torch.Generator | None = None)

Fonte

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 respectivas máscaras booleanas.

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, em que as posições mascaradas em uma correspondem às posições não mascaradas na outra.

check_block_should_continue[[diffusers.BlockRefinementScheduler.check_block_should_continue]]

check_block_should_continue(step_idx: int, masks_remaining: bool, editing_enabled: bool, editing_transfer_index: torch.BoolTensor, post_steps: int, max_post_steps: int, finished: torch.BoolTensor)

Fonte

Parâmetros:

step_idx (int) : Índice atual da etapa de refinamento dentro deste bloco.

masks_remaining (bool) : Se há tokens de máscara restantes 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) : Etapas de edição pós-máscara máximas permitidas.

finished (torch.BoolTensor) : Flags de conclusã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[[diffusers.BlockRefinementScheduler.check_eos_finished]]

check_eos_finished(cur_x: torch.LongTensor, sampled_tokens: torch.LongTensor, final_transfer: torch.BoolTensor, finished: torch.BoolTensor, eos_token_id: int, mask_token_id: int, prompt_length: int)

Fonte

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[[diffusers.BlockRefinementScheduler.get_num_transfer_tokens]]

get_num_transfer_tokens(block_length: int, num_inference_steps: int)

Fonte

Distribua uniformemente os block_length tokens entre as num_inference_steps etapas.

etapa[[diffusers.BlockRefinementScheduler.step]]

step(model_output: torch.Tensor, timestep: int | torch.Tensor, sample: torch.LongTensor, mask_token_id: int | None = None, temperature: float = 0.0, top_p: float | None = None, top_k: int | None = None, sampling_method: str = 'auto', threshold: float | None = None, editing_threshold: float | None = None, minimal_topk: int | None = None, prompt_mask: torch.BoolTensor | None = None, generator: torch.Generator | None = None, return_dict: bool = True)

Fonte

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 do passo atual dentro da agenda 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 confirmadas (reiniciando em timestep == 0) e re-ruidifica as não confirmadas 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 por 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 edição de tokens não mascarados; deve ser positivo para habilitar a edição. O valor padrão é o 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) : RNG 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[[diffusers.BlockRefinementSchedulerOutput]]

diffusers.BlockRefinementSchedulerOutput[[diffusers.BlockRefinementSchedulerOutput]]

diffusers.BlockRefinementSchedulerOutput(prev_sample: torch.LongTensor, transfer_index: torch.BoolTensor, editing_transfer_index: torch.BoolTensor, sampled_tokens: torch.LongTensor, sampled_probs: torch.Tensor, pred_logits: torch.Tensor)

Fonte

Parâmetros:

prev_sample (torch.LongTensor de forma (batch_size, block_length)) : Tokens de bloco atualizados após a etapa de refinamento atual.

transfer_index (torch.BoolTensor de forma (batch_size, block_length)) : Máscara booleana indicando 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 mascarada).

sampled_tokens (torch.LongTensor de forma (batch_size, block_length)) : IDs de tokens amostrados 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 denoiser, passados adiante para autocondicionar o próximo passo.

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