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 )
Планировщик для блочного итеративного уточнения (коммит по степени уверенности).
На каждом шаге планировщик выбирает кандидаты токены из логитов модели и фиксирует те, у которых самая высокая уверенность. Количество токенов, фиксируемых на каждом шаге, определяется равномерным распределением длины блока на количество шагов уточнения.
Необязательно поддерживает редактирование: после разрешения всех маскирующих токенов токены могут быть заменены, если модель предсказывает другой токен с уверенностью выше положительного 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)) — Чистые идентификаторы токенов. - attention_mask (
torch.LongTensorформы(batch_size, seq_len)) — Маска отсечения (1 для валидных токенов, 0 для отсечения). - prompt_length (
int) — Количество начальных токенов промпта, которые нужно оставить без маски. - block_length (
int) — Размер блока для маскирования. - mask_token_id (
int) — Идентификатор токена, используемый для замаскированных позиций. - generator (
torch.Generator, необязательный) — Генератор случайных чисел для воспроизводимости.
Возвращает
tuple[torch.LongTensor, torch.LongTensor, torch.BoolTensor, torch.BoolTensor]
(noisy, noisy_rev, masked, masked_rev) — две взаимодополняющие зашумленные последовательности и соответствующие им булевы маски.
Примените прямой (шумовой) процесс для маскирования блоков в полуавто-регрессивном режиме.
Для каждого блока после промпта случайная доля валидных (не служебных) токенов заменяется на 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) — Количество выполненных на данный момент шагов редактирования после маски. - 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формы(batch_size, seq_len)) — Текущая полная последовательность, включающая все блоки до текущего окна. - sampled_tokens (
torch.LongTensorформы(batch_size, block_length)) — Токены, выбранные планировщиком на этом шаге. - final_transfer (
torch.BoolTensorформы(batch_size, block_length)) — Совмещенная маска зафиксированных и отредактированных позиций. - finished (
torch.BoolTensorформы(batch_size,)) — Текущие флаги завершения для каждой партии. - eos_token_id (
int) — ID токена EOS. - 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)) — Сырые логиты модели для текущего блока. - timestep (
intилиtorch.Tensor) — Текущий индекс шага в расписании уточнения блока. - sample (
torch.LongTensorформы(batch_size, block_length)) — ID токенов текущего блока (в режиме на основе маски содержит маскирующие токены для неподтвержденных позиций). - mask_token_id (
int, необязательный) — Идентификатор токена, используемый для замаскированных позиций. ЕслиNone, планировщик работает в режиме равномерного повреждения: он внутренне отслеживает зафиксированные позиции (сбрасывая их приtimestep == 0) и добавляет шум к незафиксированным позициям с помощью равномерно случайных токенов, что соответствует блочному сэмплеру уточнения DiffusionGemma. - temperature (
float) — Температура выборки. - top_p (
float, необязательный) — Порог обрезки для ядерной выборки. - top_k (
int, необязательный) — Порог выборки Top-k. - sampling_method (
str) — Метод выборки (auto,greedy,multinomial). - threshold (
float, необязательный) — Порог уверенности для фиксации токенов. По умолчанию используется значение из конфигурации. - editing_threshold (
float, необязательный) — Порог уверенности для редактирования токенов вне маски; должен быть положительным, чтобы включить редактирование. По умолчанию используется значение из конфигурации. - minimal_topk (
int, необязательный) — Минимальное количество токенов для фиксации на каждом шаге. По умолчанию используется значение из конфигурации. - prompt_mask (
torch.BoolTensor, необязательный) — Булева маска формы(block_length,), гдеTrueотмечает позиции промпта (не редактируемые). - generator (
torch.Generator, необязательный) — Генератор случайных чисел для выборки. - 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формы(batch_size, block_length)) — Обновленные токены блока после текущего шага уточнения. - transfer_index (
torch.BoolTensorформы(batch_size, block_length)) — Булева маска, указывающая, какие токены были зафиксированы (mask-filling). - editing_transfer_index (
torch.BoolTensorформы(batch_size, block_length)) — Булева маска, указывающая, какие токены были отредактированы (замена без маски). - sampled_tokens (
torch.LongTensorформы(batch_size, block_length)) — ID сэмплированных токенов из логитов модели. - sampled_probs (
torch.Tensorформы(batch_size, block_length)) — Вероятности сэмплированных токенов. - pred_logits (
torch.Tensorформы(batch_size, block_length, vocab_size)) — Логиты денойзера, передаваемые для самоусловливания следующего шага.
Класс вывода для планирования уточнения блоков.