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 là một họ mô hình ngôn ngữ khuếch tán rời rạc, tạo văn bản thông qua tinh chỉnh lặp lại theo khối. Thay vì tạo token theo cách tự hồi quy từng token một, LLaDA2 bắt đầu với một chuỗi bị che hoàn toàn và dần dần bỏ che các token theo độ tin cậy qua nhiều bước tinh chỉnh.

Cách sử dụng

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

Callbacks được thực thi sau mỗi bước tinh chỉnh. Truyền callback_on_step_end_tensor_inputs để chọn các tensor nào được bao gồm trong callback_kwargs. Trong cài đặt hiện tại, block_x (cửa sổ chuỗi đang được tinh chỉnh) và transfer_index (mask cam kết lấp đầy mask) được cung cấp; trả về {"block_x": ...} từ callback để thay thế cửa sổ.

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"],
)

Các tham số được khuyến nghị

Mô hình LLaDA2.1 hỗ trợ hai chế độ:

Chế độthresholdediting_thresholdmax_post_steps
Chất lượng0.70.516
Tốc độ0.5None16

Truyền editing_threshold=None, 0.0 hoặc một giá trị âm để tắt chỉnh sửa sau mask.

Đối với các mô hình LLaDA2.0, hãy tắt tính năng chỉnh sửa bằng cách truyền editing_threshold=None hoặc 0.0.

Đối với tất cả các mô hình: block_length=32, temperature=0.0, num_inference_steps=32.

LLaDA2Pipeline

lớp diffusers.LLaDA2Pipeline

< >

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

Pipeline cho việc tạo văn bản khuếch tán rời rạc theo phong cách LLaDA2 thông qua tinh chỉnh lặp lại theo khối.

Pipeline này duy trì một chuỗi mẫu được lấp đầy bằng mask_token_id và tinh chỉnh nó theo từng khối. Ở mỗi bước tinh chỉnh, nó lấy mẫu các token ứng viên cho khối đang hoạt động và cam kết một tập con dựa trên độ tin cậy.

Mô hình được kỳ vọng sẽ nhận một attention mask và position_ids, và trả về logits có hình dạng [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

Tham số

  • prompt (str hoặc List[str], tùy chọn) — Văn bản prompt. Khi use_chat_template là True (mặc định) và có sẵn một tokenizer với chat template, prompt sẽ được bọc trong một tin nhắn chat trước khi được token hóa.
  • messages (List[Dict[str, str]], tùy chọn) — Các tin nhắn chat cần mã hóa (ví dụ: [{"role": "user", "content": "Hello"}]). Có độ ưu tiên cao hơn prompt khi được cung cấp. Yêu cầu một tokenizer có apply_chat_template.
  • input_ids (torch.LongTensor, tùy chọn) — Các ID đầu vào đã được phân tách token. Có quyền ưu tiên cao hơn prompt và messages.
  • attention_mask (torch.LongTensor, tùy chọn) — Mặt nạ theo từng token (1 cho các token prompt hợp lệ, 0 cho padding) khớp với hình dạng của input_ids. Chỉ được sử dụng khi input_ids được cung cấp. Khi bị bỏ qua (và input_ids được đưa vào), tất cả các vị trí đều được coi là hợp lệ. Khi xây dựng đầu vào từ prompt / messages, mặt nạ của tokenizer sẽ được truyền qua tự động.
  • use_chat_template (bool, mặc định là True) — Có nên bọc prompt trong một chat template hay không.
  • add_generation_prompt (bool, mặc định là True) — Có thêm prompt sinh khi sử dụng chat templates hay không.
  • gen_length (int) — Số lượng token cần tạo.
  • block_length (int, tùy chọn) — Kích thước khối cho quá trình tinh chỉnh. Nếu không được cung cấp, block_length đã cấu hình của trình lập lịch sẽ được sử dụng.
  • num_inference_steps (int) — Số bước tinh chỉnh cho mỗi khối.
  • temperature (float) — Nhiệt độ lấy mẫu.
  • top_p (float, tùy chọn) — Ngưỡng cắt của nucleus sampling.
  • top_k (int, tùy chọn) — Ngưỡng cắt cho lấy mẫu top-k.
  • sampling_method (str) — Phương pháp lấy mẫu (auto, greedy, multinomial).
  • threshold (float) — Ngưỡng độ tin cậy để cam kết các token.
  • editing_threshold (float, tùy chọn) — Ngưỡng độ tin cậy để chỉnh sửa các token đã cam kết (không phải mask). Khi có giá trị dương, sau khi tất cả các token mask trong một khối được giải quyết, pipeline tiếp tục tinh chỉnh: nếu mô hình dự đoán một token khác với độ tin cậy cao hơn ngưỡng này, token hiện tại sẽ được thay thế. Đặt thành None, 0.0, hoặc một giá trị âm để vô hiệu hóa việc chỉnh sửa. Mặc định là 0.5.
  • max_post_steps (int) — Số lần lặp cải tiến bổ sung tối đa sau khi tất cả các token mask trong một khối được giải quyết. Chỉ được sử dụng khi editing_threshold được bật. Mặc định là 16.
  • minimal_topk (int) — Số lượng token tối thiểu cần cam kết ở mỗi bước.
  • eos_early_stop (bool) — Có nên dừng lại sau khi cam kết EOS trong một khối hay không.
  • eos_token_id (int, tùy chọn) — ID của token EOS để sử dụng cho việc dừng sớm.
  • mask_token_id (int, tùy chọn) — ID của token mask để sử dụng cho template.
  • generator (torch.Generator, tùy chọn) — RNG để lấy mẫu.
  • output_type (str, mặc định là "text") — Định dạng đầu ra. "text" giải mã các chuỗi thành chuỗi ký tự (yêu cầu có tokenizer). "seq" chỉ trả về các chuỗi ID token thô.
  • return_dict (bool, tùy chọn, mặc định là True) — Có nên trả về một LLaDA2PipelineOutput thay vì một tuple hay không.
  • callback_on_step_end (Callable hoặc PipelineCallback, tùy chọn) — Callback được thực thi sau mỗi bước tinh chỉnh với chữ ký callback_on_step_end(self, step: int, timestep: int, callback_kwargs: Dict).
  • callback_on_step_end_tensor_inputs (List[str], tùy chọn) — Các khóa tensor để truyền cho callback. Các khóa được phép: block_x, transfer_index, editing_transfer_index, sampled_tokens, sampled_probs, active_block.

Trả về

LLaDA2PipelineOutput hoặc tuple

Nếu return_dict là True, LLaDA2PipelineOutput sẽ được trả về, ngoài ra một tuple sẽ được trả về trong đó phần tử đầu tiên là các ID token đã tạo (torch.LongTensor) và phần tử thứ hai là các văn bản đã giải mã (list[str]), hoặc None khi output_type là "seq".

Tạo văn bản với tinh chỉnh theo khối.

Ví dụ:

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

lớp diffusers.LLaDA2PipelineOutput

< >

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

Update on GitHub

←DiffusionGemma Allegro→