Diffusers documentation
BlockRefinementScheduler
BlockRefinementScheduler
BlockRefinementSchedulerは、離散トークン拡散におけるブロック単位の反復的リファインメントを管理します。各ステップで最も確信度の高いトークンを確定し、モデルが別のトークンを高い確信度で予測している場合、既に確定済みのトークンを編集するオプションも提供します。
このスケジューラーは、LLaDA2Pipelineで使用されます。
BlockRefinementScheduler
class diffusers.BlockRefinementScheduler
< source >( block_length: int = 32num_inference_steps: int = 32threshold: float = 0.95editing_threshold: float | None = Noneminimal_topk: int = 1 )
ブロック単位の反復的改善（コミットバイコンフィデンス）用のスケジューラー。
各ステップにおいて、スケジューラーはモデルのlogitsから候補トークンをサンプリングし、最も信頼度の高いものを確定します。1ステップあたりに確定するトークン数は、ブロック長をリファインメントステップ数で均等に分配することによって決定されます。
オプションとして編集をサポートします。すべてのマスクトークンが解決された後、モデルが正の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、形状(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、任意) — 再現性のために使用する乱数生成器。
戻り値
tuple[torch.LongTensor, torch.LongTensor, torch.BoolTensor, torch.BoolTensor]
(noisy, noisy_rev, masked, masked_rev) — 2つの補完的なノイジーシーケンスと、それに対応するブールマスク。
半自己回帰ブロックマスクングに対して、フォワード（ノイジング）プロセスを適用します。
プロンプト以降の各ブロックにおいて、有効な（パディングでない）トークンのランダムな割合がmask_token_idに置き換えられます。2つの補完的なビュー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) — これまでに実行されたマスク後編集ステップの数。 - max_post_steps (
int) — マスク後の編集ステップの最大許容値。 - 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、形状(batch_size, block_length, vocab_size)) — 現在のブロックに対するモデルからの生の logits。 - timestep (
intまたはtorch.Tensor) — ブロックの精細化スケジュールにおける現在のステップインデックス。 - sample (
torch.LongTensor、形状(batch_size, block_length)) — 現在のブロックのトークン ID（マスクベースモードでは、未確定の位置にマスクトークンを含む）。 - mask_token_id (
int, 省略可) — マスクされた位置に使用されるトークン ID。Noneの場合、スケジューラーは均一破損モードで動作します。このモードでは、コミット済み位置を内部で追跡し（timestep == 0でリセット）、未コミットの位置を均一ランダムなトークンで再ノイズ化します。これにより、DiffusionGemma のブロックリファインメントサンプラーと一致します。 - temperature (
float) — サンプリング温度。 - top_p (
float, 任意) — Nucleus sampling のカットオフ値。 - top_k (
int、省略可) — Top-k サンプリングの切り捨て値です。 - sampling_method (
str) — サンプリング手法（auto、greedy、multinomial）。 - threshold (
float, 任意) — トークンを確定するための信頼度しきい値です。デフォルトでは設定値が使用されます。 - editing_threshold (
float, 任意) — マスクされていないトークンを編集するための信頼度しきい値です。編集を有効にするには正の値である必要があります。デフォルトでは、config の値が使用されます。 - minimal_topk (
int, 省略可) — 各ステップでコミットするトークンの最小数。デフォルトでは設定値が使用されます。 - prompt_mask (
torch.BoolTensor、任意) —(block_length,)の形状を持つブールマスクで、Trueがプロンプト（編集不可）の位置を示します。 - generator (
torch.Generator、任意) — サンプリング用の乱数生成器。 - return_dict (
bool) —BlockRefinementSchedulerOutputを返すか、タプルを返すかを指定します。
1回のリファインメントステップを実行します: logitsからサンプリングし、確信度の高いトークンを確定し、必要に応じて既存のトークンを編集します。
BlockRefinementSchedulerOutput
class 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、形状(batch_size, block_length)）— 現在のリファインメントステップ後に更新されたブロックトークン。 - transfer_index (
torch.BoolTensor、shape(batch_size, block_length)) — コミットされたトークン（マスクの埋め込み）を示すブールマスク。 - editing_transfer_index (
torch.BoolTensor、shape(batch_size, block_length)) — 編集されたトークン（非マスク置換）を示すブールマスク。 - sampled_tokens (
torch.LongTensor、形状(batch_size, block_length)) — モデルの logits からサンプリングされたトークン ID です。 - sampled_probs (
torch.Tensor、形状(batch_size, block_length)) — サンプリングされたトークンの確率です。 - pred_logits (
torch.Tensor、shape(batch_size, block_length, vocab_size)) — 次のステップの自己条件付けのために渡される、denoiser の logits です。
ブロックリファインメントスケジューリング用の出力クラス。