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.42.1).
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는 각 refinement 단계 이후에 실행됩니다. callback_on_step_end_tensor_inputs를 전달하여 callback_kwargs에 포함될 텐서를 선택합니다. 현재 구현에서는 block_x(refinement 중인 시퀀스 윈도우)와 transfer_index(mask-filling 커밋 마스크)가 제공됩니다. callback에서 {"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를 받아들이고, [batch, seq, vocab_size] 형태의 logits를 반환해야 합니다.

__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, 선택 사항) — input_ids의 형상과 일치하는 토큰별 마스크(유효한 프롬프트 토큰은 1, 패딩은 0). input_ids가 제공될 때만 사용됩니다. 생략된 경우(input_ids가 주어졌을 때), 모든 위치가 유효한 것으로 처리됩니다. prompt / messages에서 입력을 구성할 때, 토크나이저의 마스크가 자동으로 전달됩니다.
  • use_chat_template (bool, 기본값 True) — 프롬프트를 채팅 템플릿으로 감싸는지 여부.
  • add_generation_prompt (bool, 기본값 True) — 채팅 템플릿을 사용할 때 생성 프롬프트를 추가할지 여부입니다.
  • gen_length (int) — 생성할 토큰의 수.
  • block_length (int, 선택 사항) — 정제(refinement)에 사용할 블록 크기입니다. 제공되지 않으면 스케줄러에 설정된 block_length가 사용됩니다.
  • num_inference_steps (int) — 블록당 정제 단계 수.
  • temperature (float) — 샘플링 온도.
  • top_p (float, 선택 사항) — Nucleus sampling cutoff.
  • top_k (int, 선택 사항) — Top-k 샘플링의 절단 값.
  • sampling_method (str) — 샘플링 방법 (auto, greedy, multinomial).
  • threshold (float) — 토큰 커밋을 위한 신뢰도 임계값.
  • editing_threshold (float, 선택 사항) — 이미 커밋된(마스크가 아닌) 토큰을 편집할 때 사용하는 신뢰도 임계값입니다. 양수인 경우, 블록 내 모든 마스크 토큰이 해결된 후 파이프라인은 정제(refining)를 계속합니다. 모델이 이 임계값보다 높은 신뢰도로 다른 토큰을 예측하면 기존 토큰이 교체됩니다. 편집을 비활성화하려면 None, 0.0 또는 음수 값을 설정하세요. 기본값은 0.5입니다.
  • max_post_steps (int) — 블록의 모든 mask 토큰이 해결된 후 수행할 추가 정제 반복의 최대 횟수. editing_threshold가 활성화된 경우에만 사용됩니다. 기본값은 16입니다.
  • minimal_topk (int) — 각 단계에서 커밋할 최소 토큰 수.
  • eos_early_stop (bool) — 블록에서 EOS를 커밋한 후 중단할지 여부.
  • eos_token_id (int, 선택 사항) — 조기 정지를 사용할 EOS 토큰 ID입니다.
  • mask_token_id (int, 선택 사항) — 템플릿에 사용할 마스크 토큰 ID입니다.
  • generator (torch.Generator, 선택 사항) — 샘플링을 위한 RNG입니다.
  • output_type (str, 기본값 "text") — 출력 형식입니다. "text"는 시퀀스를 문자열로 디코딩합니다(토크나이저 필요). "seq"는 원시 토큰 ID 시퀀스만 반환합니다.
  • return_dict (bool, 선택 사항, 기본값 True) — tuple 대신 LLaDA2PipelineOutput를 반환할지 여부.
  • callback_on_step_end (Callable 또는 PipelineCallback, 선택 사항) — 각 개선 단계 이후에 실행되는 콜백으로, 시그니처는 callback_on_step_end(self, step: int, timestep: int, callback_kwargs: Dict)입니다.
  • callback_on_step_end_tensor_inputs (List[str], 선택 사항) — 콜백에 전달할 텐서 키입니다. 허용되는 키: block_x, transfer_index, editing_transfer_index, sampled_tokens, sampled_probs, active_block.

반환값

LLaDA2PipelineOutput 또는 tuple

return_dict가 True인 경우 LLaDA2PipelineOutput가 반환되고, 그렇지 않으면 첫 번째 요소가 생성된 토큰 ID(torch.LongTensor)이고 두 번째 요소가 디코딩된 텍스트(list[str])인 tuple이 반환됩니다. output_type이 "seq"인 경우 두 번째 요소는 None입니다.

block-wise refinement을 사용하여 텍스트를 생성합니다.

예:

>>> 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→