# LLaDA2

[LLaDA2](https://huggingface.co/collections/inclusionAI/llada21)は、ブロック単位の反復的な改善を通じてテキストを生成する、離散拡散言語モデルのファミリーです。LLaDA2は、自己回帰的なトークン単位の生成ではなく、完全にマスクされたシーケンスから開始し、複数の改善ステップにわたり、信頼度に基づいてトークンを段階的にアンマスクしていきます。

## 使用方法

```py
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 は各リファインメントステップの後に実行されます。`callback_on_step_end_tensor_inputs` を渡すことで、`callback_kwargs` に含めるテンソルを選択できます。現在の実装では、`block_x`（リファインメント対象のシーケンスウィンドウ）と `transfer_index`（マスク埋込コミットマスク）が提供されます。Callback から `{"block_x": ...}` を返すことで、ウィンドウを置き換えることができます。

```py
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つのモードをサポートしています：

| モード | `threshold` | `editing_threshold` | `max_post_steps` |
|------|-------------|---------------------|------------------|
| 品質 | 0.7 | 0.5 | 16 |
| 速度 | 0.5 | `None` | 16 |

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

#### diffusers.LLaDA2Pipeline[[diffusers.LLaDA2Pipeline]]

```python
diffusers.LLaDA2Pipeline(model: Any, scheduler: BlockRefinementScheduler, tokenizer: Any | None = None)
```

[ソース](https://github.com/huggingface/diffusers/blob/v0.43.0/src/diffusers/pipelines/llada2/pipeline_llada2.py#L59)

ブロック単位の反復的な改善による、LLaDA2 風の離散拡散テキスト生成のパイプライン。

このパイプラインは、`mask_token_id`で埋められたテンプレートシーケンスを維持し、ブロック単位でそれを洗練していきます。各洗練ステップでは、アクティブなブロックの候補トークンをサンプリングし、信頼度に基づいてその一部を確定します。

モデルは、attention mask と `position_ids` を受け取り、`[batch, seq, vocab_size]` の形状の logits を返すことが期待されます。

#### __call__[[diffusers.LLaDA2Pipeline.__call__]]

```python
__call__(prompt: str | list[str] | None = None, messages: list[dict[str, str]] | None = None, input_ids: torch.LongTensor | None = None, attention_mask: torch.LongTensor | None = None, use_chat_template: bool = True, add_generation_prompt: bool = True, gen_length: int = 2048, block_length: int | None = None, num_inference_steps: int = 32, temperature: float = 0.0, top_p: float | None = None, top_k: int | None = None, sampling_method: str = 'multinomial', threshold: float = 0.7, editing_threshold: float | None = 0.5, max_post_steps: int = 16, minimal_topk: int = 1, eos_early_stop: bool = True, eos_token_id: int | None = None, mask_token_id: int | None = None, generator: torch.Generator | None = None, output_type: str = 'text', return_dict: bool = True, callback_on_step_end: Callable[[int, int, dict], None] | PipelineCallback | MultiPipelineCallbacks | None = None, callback_on_step_end_tensor_inputs: list[str] | None = None)
```

[ソース](https://github.com/huggingface/diffusers/blob/v0.43.0/src/diffusers/pipelines/llada2/pipeline_llada2.py#L242)

**パラメータ:**

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`, *任意*) : ニュクレアサンプリングのカットオフ値。

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`, *任意*) : 早期停止に使用する EOS トークン ID。

mask_token_id (`int`, *任意*) : テンプレートに使用するマスクトークンID。

generator (`torch.Generator`, *任意*) : サンプリング用の乱数生成器。

output_type (`str`、デフォルトは`"text"`) : 出力形式。`"text"`はシーケンスを文字列にデコードします（トークナイザーが必要です）。`"seq"`は生のトークンIDシーケンスのみを返します。

return_dict (`bool`、*任意*、デフォルトは`True`) : タプルではなく[LLaDA2PipelineOutput](/docs/diffusers/v0.43.0/ja/api/pipelines/llada2#diffusers.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]`、*任意*) : コールバックに渡すTensorのキー。許可されるキー: `block_x`、`transfer_index`、`editing_transfer_index`、`sampled_tokens`、`sampled_probs`、`active_block`。

**戻り値:** [LLaDA2PipelineOutput](/docs/diffusers/v0.43.0/ja/api/pipelines/llada2#diffusers.LLaDA2PipelineOutput) または `tuple`

`return_dict` が `True` の場合、[LLaDA2PipelineOutput](/docs/diffusers/v0.43.0/ja/api/pipelines/llada2#diffusers.LLaDA2PipelineOutput) が返され、それ以外の場合は、最初の要素が生成されたトークン ID（`torch.LongTensor`）、2番目の要素がデコードされたテキスト（`list[str]`）である `tuple` が返されます。`output_type` が `"seq"` の場合は `None` が返されます。

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

例:
```python
>>> 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]]

#### diffusers.LLaDA2PipelineOutput[[diffusers.LLaDA2PipelineOutput]]

```python
diffusers.LLaDA2PipelineOutput(sequences: torch.LongTensor, texts: list[str] | None = None)
```

[ソース](https://github.com/huggingface/diffusers/blob/v0.43.0/src/diffusers/pipelines/llada2/pipeline_llada2.py#L54)