Hugging Face's logo

Diffusers documentation

LLaDA2

You are viewing main version, which requires installation from source. If you'd like regular pip install, checkout the latest stable version (v0.43.0).
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→