# Jamba

[Jamba](https://huggingface.co/papers/2403.19887)は、総パラメータ数が52Bから398Bに及ぶ、TransformerとMambaを組み合わせたハイブリッドなmixture-of-experts (MoE) 言語モデルです。このモデルは、Transformerモデルの性能と、Mambaのようなstate space models (SSM) の効率性および長いコンテキスト（256Kトークン）の利点を組み合わせることを目指しています。

Jambaのアーキテクチャは、ブロックとレイヤーのアプローチを採用しており、TransformerとMambaのアーキテクチャを統合的に組み合わせることを可能にしています。各Jambaブロックには、attentionレイヤーまたはMambaレイヤーのいずれかが含まれ、その後、多層パーセプトロン（MLP）が続きます。これにより、全8レイヤーのうち1レイヤーがTransformerレイヤーという比率が実現されます。モデルの容量を拡大するために、MoEレイヤーが組み込まれています。

すべてのオリジナルのJambaチェックポイントは、[AI21](https://huggingface.co/ai21labs)組織の下で確認できます。

> [!TIP]
> Jamba をさまざまな言語タスクに適用する方法の例については、右側のサイドバーにある Jamba モデルをクリックしてください。

以下の例では、[Pipeline](/docs/transformers/main/ja/main_classes/pipelines#transformers.Pipeline)、[AutoModel](/docs/transformers/main/ja/model_doc/auto#transformers.AutoModel)、およびコマンドラインからテキストを生成する方法を示しています。

```python
# install optimized Mamba implementations
# !pip install mamba-ssm causal-conv1d>=1.2.0
from transformers import pipeline

pipeline = pipeline(
    task="text-generation",
    model="ai21labs/AI21-Jamba-Mini-1.6",
    device=0
)
pipeline("Plants create energy through a process known as")
```

```python
from transformers import AutoModelForCausalLM, AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained(
    "ai21labs/AI21-Jamba-Large-1.6",
)
model = AutoModelForCausalLM.from_pretrained(
    "ai21labs/AI21-Jamba-Large-1.6",
    device_map="auto",
    attn_implementation="sdpa"
)
input_ids = tokenizer("Plants create energy through a process known as", return_tensors="pt").to(model.device)

output = model.generate(**input_ids, cache_implementation="static")
print(tokenizer.decode(output[0], skip_special_tokens=True))
```

Quantization は、重みを低精度で表現することで、大規模モデルのメモリ負荷を軽減します。利用可能な Quantization バックエンドの詳細については、[Quantization](../quantization/overview) の概要を参照してください。

以下の例では、[bitsandbytes](../quantization/bitsandbytes) を使用して、重みを 8 ビットにのみ量子化しています。

```python
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig

quantization_config = BitsAndBytesConfig(load_in_8bit=True,
                                         llm_int8_skip_modules=["mamba"])

# a device map to distribute the model evenly across 8 GPUs
device_map = {'model.embed_tokens': 0, 'model.layers.0': 0, 'model.layers.1': 0, 'model.layers.2': 0, 'model.layers.3': 0, 'model.layers.4': 0, 'model.layers.5': 0, 'model.layers.6': 0, 'model.layers.7': 0, 'model.layers.8': 0, 'model.layers.9': 1, 'model.layers.10': 1, 'model.layers.11': 1, 'model.layers.12': 1, 'model.layers.13': 1, 'model.layers.14': 1, 'model.layers.15': 1, 'model.layers.16': 1, 'model.layers.17': 1, 'model.layers.18': 2, 'model.layers.19': 2, 'model.layers.20': 2, 'model.layers.21': 2, 'model.layers.22': 2, 'model.layers.23': 2, 'model.layers.24': 2, 'model.layers.25': 2, 'model.layers.26': 2, 'model.layers.27': 3, 'model.layers.28': 3, 'model.layers.29': 3, 'model.layers.30': 3, 'model.layers.31': 3, 'model.layers.32': 3, 'model.layers.33': 3, 'model.layers.34': 3, 'model.layers.35': 3, 'model.layers.36': 4, 'model.layers.37': 4, 'model.layers.38': 4, 'model.layers.39': 4, 'model.layers.40': 4, 'model.layers.41': 4, 'model.layers.42': 4, 'model.layers.43': 4, 'model.layers.44': 4, 'model.layers.45': 5, 'model.layers.46': 5, 'model.layers.47': 5, 'model.layers.48': 5, 'model.layers.49': 5, 'model.layers.50': 5, 'model.layers.51': 5, 'model.layers.52': 5, 'model.layers.53': 5, 'model.layers.54': 6, 'model.layers.55': 6, 'model.layers.56': 6, 'model.layers.57': 6, 'model.layers.58': 6, 'model.layers.59': 6, 'model.layers.60': 6, 'model.layers.61': 6, 'model.layers.62': 6, 'model.layers.63': 7, 'model.layers.64': 7, 'model.layers.65': 7, 'model.layers.66': 7, 'model.layers.67': 7, 'model.layers.68': 7, 'model.layers.69': 7, 'model.layers.70': 7, 'model.layers.71': 7, 'model.final_layernorm': 7, 'lm_head': 7}
model = AutoModelForCausalLM.from_pretrained("ai21labs/AI21-Jamba-Large-1.6",
                                                             attn_implementation="flash_attention_2",
                                             quantization_config=quantization_config,
                                             device_map=device_map)

tokenizer = AutoTokenizer.from_pretrained("ai21labs/AI21-Jamba-Large-1.6")

messages = [
   {"role": "system", "content": "You are an ancient oracle who speaks in cryptic but wise phrases, always hinting at deeper meanings."},
   {"role": "user", "content": "Hello!"},
]

input_ids = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors='pt').to(model.device)

outputs = model.generate(input_ids, max_new_tokens=216)

# Decode the output
conversation = tokenizer.decode(outputs[0], skip_special_tokens=True)

# Split the conversation to get only the assistant's response
assistant_response = conversation.split(messages[-1]['content'])[1].strip()
print(assistant_response)
# Output: Seek and you shall find. The path is winding, but the journey is enlightening. What wisdom do you seek from the ancient echoes?
```

## 注記

- モデルの性能低下を防ぐため、Mambaブロックを量子化しないでください。
- 利用可能な場合、最適化された Mamba カーネルが自動的に使用されます（`use_kernels=True` で読み込む場合は Hub カーネル、それ以外の場合はインストール済みの `mamba-ssm` および `causal-conv1d` パッケージ）。利用できない場合、モデルは PyTorch 実装にフォールバックし、レイテンシが大幅に高くなります。
- バッチ生成には左パディングを使用します。[Padding side](../llm_tutorial#padding-side)を参照してください。

## JambaConfig[[transformers.JambaConfig]]

#### transformers.JambaConfig[[transformers.JambaConfig]]

```python
transformers.JambaConfig(transformers_version: str | None = None, architectures: list[str] | None = None, output_hidden_states: bool | None = False, return_dict: bool | None = True, dtype: str | torch.dtype | None = None, chunk_size_feed_forward: int = 0, is_encoder_decoder: bool = False, id2label: dict[int, str] | dict[str, str] | None = None, label2id: dict[str, int] | dict[str, str] | None = None, problem_type: Literal['regression', 'single_label_classification', 'multi_label_classification'] | None = None, vocab_size: int = 65536, tie_word_embeddings: bool = False, hidden_size: int = 4096, intermediate_size: int = 14336, num_hidden_layers: int = 32, num_attention_heads: int = 32, num_key_value_heads: int = 8, hidden_act: str = 'silu', initializer_range: float = 0.02, rms_norm_eps: float = 1e-06, use_cache: bool = True, output_router_logits: bool = False, router_aux_loss_coef: float = 0.001, pad_token_id: int | None = 0, bos_token_id: int | None = 1, eos_token_id: int | list[int] | None = 2, max_position_embeddings: int = 262144, attention_dropout: float | int = 0.0, num_experts_per_tok: int = 2, num_experts: int = 16, expert_layer_period: int = 2, expert_layer_offset: int = 1, attn_layer_period: int = 8, attn_layer_offset: int = 4, use_mamba_kernels: bool = True, mamba_d_state: int = 16, mamba_d_conv: int = 4, mamba_expand: int = 2, mamba_dt_rank: int | str = 'auto', mamba_conv_bias: bool = True, mamba_proj_bias: bool = False, use_mambapy: bool = False, use_associative_scan: bool = True)
```

[ソース](https://github.com/huggingface/transformers/blob/main/src/transformers/models/jamba/configuration_jamba.py#L26)

**パラメータ:**

vocab_size (`int`、*任意*、デフォルトは `65536`) : モデルの語彙サイズ。`input_ids` で表すことができるトークンの種類数を定義します。

tie_word_embeddings (`bool`, *任意*, デフォルトは `False`) : モデルの `tied_weights_keys` マッピングに従って重み埋め込みを共有するかどうか。

hidden_size (`int`、*任意*、デフォルトは`4096`) : 隠れ表現の次元数。

intermediate_size (`int`、*optional*、デフォルトは`14336`) : MLP 表現の次元数。

num_hidden_layers (`int`、*任意*、デフォルトは`32`) : Transformerデコーダーの隠れ層の数。

num_attention_heads (`int`、*任意*、デフォルトは `32`) : Transformer デコーダーの各アテンション層におけるアテンションヘッドの数。

num_key_value_heads (`int`, *optional*, デフォルトは `8`) : Grouped Query Attention を実装するために使用する key_value ヘッドの数です。`num_key_value_heads=num_attention_heads` の場合、モデルは Multi Head Attention (MHA) を使用し、`num_key_value_heads=1` の場合は Multi Query Attention (MQA) を使用し、それ以外の場合は GQA が使用されます。マルチヘッドのチェックポイントを GQA のチェックポイントに変換する際は、各グループの key ヘッドと value ヘッドは、そのグループ内の元のすべてのヘッドを平均プーリングして構築する必要があります。詳細については、[この論文](https://huggingface.co/papers/2305.13245)を参照してください。指定されない場合、デフォルトで `num_attention_heads` が使用されます。

hidden_act (`str`、*任意*、デフォルトは`silu`) : デコーダーで使用する非線形活性化関数（関数または文字列）。例えば、`"gelu"`、`"relu"`、`"silu"`など。

initializer_range (`float`、*任意*、デフォルトは`0.02`) : すべての重み行列を初期化するためのtruncated_normal_initializerの標準偏差です。

rms_norm_eps (`float`、*任意*、デフォルトは`1e-06`) : rms 正規化レイヤーで使用する epsilon。

use_cache (`bool`、*任意*、デフォルトは `True`) : モデルが最後のキー/バリューアテンションを返すかどうか（すべてのモデルで利用されるわけではありません）。`config.is_decoder=True` の場合、またはモデルがデコーダー専用の生成モデルである場合にのみ関連します。

output_router_logits (`bool`, *任意*, デフォルトは `False`) : モデルがルーターのロジスティック値を返すかどうかを指定します。これを有効にすると、モデルはロードバランシングロスやルーターのzロスを含む補助ロスの出力も可能になります。

router_aux_loss_coef (`float`、*任意*、デフォルトは `0.001`) : 補助的なロードバランシング損失の係数です。MoE モデルにおける不均一なエキスパートルーティングをペナルティするために使用されます。

pad_token_id (`int`, *任意*, デフォルトは `0`) : 語彙内でパディングに使用されるトークン ID。

bos_token_id (`int`, *任意*, デフォルトは `1`) : 語彙内でストリームの開始に使用されるトークン ID。

eos_token_id (`Union[int, list[int]]`、*任意*、デフォルトは`2`) : 語彙内でストリーム終了に使用されるトークン ID。

max_position_embeddings (`int`、*optional*、デフォルトは`262144`) : このモデルが使用される可能性のある最大シーケンス長です。

attention_dropout (`Union[float, int]`、*任意*、デフォルトは`0.0`) : 注意確率のドロップアウト率。

num_experts_per_tok (`int`、*任意*、デフォルトは`2`) : 各トークンをルーティングするエキスパートの数です。これは、トークン選択ルーティングにおけるtop-kの値です。

num_experts (`int`、*任意*、デフォルトは `16`) : MoE レイヤーにおけるルーティング対象のエクスパート数。

expert_layer_period (`int`、*任意*、デフォルトは 2) : この数だけの間隔で、エキスパート層を 1 つ配置します

expert_layer_offset (`int`、*任意*、デフォルトは 1) : expert mlp レイヤーを含む最初のレイヤーのインデックス

attn_layer_period (`int`、*optional*、デフォルトは 8) : この数だけの間隔で、バニラなアテンション層を 1 つ配置します

attn_layer_offset (`int`、*任意*、デフォルトは 4) : バニラな attention mlp レイヤーを含む最初のレイヤーのインデックス

use_mamba_kernels (`bool`、*任意*、デフォルトは`True`)：非推奨であり、効果はありません。利用可能な場合、高速な Mamba カーネルが自動的に使用されます。

mamba_d_state (`int`、*任意*、デフォルトは`16`) : Mamba レイヤーにおける SSM 状態（潜在状態次元）のサイズ。

mamba_d_conv (`int`、*任意*、デフォルトは`4`) : mamba の畳み込みカーネルのサイズ

mamba_expand (`int`、*任意*、デフォルトは`2`) : mamba の中間サイズを決定するために使用される拡張係数（hidden_size に対する相対値）

mamba_dt_rank (`Union[int,str]`、*任意*、デフォルトは`"auto"`) : mamba離散化射影行列のランク。`"auto"`の場合、`math.ceil(self.hidden_size / 16)`がデフォルトとして使用されます。

mamba_conv_bias (`bool`, *任意*, デフォルトは `True`) : mamba mixer ブロックの畳み込み層でバイアスを使用するかどうかを示すフラグ。

mamba_proj_bias (`bool`, *任意*, デフォルトは `False`) : mamba mixer ブロックの入力および出力プロジェクション (["in_proj", "out_proj"]) でバイアスを使用するかどうかを示すフラグ

use_mambapy (`bool`, *任意*, デフォルトは `False`) : Mamba の CUDA ベースの公式実装が利用できない場合の、学習時のフォールバック戦略を決定します。`True` の場合、mamba.py の実装が使用されます。`False` の場合、単純で遅い実装が使用されます。メモリが限られている場合は、単純なバージョンへの切り替えを検討してください。

use_associative_scan (`bool`、*任意*、デフォルトは`True`) : 単純な逐次実装の代わりに、並列スキャンに PyTorch の `torch._higher_order_ops.associative_scan` を使用するかどうかを指定します。associative scan は `torch.compile` のトレース中のみ有効であり、torch >= 2.9.0 が必要です。両方のパスが数値的に同一の結果を生成することがテストで確認されています（`test_associative_scan_matches_sequential` を参照）。`False` に設定すると、逐次ループにフォールバックします。

これはJambaModelの設定を保存するための設定クラスです。指定された引数に基づいてJambaモデルをインスタンス化し、モデルのアーキテクチャを定義するために使用されます。デフォルト値で設定をインスタンス化すると、[ai21labs/Jamba-v0.1](https://huggingface.co/ai21labs/Jamba-v0.1) と同様の設定が得られます。

設定オブジェクトは [PreTrainedConfig](/docs/transformers/main/ja/main_classes/configuration#transformers.PreTrainedConfig) を継承し、モデルの出力を制御するために使用できます。詳細については、[PreTrainedConfig](/docs/transformers/main/ja/main_classes/configuration#transformers.PreTrainedConfig) のドキュメントを参照してください。

## JambaModel[[transformers.JambaModel]]

#### transformers.JambaModel[[transformers.JambaModel]]

```python
transformers.JambaModel(config: JambaConfig)
```

[ソース](https://github.com/huggingface/transformers/blob/main/src/transformers/models/jamba/modeling_jamba.py#L835)

**パラメータ:**

config ([JambaConfig](/docs/transformers/main/ja/model_doc/jamba#transformers.JambaConfig)) : モデルのパラメータをすべて含むモデル設定クラスです。設定ファイルで初期化しても、モデルに関連する重みは読み込まれず、設定のみが読み込まれます。モデルの重みをロードするには、[from_pretrained()](/docs/transformers/main/ja/main_classes/model#transformers.PreTrainedModel.from_pretrained) メソッドを確認してください。

特定のヘッドを付加せずに、生の隠れ状態を出力する素のJamba Model。

このモデルは [PreTrainedModel](/docs/transformers/main/ja/main_classes/model#transformers.PreTrainedModel) を継承しています。ライブラリがすべてのモデルに対して実装している汎用的なメソッド（ダウンロードや保存、入力埋め込みのリサイズ、ヘッドのプルーニングなど）については、スーパークラスのドキュメントを参照してください。

このモデルは、PyTorchの[torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module)のサブクラスでもあります。通常のPyTorch Moduleとして使用し、一般的な使用方法や動作に関する詳細については、PyTorchのドキュメントを参照してください。

#### forward[[transformers.JambaModel.forward]]

```python
forward(input_ids: typing.Optional[torch.LongTensor] = None, attention_mask: typing.Optional[torch.Tensor] = None, position_ids: typing.Optional[torch.LongTensor] = None, past_key_values: transformers.cache_utils.Cache | None = None, inputs_embeds: typing.Optional[torch.FloatTensor] = None, use_cache: bool | None = None, **kwargs: *<class 'transformers.utils.generic.TransformersKwargs'>)
```

[ソース](https://github.com/huggingface/transformers/blob/main/src/transformers/models/jamba/modeling_jamba.py#L854)

**パラメータ:**

input_ids (`torch.LongTensor`、shape `(batch_size, sequence_length)`、*任意*) : 語彙における入力シーケンスのトークンインデックスです。デフォルトでは、パディングは無視されます。インデックスは [AutoTokenizer](/docs/transformers/main/ja/model_doc/auto#transformers.AutoTokenizer) を使用して取得できます。詳細については、[PreTrainedTokenizer.encode()](/docs/transformers/main/ja/internal/tokenization_utils#transformers.PreTrainedTokenizerBase.encode) および [PreTrainedTokenizer.__call__()](/docs/transformers/main/ja/internal/tokenization_utils#transformers.PreTrainedTokenizerBase.__call__) を参照してください。 [入力IDとは何ですか？](../glossary#input-ids)

attention_mask (`torch.Tensor`、shape `(batch_size, sequence_length)`、*任意*) : padding トークンのインデックスに対する attention を回避するためのマスクです。マスクの値は `[0, 1]` から選択されます:  - **マスクされない**トークンは 1、- **マスクされる**トークンは 0。  [attention マスクとは?](../glossary#attention-mask)

position_ids (`torch.LongTensor`、shape `(batch_size, sequence_length)`、*任意*) : 各入力シーケンスのトークンが位置埋め込み内で占める位置のインデックスです。`[0, config.n_positions - 1]` の範囲から選択されます。 [位置 ID とは何か？](../glossary#position-ids)

past_key_values (`~cache_utils.Cache`, *任意*) : 逐次デコードを高速化するために使用できる、事前計算された hidden-states（self-attention ブロックおよび cross-attention ブロック内の key と value）です。通常、`use_cache=True` または `config.use_cache=True` の場合、モデルがデコードの以前の段階で返した `past_key_values` で構成されます。入力として許可されるのは [Cache](/docs/transformers/main/ja/internal/generation_utils#transformers.Cache) インスタンスのみです。[kv cache ガイド](https://huggingface.co/docs/transformers/en/kv_cache)を参照してください。`past_key_values` が渡されない場合、デフォルトで [DynamicCache](/docs/transformers/main/ja/internal/generation_utils#transformers.DynamicCache) が初期化されます。モデルは、入力として与えられたのと同じキャッシュ形式を出力します。`past_key_values` を使用する場合は、すべての `input_ids`（形状 `(batch_size, sequence_length)`）の代わりに、このモデルに過去の key value 状態が与えられていない未処理の `input_ids`（形状 `(batch_size, unprocessed_length)`）のみを入力する必要があります。

inputs_embeds (`torch.FloatTensor`、shape `(batch_size, sequence_length, hidden_size)`、*任意*) : `input_ids` を渡す代わりに、埋め込み表現を直接渡すこともできます。これは、モデル内部の埋め込み参照行列よりも、`input_ids` のインデックスを対応するベクトルに変換する方法をより細かく制御したい場合に有用です。

use_cache (`bool`, *任意*) : `True`に設定すると、`past_key_values`のキーバリュー状態が返され、デコーディングの高速化に利用できます（`past_key_values`を参照）。

**戻り値:** `MoeModelOutputWithPast` または `tuple(torch.FloatTensor)`

`MoeModelOutputWithPast`、または`torch.FloatTensor`のタプル（`return_dict=False`が渡された場合、または`config.return_dict=False`の場合）。設定（[JambaConfig](/docs/transformers/main/ja/model_doc/jamba#transformers.JambaConfig)）や入力に応じて、さまざまな要素で構成されます。

[JambaModel](/docs/transformers/main/ja/model_doc/jamba#transformers.JambaModel)のforwardメソッドは、`__call__`特殊メソッドをオーバーライドします。

この関数内でフォワードパスの手順を定義する必要がありますが、その後、この関数ではなく`Module`インスタンスを呼び出すべきです。前者は前処理と後処理の実行を適切に処理するのに対し、後者はそれらを無視してしまいます。

- **last_hidden_state** (`torch.FloatTensor`、shape `(batch_size, sequence_length, hidden_size)`) -- モデルの最終層の出力における隠れ状態のシーケンス。
- **past_key_values** (`Cache`、*任意*、`use_cache=True` が渡された場合または `config.use_cache=True` の場合に返されます) -- [Cache](/docs/transformers/main/ja/internal/generation_utils#transformers.Cache) のインスタンスです。詳細については、[kv cache ガイド](https://huggingface.co/docs/transformers/en/kv_cache)を参照してください。

  自己アテンションブロック（および `config.is_encoder_decoder=True` の場合、クロスカレンションブロック）の事前計算済みヒドゥンステート（キーとバリュー）を保持します。これらは逐次デコーディングを高速化するために使用できます（`past_key_values` 入力参照）。
- **hidden_states** (`tuple(torch.FloatTensor)`、*optional*、`output_hidden_states=True` が渡された場合または `config.output_hidden_states=True` の場合に返されます) -- `torch.FloatTensor` のタプル（埋め込み層を持つモデルの場合、埋め込みの出力用が 1 つ、各層の出力用が 1 つずつ）で、形状は `(batch_size, sequence_length, hidden_size)` です。

  各層の出力におけるモデルのHidden-states、および任意の初期埋め込み出力。
- **attentions** (`tuple(torch.FloatTensor)`、*optional*、`output_attentions=True` が渡された場合または `config.output_attentions=True` の場合に返されます) -- `torch.FloatTensor` のタプル（各レイヤーごとに1つ）で、形状は `(batch_size, num_heads, sequence_length, sequence_length)` です。

  アテンションのsoftmax後のアテンション重み。自己アテンションヘッド内で加重平均を計算するために使用されます。
- **router_logits** (`tuple(torch.FloatTensor)`、*optional*、`output_router_probs=True` と `config.add_router_probs=True` が渡された場合、または `config.output_router_probs=True` の場合に返されます) -- `torch.FloatTensor` のタプル（レイヤーごとに1つ）、形状は `(batch_size, sequence_length, num_experts)`。

  MoE ルーターによって計算される生（raw）のルーター logits（softmax 後）であり、Mixture of Experts モデルの補助損失を計算するためにこれらの用語が使用されます。

## JambaForCausalLM[[transformers.JambaForCausalLM]]

#### transformers.JambaForCausalLM[[transformers.JambaForCausalLM]]

```python
transformers.JambaForCausalLM(config: JambaConfig)
```

[ソース](https://github.com/huggingface/transformers/blob/main/src/transformers/models/jamba/modeling_jamba.py#L984)

**パラメータ:**

config ([JambaConfig](/docs/transformers/main/ja/model_doc/jamba#transformers.JambaConfig)) : モデルのパラメータをすべて含むモデル設定クラスです。設定ファイルで初期化しても、モデルに関連する重みは読み込まれず、設定のみが読み込まれます。モデルの重みをロードするには、[from_pretrained()](/docs/transformers/main/ja/main_classes/model#transformers.PreTrainedModel.from_pretrained) メソッドを確認してください。

因果言語モデリング用のJambaモデル。

このモデルは [PreTrainedModel](/docs/transformers/main/ja/main_classes/model#transformers.PreTrainedModel) を継承しています。ライブラリがすべてのモデルに対して実装している汎用的なメソッド（ダウンロードや保存、入力埋め込みのリサイズ、ヘッドのプルーニングなど）については、スーパークラスのドキュメントを参照してください。

このモデルは、PyTorchの[torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module)のサブクラスでもあります。通常のPyTorch Moduleとして使用し、一般的な使用方法や動作に関する詳細については、PyTorchのドキュメントを参照してください。

#### forward[[transformers.JambaForCausalLM.forward]]

```python
forward(input_ids: typing.Optional[torch.LongTensor] = None, attention_mask: typing.Optional[torch.Tensor] = None, position_ids: typing.Optional[torch.LongTensor] = None, past_key_values: transformers.cache_utils.Cache | None = None, inputs_embeds: typing.Optional[torch.FloatTensor] = None, labels: typing.Optional[torch.LongTensor] = None, use_cache: bool | None = None, output_router_logits: bool | None = None, logits_to_keep: typing.Union[int, torch.Tensor] = 0, **kwargs: *<class 'transformers.utils.generic.TransformersKwargs'>)
```

[ソース](https://github.com/huggingface/transformers/blob/main/src/transformers/models/jamba/modeling_jamba.py#L1002)

**パラメータ:**

input_ids (`torch.LongTensor`、shape `(batch_size, sequence_length)`、*任意*) : 語彙における入力シーケンスのトークンインデックスです。デフォルトでは、パディングは無視されます。インデックスは [AutoTokenizer](/docs/transformers/main/ja/model_doc/auto#transformers.AutoTokenizer) を使用して取得できます。詳細については、[PreTrainedTokenizer.encode()](/docs/transformers/main/ja/internal/tokenization_utils#transformers.PreTrainedTokenizerBase.encode) および [PreTrainedTokenizer.__call__()](/docs/transformers/main/ja/internal/tokenization_utils#transformers.PreTrainedTokenizerBase.__call__) を参照してください。 [入力IDとは何ですか？](../glossary#input-ids)

attention_mask (`torch.Tensor`、shape `(batch_size, sequence_length)`、*任意*) : padding トークンのインデックスに対する attention を回避するためのマスクです。マスクの値は `[0, 1]` から選択されます:  - **マスクされない**トークンは 1、- **マスクされる**トークンは 0。  [attention マスクとは?](../glossary#attention-mask)

position_ids (`torch.LongTensor`、shape `(batch_size, sequence_length)`、*任意*) : 各入力シーケンスのトークンが位置埋め込み内で占める位置のインデックスです。`[0, config.n_positions - 1]` の範囲から選択されます。 [位置 ID とは何か？](../glossary#position-ids)

past_key_values (`~cache_utils.Cache`, *任意*) : 逐次デコードを高速化するために使用できる、事前計算された hidden-states（self-attention ブロックおよび cross-attention ブロック内の key と value）です。通常、`use_cache=True` または `config.use_cache=True` の場合、モデルがデコードの以前の段階で返した `past_key_values` で構成されます。入力として許可されるのは [Cache](/docs/transformers/main/ja/internal/generation_utils#transformers.Cache) インスタンスのみです。[kv cache ガイド](https://huggingface.co/docs/transformers/en/kv_cache)を参照してください。`past_key_values` が渡されない場合、デフォルトで [DynamicCache](/docs/transformers/main/ja/internal/generation_utils#transformers.DynamicCache) が初期化されます。モデルは、入力として与えられたのと同じキャッシュ形式を出力します。`past_key_values` を使用する場合は、すべての `input_ids`（形状 `(batch_size, sequence_length)`）の代わりに、このモデルに過去の key value 状態が与えられていない未処理の `input_ids`（形状 `(batch_size, unprocessed_length)`）のみを入力する必要があります。

inputs_embeds (`torch.FloatTensor`、shape `(batch_size, sequence_length, hidden_size)`、*任意*) : `input_ids` を渡す代わりに、埋め込み表現を直接渡すこともできます。これは、モデル内部の埋め込み参照行列よりも、`input_ids` のインデックスを対応するベクトルに変換する方法をより細かく制御したい場合に有用です。

labels (`torch.LongTensor`、形状 `(batch_size, sequence_length)`、*任意*) : マスク言語モデリング損失を計算するためのラベル。インデックスは `[0, ..., config.vocab_size]` の範囲内であるか、-100 である必要があります（`input_ids` のドキュメント文字列を参照）。インデックスが `-100` に設定されたトークンは無視（マスク）され、損失は `[0, ..., config.vocab_size]` の範囲内のラベルを持つトークンに対してのみ計算されます。

use_cache (`bool`, *任意*) : `True`に設定すると、`past_key_values`のキーバリュー状態が返され、デコーディングの高速化に利用できます（`past_key_values`を参照）。

output_router_logits (`bool`, *任意*) : すべてのルーターの logits を返すかどうか。ルーター損失の計算に有用であり、推論時には返すべきではありません。

logits_to_keep (`Union[int, torch.Tensor]`、*任意*、デフォルトは `0`) : `int` の場合、最後の `logits_to_keep` トークンについて logits を計算します。`0` の場合、すべての `input_ids` について logits を計算します（特殊ケース）。生成には最後のトークンの logits だけで十分であり、そのトークンのみについて計算することでメモリを節約できます。これは、長いシーケンスや大規模な語彙サイズの場合に特に顕著な効果があります。`torch.Tensor` の場合、シーケンス長次元で保持するインデックスに対応する 1 次元である必要があります。これは、パックされたテンソル形式（バッチとシーケンス長を単一の次元にまとめる形式）を使用する場合に有用です。

**戻り値:** `MoeCausalLMOutputWithPast` または `tuple(torch.FloatTensor)`

`MoeCausalLMOutputWithPast`、または`torch.FloatTensor`のタプル（`return_dict=False`が渡された場合、または`config.return_dict=False`の場合）。設定（[JambaConfig](/docs/transformers/main/ja/model_doc/jamba#transformers.JambaConfig)）や入力に応じて、さまざまな要素で構成されます。

[JambaForCausalLM](/docs/transformers/main/ja/model_doc/jamba#transformers.JambaForCausalLM)のforwardメソッドは、`__call__`特殊メソッドをオーバーライドします。

この関数内でフォワードパスの手順を定義する必要がありますが、その後、この関数ではなく`Module`インスタンスを呼び出すべきです。前者は前処理と後処理の実行を適切に処理するのに対し、後者はそれらを無視してしまいます。

- **loss** (`torch.FloatTensor`、shape `(1,)`、*optional*、`labels` が指定された場合に返されます) -- 言語モデリングの損失(次のトークンの予測用)。
- **logits** (`torch.FloatTensor`、shape `(batch_size, sequence_length, config.vocab_size)`) -- 言語モデリングヘッドの予測スコア（SoftMax 適用前の各語彙トークンのスコア）。
- **aux_loss** (`torch.FloatTensor`、*optional*、`output_router_logits=True` が渡された場合、または `config.output_router_logits=True` の場合、かつモデルがロードバランシング損失でルーターを学習する場合に返されます) -- スパースモジュール用のロードバランシング補助損失。ルーターバイアスでエキスパートをバランスさせるモデル（DeepSeek-V3 とそこから派生したアーキテクチャ）は、`router_logits` を返しながらも、ここでは `None` を返します。
- **router_logits** (`tuple(torch.FloatTensor)`、*optional*、`output_router_logits=True` が渡された場合または `config.output_router_logits=True` の場合に返されます) -- `torch.FloatTensor` のタプル（スパース層ごとに1つ）、形状は `(batch_size * sequence_length, num_experts)`。

  MoE ルーターによって計算される生のルーター logits です。ロードバランシング損失の計算や、トークンがどのようにエキスパートにディスパッチされるかの確認に使用できます。
- **past_key_values** (`Cache`、*任意*、`use_cache=True` が渡された場合または `config.use_cache=True` の場合に返されます) -- [Cache](/docs/transformers/main/ja/internal/generation_utils#transformers.Cache) のインスタンスです。詳細については、[kv cache ガイド](https://huggingface.co/docs/transformers/en/kv_cache)を参照してください。

  逐次デコーディングを高速化するために使用できる（`past_key_values` 入力参照）、事前計算済みの隠れ状態（自己注意ブロックのキーとバリュー）が含まれます。
- **hidden_states** (`tuple(torch.FloatTensor)`、*optional*、`output_hidden_states=True` が渡された場合または `config.output_hidden_states=True` の場合に返されます) -- `torch.FloatTensor` のタプル（埋め込み層を持つモデルの場合、埋め込みの出力用が 1 つ、各層の出力用が 1 つずつ）で、形状は `(batch_size, sequence_length, hidden_size)` です。

  各層の出力におけるモデルのHidden-states、および任意の初期埋め込み出力。
- **attentions** (`tuple(torch.FloatTensor)`、*optional*、`output_attentions=True` が渡された場合または `config.output_attentions=True` の場合に返されます) -- `torch.FloatTensor` のタプル（各レイヤーごとに1つ）で、形状は `(batch_size, num_heads, sequence_length, sequence_length)` です。

  アテンションのsoftmax後のアテンション重み。自己アテンションヘッド内で加重平均を計算するために使用されます。

例:

```python
>>> from transformers import AutoTokenizer, JambaForCausalLM

>>> model = JambaForCausalLM.from_pretrained("ai21labs/Jamba-v0.1")
>>> tokenizer = AutoTokenizer.from_pretrained("ai21labs/Jamba-v0.1")

>>> prompt = "Hey, are you conscious? Can you talk to me?"
>>> inputs = tokenizer(prompt, return_tensors="pt")

>>> # Generate
>>> generate_ids = model.generate(inputs.input_ids, max_length=30)
>>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
"Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
```

## JambaForSequenceClassification[[transformers.JambaForSequenceClassification]]

#### transformers.JambaForSequenceClassification[[transformers.JambaForSequenceClassification]]

```python
transformers.JambaForSequenceClassification(config)
```

[ソース](https://github.com/huggingface/transformers/blob/main/src/transformers/models/jamba/modeling_jamba.py#L1081)

#### forward[[transformers.JambaForSequenceClassification.forward]]

```python
forward(input_ids: torch.LongTensor | None = None, attention_mask: torch.Tensor | None = None, position_ids: torch.LongTensor | None = None, past_key_values: Cache | None = None, inputs_embeds: torch.FloatTensor | None = None, labels: torch.LongTensor | None = None, use_cache: bool | None = None, **kwargs: Unpack[TransformersKwargs])
```

[ソース](https://github.com/huggingface/transformers/blob/main/src/transformers/modeling_layers.py#L128)

**パラメータ:**

input_ids (`torch.LongTensor`、shape `(batch_size, sequence_length)`、*任意*) : 語彙における入力シーケンスのトークンインデックスです。デフォルトでは、パディングは無視されます。インデックスは [AutoTokenizer](/docs/transformers/main/ja/model_doc/auto#transformers.AutoTokenizer) を使用して取得できます。詳細については、[PreTrainedTokenizer.encode()](/docs/transformers/main/ja/internal/tokenization_utils#transformers.PreTrainedTokenizerBase.encode) および [PreTrainedTokenizer.__call__()](/docs/transformers/main/ja/internal/tokenization_utils#transformers.PreTrainedTokenizerBase.__call__) を参照してください。 [入力IDとは何ですか？](../glossary#input-ids)

attention_mask (`torch.Tensor`、shape `(batch_size, sequence_length)`、*任意*) : padding トークンのインデックスに対する attention を回避するためのマスクです。マスクの値は `[0, 1]` から選択されます:  - **マスクされない**トークンは 1、- **マスクされる**トークンは 0。  [attention マスクとは?](../glossary#attention-mask)

position_ids (`torch.LongTensor`、shape `(batch_size, sequence_length)`、*任意*) : 各入力シーケンスのトークンが位置埋め込み内で占める位置のインデックスです。`[0, config.n_positions - 1]` の範囲から選択されます。 [位置 ID とは何か？](../glossary#position-ids)

past_key_values (`Cache`, *任意*) : 逐次デコードを高速化するために使用できる、事前計算された隠れ状態（自己アテンションブロックおよびクロスアテンションブロック内のキーとバリュー）です。通常、`use_cache=True` または `config.use_cache=True` が設定されている場合、デコードの以前の段階でモデルが返した `past_key_values` で構成されます。入力として許可されるのは [Cache](/docs/transformers/main/ja/internal/generation_utils#transformers.Cache) インスタンスのみです。[kv cache ガイド](https://huggingface.co/docs/transformers/en/kv_cache) を参照してください。`past_key_values` が渡されない場合、デフォルトで [DynamicCache](/docs/transformers/main/ja/internal/generation_utils#transformers.DynamicCache) が初期化されます。モデルは、入力として与えられたのと同じキャッシュ形式を出力します。`past_key_values` を使用する場合は、すべての `input_ids`（形状 `(batch_size, sequence_length)`）の代わりに、このモデルに過去のキーバリュー状態が与えられていない未処理の `input_ids`（形状 `(batch_size, unprocessed_length)`）のみを入力する必要があります。

inputs_embeds (`torch.FloatTensor`、shape `(batch_size, sequence_length, hidden_size)`、*任意*) : `input_ids` を渡す代わりに、埋め込み表現を直接渡すこともできます。これは、モデル内部の埋め込み参照行列よりも、`input_ids` のインデックスを対応するベクトルに変換する方法をより細かく制御したい場合に有用です。

labels (`torch.LongTensor`、形状 `(batch_size,)`、*任意*) : シーケンス分類/回帰の損失を計算するためのラベル。インデックスは `[0, ..., config.num_labels - 1]` の範囲である必要があります。`config.num_labels == 1` の場合は回帰損失（平均二乗誤差）が計算され、`config.num_labels > 1` の場合は分類損失（クロスエントロピー）が計算されます。

use_cache (`bool`, *任意*) : `True`に設定すると、`past_key_values`のキーバリュー状態が返され、デコーディングの高速化に利用できます（`past_key_values`を参照）。

**戻り値:** `SequenceClassifierOutputWithPast`

`GenericForSequenceClassification`のforwardメソッドは、`__call__`特殊メソッドをオーバーライドします。

この関数内でフォワードパスの手順を定義する必要がありますが、その後、この関数ではなく`Module`インスタンスを呼び出すべきです。前者は前処理と後処理の実行を適切に処理するのに対し、後者はそれらを無視してしまいます。