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

コールバック

コールバックは、各リファインメントステップの後に実行されます。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モデルは2つのモードをサポートしています：

モード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

class 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, 任意) — リファインメント用のブロックサイズ。指定しない場合、スケジューラーに設定された block_length が使用されます。
  • num_inference_steps (int) — ブロックごとのリファインメントステップ数。
  • temperature (float) — サンプリング温度。
  • top_p (float, 任意) — Nucleus sampling のカットオフ値。
  • 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) — 1ステップごとにコミットするトークンの最小数。
  • eos_early_stop (bool) — ブロック内で EOS をコミットした後に停止するかどうか。
  • eos_token_id (int, 省略可) — 早期停止に使用する EOS トークン ID。
  • 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 または 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)、2番目の要素がデコードされたテキスト (list[str])、または output_type が "seq" の場合は None となる tuple が返されます。

ブロック単位の改善でテキストを生成します。

例:

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

class diffusers.LLaDA2PipelineOutput

< >

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

Update on GitHub

←DiffusionGemma Allegro→