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

BlockRefinementScheduler

BlockRefinementScheduler управляет блочным итеративным уточнением для дискретной диффузии токенов. На каждом шаге он фиксирует наиболее уверенные токены и при необходимости редактирует уже зафиксированные токены, если модель предсказывает другой токен с высокой степенью уверенности.

Этот планировщик используется в 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)

Источник

Планировщик для блочного итеративного уточнения (коммит по степени уверенности).

На каждом шаге планировщик выбирает кандидаты токены из логитов модели и фиксирует те, у которых уверенность максимальна. Количество токенов, фиксируемых на каждом шаге, определяется равномерным распределением длины блока на количество шагов уточнения.

Необязательно поддерживает редактирование: после разрешения всех маскирующих токенов токены могут быть заменены, если модель предсказывает другой токен с уверенностью выше положительного editing_threshold (None, 0.0 или отрицательное значение отключает редактирование).

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)

Источник

Параметры:

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[[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)

Источник

Параметры:

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[[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)

Источник

Параметры:

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.

get_num_transfer_tokens[[diffusers.BlockRefinementScheduler.get_num_transfer_tokens]]

get_num_transfer_tokens(block_length: int, num_inference_steps: int)

Источник

Равномерно распределите коммиты токенов block_length по num_inference_steps шагам.

шаг[[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)

Источник

Параметры:

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, optional) : Идентификатор токена, используемый для замаскированных позиций. Если None, планировщик работает в режиме равномерного повреждения: он внутренне отслеживает зафиксированные позиции (сбрасывая их при timestep == 0) и добавляет шум к незафиксированным позициям с помощью равномерно случайных токенов, что соответствует блочному сэмплеру уточнения DiffusionGemma.

temperature (float) : Температура выборки.

top_p (float, необязательный) : Порог обрезки для ядерной выборки.

top_k (int, optional) : Порог выборки Top-k.

sampling_method (str) : Метод выборки (auto, greedy, multinomial).

threshold (float, необязательный) : Порог уверенности для фиксации токенов. По умолчанию используется значение из конфигурации.

editing_threshold (float, optional) : Порог уверенности для редактирования токенов вне маски; должен быть положительным для включения редактирования. По умолчанию используется значение из конфигурации.

minimal_topk (int, optional) : Минимальное количество токенов для фиксации на каждом шаге. По умолчанию используется значение из конфигурации.

prompt_mask (torch.BoolTensor, необязательный) : Булева маска формы (block_length,), где True отмечает позиции промпта (не редактируемые).

generator (torch.Generator, необязательный) : генератор случайных чисел для выборки.

return_dict (bool) : Нужно ли возвращать BlockRefinementSchedulerOutput или кортеж.

Выполните один шаг уточнения: сгенерируйте токены из логитов, зафиксируйте уверенные токены и при необходимости отредактируйте существующие.

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)

Источник

Параметры:

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)) : Сэмплированные идентификаторы токенов из логитов модели.

sampled_probs (torch.Tensor формы (batch_size, block_length)) : Вероятности сэмплированных токенов.

pred_logits (torch.Tensor формы (batch_size, block_length, vocab_size)) : Логиты денойзера, передаваемые для самоусловливания следующего шага.

Класс вывода для планирования уточнения блоков.