Diffusers documentation
BlockRefinementScheduler
BlockRefinementScheduler
BlockRefinementScheduler는 이산 토큰 확산에서 블록 단위 반복적 정화를 관리합니다. 각 단계에서 가장 확신도가 높은 토큰을 확정하고, 모델이 다른 토큰을 높은 확신도로 예측하는 경우 이미 확정된 토큰을 선택적으로 수정합니다.
이 스케줄러는 LLaDA2Pipeline에서 사용됩니다.
BlockRefinementScheduler
클래스 diffusers.BlockRefinementScheduler
< source >( block_length: int = 32num_inference_steps: int = 32threshold: float = 0.95editing_threshold: float | None = Noneminimal_topk: int = 1 )
블록 단위 반복적 정제(신뢰도 기반 커밋)를 위한 스케줄러
각 단계에서 스케줄러는 모델 로짓에서 후보 토큰을 샘플링하고 가장 높은 신뢰도를 가진 토큰을 확정합니다. 단계별로 확정할 토큰의 수는 블록 길이를 리파인먼트 단계 수에 걸쳐 균등하게 분배하여 결정됩니다.
선택적으로 편집을 지원합니다: 모든 mask 토큰이 해결된 후, 모델이 양의 editing_threshold보다 높은 확률로 다른 토큰을 예측하면 해당 토큰을 교체할 수 있습니다(None, 0.0 또는 음수는 편집을 비활성화합니다).
add_noise
< source >( 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]
매개변수
- original_samples (
torch.LongTensor, shape(batch_size, seq_len)) — 깨끗한 토큰 ID입니다. - attention_mask (
torch.LongTensor, shape(batch_size, seq_len)) — 패딩 마스크(유효 토큰은 1, 패딩은 0). - prompt_length (
int) — 마스크를 적용하지 않고 유지할 앞부분 프롬프트 토큰의 수. - block_length (
int) — 마스킹에 사용할 블록 크기입니다. - mask_token_id (
int) — 마스크된 위치에 사용할 토큰 ID입니다. - generator (
torch.Generator, 선택 사항) — 재현성을 위한 RNG입니다.
반환값
tuple[torch.LongTensor, torch.LongTensor, torch.BoolTensor, torch.BoolTensor]
(noisy, noisy_rev, masked, masked_rev) — 두 개의 상보적 노이즈 시퀀스와 이에 해당하는 불리언 마스크입니다.
반전파(noising) 과정을 반자동 회귀 블록 마스킹에 적용합니다.
프롬프트 이후의 각 블록에서, 유효한(패딩이 아닌) 토큰의 무작위 비율을 mask_token_id로 대체합니다. 두 개의 상보적인 뷰가 반환됩니다: noisy와 noisy_rev. 여기서 한 뷰의 마스킹된 위치는 다른 뷰의 마스킹되지 않은 위치입니다.
check_block_should_continue
< source >( step_idx: intmasks_remaining: boolediting_enabled: boolediting_transfer_index: torch.BoolTensorpost_steps: intmax_post_steps: intfinished: torch.BoolTensor ) → bool
매개변수
- step_idx (
int) — 이 블록 내에서 현재 정제 단계 인덱스. - masks_remaining (
bool) — 블록에 마스킹 토큰이 남아 있는지 여부. - editing_enabled (
bool) — 편집 모드가 활성화되어 있는지 여부. - editing_transfer_index (
torch.BoolTensor) — 이 단계에서 편집된 토큰을 나타냅니다. - post_steps (
int) — 지금까지 수행된 post-mask 편집 단계 수. - max_post_steps (
int) — 허용되는 최대 post-mask 편집 단계입니다. - finished (
torch.BoolTensor) — EOS 감지로부터 생성된 배치별 완료 플래그입니다.
반환값
bool
True이면 정밀화가 계속되고, False이면 중단됩니다.
현재 블록에 대해 내부 정제 루프를 계속할지 여부를 결정합니다.
check_eos_finished
< source >( cur_x: torch.LongTensorsampled_tokens: torch.LongTensorfinal_transfer: torch.BoolTensorfinished: torch.BoolTensoreos_token_id: intmask_token_id: intprompt_length: int ) → torch.BoolTensor
매개변수
- cur_x (
torch.LongTensor, shape(batch_size, seq_len)) — 현재 윈도우까지의 모든 블록을 포함한 현재 전체 시퀀스. - sampled_tokens (
torch.LongTensor, shape(batch_size, block_length)) — 해당 단계에서 스케줄러가 샘플링한 토큰. - final_transfer (
torch.BoolTensor, shape(batch_size, block_length)) — 커밋된 위치와 편집된 위치의 결합된 마스크. - finished (
torch.BoolTensor, shape(batch_size,)) — 현재 배치별 완료 플래그입니다. - eos_token_id (
int) — EOS 토큰 ID입니다. - mask_token_id (
int) — 마스크 토큰 ID입니다. - prompt_length (
int) — 시퀀스 시작 부분의 프롬프트 토큰 수.
반환값
torch.BoolTensor
finished 플래그를 업데이트했습니다.
EOS 토큰이 커밋되면 배치별 완료 플래그를 업데이트합니다.
block_length 토큰 커밋을 num_inference_steps 단계에 걸쳐 균등하게 분배합니다.
단계
< source >( 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 )
매개변수
- model_output (
torch.Tensor, shape(batch_size, block_length, vocab_size)) — 현재 블록에 대한 모델의 원시 로짓(raw logits). - timestep (
int또는torch.Tensor) — 블록의 정제 스케줄 내에서 현재 단계 인덱스입니다. - sample (
torch.LongTensor, shape(batch_size, block_length)) — 현재 블록의 토큰 ID입니다. (mask 기반 모드에서는 미확정 위치의 mask 토큰을 포함합니다.) - mask_token_id (
int, 선택 사항) — 마스크된 위치에 사용되는 토큰 ID입니다.None인 경우 스케줄러는 균일한 손상 모드에서 실행됩니다. 이 모드에서는 내부적으로 커밋된 위치를 추적하고(timestep == 0에서 초기화됨) 커밋되지 않은 위치를 균일하게 무작위로 선택된 토큰으로 재노이즈하여 DiffusionGemma의 블록 리파인먼트 샘플러와 일치합니다. - temperature (
float) — 샘플링 온도. - top_p (
float, 선택 사항) — Nucleus sampling cutoff. - top_k (
int, 선택 사항) — Top-k 샘플링의 절단 값. - sampling_method (
str) — 샘플링 방법 (auto,greedy,multinomial). - threshold (
float, 선택 사항) — 토큰 커밋을 위한 신뢰도 임계값. 기본값은 config 값입니다. - editing_threshold (
float, 선택 사항) — 마스크가 아닌 토큰을 편집할 때의 신뢰도 임계값입니다. 편집을 활성화하려면 양수여야 합니다. 기본값은 config 값입니다. - minimal_topk (
int, 선택 사항) — 각 단계에서 커밋할 최소 토큰 수. 기본값은 config 값입니다. - prompt_mask (
torch.BoolTensor, 선택 사항) —(block_length,)모양의 불리언 마스크로,True는 프롬프트(편집 불가) 위치를 나타냅니다. - generator (
torch.Generator, 선택 사항) — 샘플링을 위한 RNG입니다. - return_dict (
bool) —BlockRefinementSchedulerOutput를 반환할지, 아니면 튜플을 반환할지를 결정합니다.
…
BlockRefinementSchedulerOutput
클래스 diffusers.BlockRefinementSchedulerOutput
< source >( prev_sample: torch.LongTensortransfer_index: torch.BoolTensorediting_transfer_index: torch.BoolTensorsampled_tokens: torch.LongTensorsampled_probs: torch.Tensorpred_logits: torch.Tensor )
매개변수
- prev_sample (
torch.LongTensor, shape(batch_size, block_length)) — 현재 refinement 단계 이후 업데이트된 블록 토큰. - transfer_index (
torch.BoolTensor, shape(batch_size, block_length)) — 커밋된 토큰을 나타내는 불리언 마스크(mask-filling). - editing_transfer_index (
torch.BoolTensor, shape(batch_size, block_length)) — 편집된 토큰(비마스크 치환)을 나타내는 불리언 마스크. - sampled_tokens (
torch.LongTensor, shape(batch_size, block_length)) — 모델 logits에서 샘플링된 토큰 ID입니다. - sampled_probs (
torch.Tensor, shape(batch_size, block_length)) — 샘플링된 토큰의 확률입니다. - pred_logits (
torch.Tensor, shape(batch_size, block_length, vocab_size)) — 다음 단계의 self-conditioning을 위해 전달되는 denoiser logits입니다.
블록 리파인먼트 스케줄링을 위한 출력 클래스.