Hugging Face's logo

Diffusers documentation

LLaDA2

Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

LLaDA2

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

Использование

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

from diffusers import BlockRefinementScheduler, LLaDA2Pipeline

model_id = "inclusionAI/LLaDA2.1-mini"
model = AutoModelForCausalLM.from_pretrained(
    model_id, trust_remote_code=True, dtype=torch.bfloat16, device_map="auto"
)
tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
scheduler = BlockRefinementScheduler()

pipe = LLaDA2Pipeline(model=model, scheduler=scheduler, tokenizer=tokenizer)
output = pipe(
    prompt="Write a short poem about the ocean.",
    gen_length=256,
    block_length=32,
    num_inference_steps=32,
    threshold=0.7,
    editing_threshold=0.5,
    max_post_steps=16,
    temperature=0.0,
)
print(output.texts[0])

Обработчики событий

Обработчики (callbacks) выполняются после каждого шага уточнения. Передайте callback_on_step_end_tensor_inputs, чтобы выбрать, какие тензоры будут включены в callback_kwargs. В текущей реализации предоставляются block_x (окно последовательности, которое уточняется) и transfer_index (маска фиксации заполнения маски); верните из обработчика {"block_x": ...}, чтобы заменить окно.

def on_step_end(pipe, step, timestep, callback_kwargs):
    block_x = callback_kwargs["block_x"]
    # Inspect or modify `block_x` here.
    return {"block_x": block_x}

out = pipe(
    prompt="Write a short poem.",
    callback_on_step_end=on_step_end,
    callback_on_step_end_tensor_inputs=["block_x"],
)

Рекомендуемые параметры

Модели LLaDA2.1 поддерживают два режима:

Режимthresholdediting_thresholdmax_post_steps
Качество0.70.516
Скорость0.5None16

Передайте editing_threshold=None, 0.0 или отрицательное значение, чтобы отключить редактирование после маски.

Для моделей LLaDA2.0 отключите редактирование, передав editing_threshold=None или 0.0.

Для всех моделей: block_length=32, temperature=0.0, num_inference_steps=32.

LLaDA2Pipeline

класс diffusers.LLaDA2Pipeline

< >

( model: Anyscheduler: BlockRefinementSchedulertokenizer: Any | None = None )

Пайплайн для генерации дискретного диффузионного текста в стиле LLaDA2 с помощью блочного итеративного уточнения.

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

Ожидается, что модель будет принимать attention mask и position_ids и возвращать logits формы [batch, seq, vocab_size].

__call__

< >

( prompt: str | list[str] | None = Nonemessages: list[dict[str, str]] | None = Noneinput_ids: torch.LongTensor | None = Noneattention_mask: torch.LongTensor | None = Noneuse_chat_template: bool = Trueadd_generation_prompt: bool = Truegen_length: int = 2048block_length: int | None = Nonenum_inference_steps: int = 32temperature: float = 0.0top_p: float | None = Nonetop_k: int | None = Nonesampling_method: str = 'multinomial'threshold: float = 0.7editing_threshold: float | None = 0.5max_post_steps: int = 16minimal_topk: int = 1eos_early_stop: bool = Trueeos_token_id: int | None = Nonemask_token_id: int | None = Nonegenerator: torch.Generator | None = Noneoutput_type: str = 'text'return_dict: bool = Truecallback_on_step_end: Callable[[int, int, dict], None] | PipelineCallback | MultiPipelineCallbacks | None = Nonecallback_on_step_end_tensor_inputs: list[str] | None = None ) → LLaDA2PipelineOutput or tuple

Параметры

  • prompt (str или List[str], необязательный) — Текст запроса. Если use_chat_template установлен в True (по умолчанию) и доступен токенизатор с шаблоном чата, запрос оборачивается в сообщение чата перед токенизацией.
  • messages (List[Dict[str, str]], необязательный) — Сообщения чата для кодирования (например, [{"role": "user", "content": "Hello"}]). Имеет приоритет над prompt, если задано. Требует токенизатора с apply_chat_template.
  • input_ids (torch.LongTensor, необязательный) — Предварительно токенизированные ID входных данных. Имеет приоритет над prompt и messages.
  • attention_mask (torch.LongTensor, необязательный) — Маска для каждого токена (1 для валидных токенов промпта, 0 для заполнителей), совпадающая по форме с input_ids. Используется только при передаче input_ids. Если параметр опущен (а input_ids передан), все позиции считаются валидными. При формировании входных данных из prompt / messages маска токенизатора передаётся автоматически.
  • use_chat_template (bool, по умолчанию True) — Определяет, нужно ли обернуть запрос в шаблон чата.
  • add_generation_prompt (bool, по умолчанию True) — Добавлять ли промпт генерации при использовании шаблонов чата.
  • gen_length (int) — Количество токенов для генерации.
  • block_length (int, необязательный) — Размер блока для уточнения. Если не указан, используется block_length, настроенный в планировщике.
  • num_inference_steps (int) — Количество шагов уточнения на блок.
  • temperature (float) — Температура выборки.
  • top_p (float, необязательный) — Порог обрезки для ядерной выборки.
  • top_k (int, необязательный) — Порог выборки Top-k.
  • sampling_method (str) — Метод выборки (auto, greedy, multinomial).
  • threshold (float) — Порог уверенности для фиксации токенов.
  • editing_threshold (float, необязательный) — Порог уверенности для редактирования уже зафиксированных (не маскирующих) токенов. Если значение положительное, после разрешения всех маскирующих токенов в блоке конвейер продолжает уточнение: если модель предсказывает другой токен с уверенностью выше этого порога, существующий токен заменяется. Установите None, 0.0 или отрицательное значение, чтобы отключить редактирование. По умолчанию 0.5.
  • max_post_steps (int) — Максимальное количество дополнительных итераций уточнения после разрешения всех токенов маски в блоке. Используется только при включённом editing_threshold. По умолчанию 16.
  • minimal_topk (int) — Минимальное количество токенов, которые нужно зафиксировать на каждом шаге.
  • eos_early_stop (bool) — Определяет, следует ли останавливаться после фиксации EOS в блоке.
  • eos_token_id (int, необязательный) — ID токена EOS, используемого для раннего завершения.
  • mask_token_id (int, необязательный) — ID маскирующего токена, используемого в шаблоне.
  • generator (torch.Generator, необязательный) — Генератор случайных чисел для выборки.
  • output_type (str, по умолчанию "text") — Формат вывода. "text" декодирует последовательности в строки (требуется токенизатор). "seq" возвращает только исходные последовательности ID токенов.
  • return_dict (bool, необязательный, по умолчанию True) — Определяет, следует ли возвращать LLaDA2PipelineOutput вместо кортежа.
  • callback_on_step_end (Callable or PipelineCallback, необязательный) — Callback, вызываемый после каждого шага уточнения с сигнатурой callback_on_step_end(self, step: int, timestep: int, callback_kwargs: Dict).
  • callback_on_step_end_tensor_inputs (List[str], необязательный) — Ключи тензоров, передаваемые в callback. Допустимые ключи: block_x, transfer_index, editing_transfer_index, sampled_tokens, sampled_probs, active_block.

Возвращает

LLaDA2PipelineOutput или tuple

Если return_dict равно True, возвращается LLaDA2PipelineOutput, в противном случае возвращается tuple, где первый элемент — это сгенерированные идентификаторы токенов (torch.LongTensor), а второй — декодированные тексты (list[str]), либо None, если output_type равно "seq".

Генерация текста с покомпонентным уточнением.

Примеры:

>>> import torch
>>> from transformers import AutoModelForCausalLM, AutoTokenizer
>>> from diffusers import BlockRefinementScheduler, LLaDA2Pipeline

>>> model_id = "inclusionAI/LLaDA2.1-mini"
>>> model = AutoModelForCausalLM.from_pretrained(
...     model_id, trust_remote_code=True, dtype=torch.bfloat16, device_map="auto"
... )
>>> tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
>>> scheduler = BlockRefinementScheduler()

>>> pipe = LLaDA2Pipeline(model=model, scheduler=scheduler, tokenizer=tokenizer)
>>> output = pipe(prompt="What is the meaning of life?", gen_length=256)
>>> print(output.texts[0])

LLaDA2PipelineOutput

класс diffusers.LLaDA2PipelineOutput

< >

( sequences: torch.LongTensortexts: list[str] | None = None )

Update on GitHub

←DiffusionGemma Allegro→