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

コールバック

コールバックは、各リファインメントステップの後に実行されます。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→