Transformers 文件
Jamba
並獲得增強的文件體驗
開始使用
Jamba
Jamba 是一款混合了 Transformer 和 Mamba 的混合專家(MoE)語言模型,其總引數量從 520 億到 3980 億不等。該模型旨在結合這兩個模型家族的優點:Transformer 模型的效能,以及像 Mamba 這類狀態空間模型(SSM)的效率和更長的上下文(256K 詞元)。
Jamba 的架構採用了一種塊與層的方法,使其能夠成功地將 Transformer 和 Mamba 架構整合在一起。每個 Jamba 塊包含一個注意力層或一個 Mamba 層,後面跟著一個多層感知機(MLP),從而使得每八個總層中就有一個是 Transformer 層。同時,混合了 MoE 層以增加模型容量。
你可以在 AI21 組織下找到所有原始的 Jamba 模型檢查點。
點選右側邊欄中的 Jamba 模型,檢視更多關於如何將 Jamba 應用於不同語言任務的示例。
下面的示例演示瞭如何使用 Pipeline、AutoModel 以及從命令列生成文字。
# install optimized Mamba implementations
# !pip install mamba-ssm causal-conv1d>=1.2.0
import torch
from transformers import pipeline
pipeline = pipeline(
task="text-generation",
model="ai21labs/AI21-Jamba-Mini-1.6",
torch_dtype=torch.float16,
device=0
)
pipeline("Plants create energy through a process known as")
量化透過以較低精度表示權重來減少大型模型的記憶體負擔。有關更多可用量化後端,請參閱量化概述。
下面的示例使用 bitsandbytes 將權重僅量化為 8 位。
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, 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",
torch_dtype=torch.bfloat16,
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 核心的情況下使用 Mamba,因為這會導致延遲顯著增加。如果你仍想在沒有核心的情況下使用 Mamba,請在 from_pretrained() 中設定
use_mamba_kernels=False
。import torch from transformers import AutoModelForCausalLM model = AutoModelForCausalLM.from_pretrained("ai21labs/AI21-Jamba-1.5-Large", use_mamba_kernels=False)
JambaConfig
class transformers.JambaConfig
< 原始碼 >( vocab_size = 65536 tie_word_embeddings = False hidden_size = 4096 intermediate_size = 14336 num_hidden_layers = 32 num_attention_heads = 32 num_key_value_heads = 8 hidden_act = 'silu' initializer_range = 0.02 rms_norm_eps = 1e-06 use_cache = True num_logits_to_keep = 1 output_router_logits = False router_aux_loss_coef = 0.001 pad_token_id = 0 bos_token_id = 1 eos_token_id = 2 sliding_window = None max_position_embeddings = 262144 attention_dropout = 0.0 num_experts_per_tok = 2 num_experts = 16 expert_layer_period = 2 expert_layer_offset = 1 attn_layer_period = 8 attn_layer_offset = 4 use_mamba_kernels = True mamba_d_state = 16 mamba_d_conv = 4 mamba_expand = 2 mamba_dt_rank = 'auto' mamba_conv_bias = True mamba_proj_bias = False **kwargs )
引數
- vocab_size (
int
, 可選, 預設為 65536) — Jamba 模型的詞彙表大小。定義了在呼叫 JambaModel 時傳入的inputs_ids
可以表示的不同詞元的數量。 - tie_word_embeddings (
bool
, 可選, 預設為False
) — 是否應將模型的輸入和輸出詞嵌入繫結。請注意,這僅在模型具有輸出詞嵌入層時才相關。 - hidden_size (
int
, 可選, 預設為 4096) — 隱藏表示的維度。 - intermediate_size (
int
, 可選, 預設為 14336) — MLP 表示的維度。 - num_hidden_layers (
int
, 可選, 預設為 32) — Transformer 編碼器中的隱藏層數量。 - num_attention_heads (
int
, 可選, 預設為 32) — Transformer 編碼器中每個注意力層的注意力頭數量。 - num_key_value_heads (
int
, 可選, 預設為 8) — 這是實現分組查詢注意力(Grouped Query Attention)時應使用的鍵值頭(key_value heads)的數量。如果num_key_value_heads=num_attention_heads
,模型將使用多頭注意力(MHA);如果num_key_value_heads=1
,模型將使用多查詢注意力(MQA);否則,使用 GQA。當將一個多頭檢查點轉換為 GQA 檢查點時,每個組的鍵和值頭應透過對該組內所有原始頭進行平均池化來構建。更多細節,請檢視這篇論文。如果未指定,將預設為8
。 - hidden_act (
str
或function
, 可選, 預設為"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
時相關。 - num_logits_to_keep (
int
或None
, 可選, 預設為 1) — 在生成過程中需要計算的提示(prompt) logits 的數量。如果為None
,則計算所有 logits。如果為整數值,則只計算最後num_logits_to_keep
個 logits。預設為 1,因為生成時只需要最後一個提示詞元的 logits。對於長序列,整個序列的 logits 可能會佔用大量記憶體,因此設定num_logits_to_keep=1
將顯著減少記憶體佔用。 - output_router_logits (
bool
, 可選, 預設為False
) — 是否應由模型返回路由器(router)的 logits。啟用此項也將允許模型輸出輔助損失。更多詳情請參閱此處。 - router_aux_loss_coef (
float
, 可選, 預設為 0.001) — 總損失中的輔助損失因子。 - pad_token_id (
int
, 可選, 預設為 0) — 填充詞元的 ID。 - bos_token_id (
int
, 可選, 預設為 1) — “序列開始”詞元的 ID。 - eos_token_id (
int
, 可選, 預設為 2) — “序列結束”詞元的 ID。 - sliding_window (
int
, 可選) — 滑動視窗注意力的視窗大小。如果未指定,將預設為None
。 - max_position_embeddings (
int
, 可選, 預設為 262144) — 此值沒有實際效果。該模型預期使用的最大序列長度。它可以用於更長的序列,但效能可能會下降。 - attention_dropout (
float
, 可選, 預設為 0.0) — 注意力機率的 dropout 比率。 - num_experts_per_tok (
int
, 可選, 預設為 2) — 每個詞元要路由到的專家數量,也可以解釋為top-p
路由引數。 - num_experts (
int
, 可選, 預設為 16) — 每個稀疏 MLP 層的專家數量。 - expert_layer_period (
int
, 可選, 預設為 2) — 每隔這麼多層,就會有一個專家層。 - expert_layer_offset (
int
, 可選, 預設為 1) — 包含專家 MLP 層的第一個層的索引。 - attn_layer_period (
int
, 可選, 預設為 8) — 每隔這麼多層,就會有一個普通的注意力層。 - attn_layer_offset (
int
, 可選, 預設為 4) — 包含普通注意力 MLP 層的第一個層的索引。 - use_mamba_kernels (
bool
, 可選, 預設為True
) — 標誌,指示是否使用快速 Mamba 核心。這些核心僅在安裝了mamba-ssm
和causal-conv1d
,並且 Mamba 模組在 CUDA 裝置上執行時才可用。如果為True
但核心不可用,則會引發 ValueError。 - mamba_d_state (
int
, 可選, 預設為 16) — Mamba 狀態空間潛變數的維度。 - mamba_d_conv (
int
, optional, 預設為 4) — Mamba 卷積核的大小 - mamba_expand (
int
, optional, 預設為 2) — 用於確定 Mamba 中間大小的擴充套件因子(相對於 hidden_size) - mamba_dt_rank (
Union[int,str]
, optional, 預設為"auto"
) — Mamba 離散化投影矩陣的秩。"auto"
表示它將預設為math.ceil(self.hidden_size / 16)
- mamba_conv_bias (
bool
, optional, 預設為True
) — 指示是否在 Mamba 混合器塊的卷積層中使用偏置的標誌。 - mamba_proj_bias (
bool
, optional, 預設為False
) — 指示是否在 Mamba 混合器塊的輸入和輸出投影([“in_proj”, “out_proj”])中使用偏置的標誌
這是一個配置類,用於儲存 JambaModel 的配置。它用於根據指定的引數例項化 Jamba 模型,定義模型架構。使用預設值例項化配置將產生與 Jamba-v0.1 模型類似的配置。
配置物件繼承自 PretrainedConfig,可用於控制模型輸出。請閱讀 PretrainedConfig 的文件以獲取更多資訊。
JambaModel
class transformers.JambaModel
< 原始檔 >( config: JambaConfig )
引數
- config (JambaConfig) — 包含模型所有引數的模型配置類。使用配置檔案初始化不會載入與模型相關的權重,只會載入配置。請檢視 from_pretrained() 方法來載入模型權重。
基礎的 Jamba 模型,輸出原始的隱藏狀態,頂部沒有任何特定的頭部。
該模型繼承自 PreTrainedModel。請檢視超類文件,瞭解該庫為其所有模型實現的通用方法(例如下載或儲存、調整輸入嵌入大小、修剪頭部等)。
該模型也是一個 PyTorch torch.nn.Module 的子類。可以像常規的 PyTorch 模組一樣使用它,並參考 PyTorch 文件瞭解所有與通用用法和行為相關的事項。
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: typing.Optional[transformers.models.jamba.modeling_jamba.HybridMambaAttentionDynamicCache] = None inputs_embeds: typing.Optional[torch.FloatTensor] = None use_cache: typing.Optional[bool] = None output_attentions: typing.Optional[bool] = None output_hidden_states: typing.Optional[bool] = None output_router_logits: typing.Optional[bool] = None cache_position: typing.Optional[torch.LongTensor] = None ) → transformers.modeling_outputs.MoeModelOutputWithPast
或 tuple(torch.FloatTensor)
引數
- input_ids (
torch.LongTensor
,形狀為(batch_size, sequence_length)
, optional) — 詞彙表中輸入序列標記的索引。預設情況下,填充將被忽略。索引可以使用 AutoTokenizer 獲取。詳情請參閱 PreTrainedTokenizer.encode() 和 PreTrainedTokenizer.call()。
- attention_mask (
torch.Tensor
,形狀為(batch_size, sequence_length)
, optional) — 用於避免對填充標記索引執行注意力的掩碼。掩碼值在[0, 1]
中選擇:- 1 表示標記未被掩碼,
- 0 表示標記被掩碼。
- position_ids (
torch.LongTensor
,形狀為(batch_size, sequence_length)
, optional) — 每個輸入序列標記在位置嵌入中的位置索引。在[0, config.n_positions - 1]
範圍內選擇。 - past_key_values (
~models.jamba.modeling_jamba.HybridMambaAttentionDynamicCache
, optional) — 預計算的隱藏狀態(自注意力塊和交叉注意力塊中的鍵和值),可用於加速順序解碼。這通常是在解碼的先前階段,當use_cache=True
或config.use_cache=True
時,由模型返回的past_key_values
。允許兩種格式:
- 一個 Cache 例項,請參閱我們的 kv 快取指南;
- 長度為
config.n_layers
的tuple(torch.FloatTensor)
元組,每個元組包含 2 個形狀為(batch_size, num_heads, sequence_length, embed_size_per_head)
的張量。這也稱為舊版快取格式。
模型將輸出與輸入相同的快取格式。如果沒有傳遞
past_key_values
,將返回舊版快取格式。如果使用
past_key_values
,使用者可以選擇只輸入最後一個input_ids
(那些沒有給出過去鍵值狀態的標記),形狀為(batch_size, 1)
,而不是所有形狀為(batch_size, sequence_length)
的input_ids
。 - inputs_embeds (
torch.FloatTensor
,形狀為(batch_size, sequence_length, hidden_size)
, optional) — 可選地,您可以選擇直接傳遞嵌入表示,而不是傳遞input_ids
。如果您希望比模型的內部嵌入查詢矩陣更好地控制如何將input_ids
索引轉換為關聯向量,這會很有用。 - use_cache (
bool
, optional) — 如果設定為True
,則返回past_key_values
鍵值狀態,可用於加速解碼(參見past_key_values
)。 - output_attentions (
bool
, optional) — 是否返回所有注意力層的注意力張量。有關更多詳細資訊,請參閱返回張量下的attentions
。 - output_hidden_states (
bool
, optional) — 是否返回所有層的隱藏狀態。有關更多詳細資訊,請參閱返回張量下的hidden_states
。 - output_router_logits (
bool
, optional) — 是否返回所有路由器的 logits。它們對於計算路由器損失很有用,在推理期間不應返回。 - cache_position (
torch.LongTensor
,形狀為(sequence_length)
, optional) — 描述輸入序列標記在序列中位置的索引。與position_ids
不同,此張量不受填充影響。它用於在正確的位置更新快取並推斷完整的序列長度。
返回
transformers.modeling_outputs.MoeModelOutputWithPast
或 tuple(torch.FloatTensor)
一個 transformers.modeling_outputs.MoeModelOutputWithPast
或一個 torch.FloatTensor
的元組(如果傳遞了 return_dict=False
或當 config.return_dict=False
時),根據配置(JambaConfig)和輸入,包含各種元素。
-
last_hidden_state (
torch.FloatTensor
, 形狀為(batch_size, sequence_length, hidden_size)
) — 模型最後一層輸出的隱藏狀態序列。 -
past_key_values (
Cache
, optional, 當傳遞use_cache=True
或config.use_cache=True
時返回) — 這是一個 Cache 例項。更多詳情,請參閱我們的 kv 快取指南。包含預計算的隱藏狀態(自注意力塊中的鍵和值,如果
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
的元組(一個用於嵌入層的輸出,如果模型有嵌入層,+ 一個用於每層的輸出),形狀為(batch_size, sequence_length, hidden_size)
。模型在每個層輸出的隱藏狀態以及可選的初始嵌入輸出。
-
attentions (
tuple(torch.FloatTensor)
, optional, 當傳遞output_attentions=True
或config.output_attentions=True
時返回) —torch.FloatTensor
的元組(每層一個),形狀為(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
的元組(每層一個),形狀為(batch_size, sequence_length, num_experts)
。由 MoE 路由器計算的原始路由器對數(softmax 後),這些術語用於計算專家混合模型的輔助損失。
JambaModel 的前向方法,覆蓋了 __call__
特殊方法。
儘管前向傳遞的邏輯需要在此函式內定義,但之後應呼叫 Module
例項而不是此函式,因為前者會處理執行前處理和後處理步驟,而後者會靜默地忽略它們。
JambaForCausalLM
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: typing.Optional[transformers.models.jamba.modeling_jamba.HybridMambaAttentionDynamicCache] = None inputs_embeds: typing.Optional[torch.FloatTensor] = None labels: typing.Optional[torch.LongTensor] = None use_cache: typing.Optional[bool] = None output_attentions: typing.Optional[bool] = None output_hidden_states: typing.Optional[bool] = None output_router_logits: typing.Optional[bool] = None cache_position: typing.Optional[torch.LongTensor] = None logits_to_keep: typing.Union[int, torch.Tensor] = 0 **loss_kwargs ) → transformers.modeling_outputs.MoeCausalLMOutputWithPast
或 tuple(torch.FloatTensor)
引數
- input_ids (
torch.LongTensor
,形狀為(batch_size, sequence_length)
, optional) — 詞彙表中輸入序列標記的索引。預設情況下,填充將被忽略。索引可以使用 AutoTokenizer 獲取。詳情請參閱 PreTrainedTokenizer.encode() 和 PreTrainedTokenizer.call()。
- attention_mask (
torch.Tensor
,形狀為(batch_size, sequence_length)
, optional) — 用於避免對填充標記索引執行注意力的掩碼。掩碼值在[0, 1]
中選擇:- 1 表示標記未被掩碼,
- 0 表示標記被掩碼。
- position_ids (
torch.LongTensor
,形狀為(batch_size, sequence_length)
, optional) — 每個輸入序列標記在位置嵌入中的位置索引。在[0, config.n_positions - 1]
範圍內選擇。 - past_key_values (
~models.jamba.modeling_jamba.HybridMambaAttentionDynamicCache
, optional) — 預計算的隱藏狀態(自注意力塊和交叉注意力塊中的鍵和值),可用於加速順序解碼。這通常是在解碼的先前階段,當use_cache=True
或config.use_cache=True
時,由模型返回的past_key_values
。允許兩種格式:
- 一個 Cache 例項,請參閱我們的 kv 快取指南;
- 長度為
config.n_layers
的tuple(torch.FloatTensor)
元組,每個元組包含 2 個形狀為(batch_size, num_heads, sequence_length, embed_size_per_head)
的張量。這也稱為舊版快取格式。
模型將輸出與輸入相同的快取格式。如果沒有傳遞
past_key_values
,將返回舊版快取格式。如果使用
past_key_values
,使用者可以選擇只輸入最後一個input_ids
(那些沒有給出過去鍵值狀態的標記),形狀為(batch_size, 1)
,而不是所有形狀為(batch_size, sequence_length)
的input_ids
。 - inputs_embeds (
torch.FloatTensor
,形狀為(batch_size, sequence_length, hidden_size)
, optional) — 可選地,您可以選擇直接傳遞嵌入表示,而不是傳遞input_ids
。如果您希望比模型的內部嵌入查詢矩陣更好地控制如何將input_ids
索引轉換為關聯向量,這會很有用。 - labels (
torch.LongTensor
,形狀為(batch_size, sequence_length)
, optional) — 用於計算掩碼語言建模損失的標籤。索引應在[0, ..., config.vocab_size]
或 -100 之間(參見input_ids
文件字串)。索引設定為-100
的標記將被忽略(掩碼),損失僅針對標籤在[0, ..., config.vocab_size]
中的標記計算。 - use_cache (
bool
, optional) — 如果設定為True
,則返回past_key_values
鍵值狀態,可用於加速解碼(參見past_key_values
)。 - output_attentions (
bool
, optional) — 是否返回所有注意力層的注意力張量。有關更多詳細資訊,請參閱返回張量下的attentions
。 - output_hidden_states (
bool
, optional) — 是否返回所有層的隱藏狀態。有關更多詳細資訊,請參閱返回張量下的hidden_states
。 - output_router_logits (
bool
, optional) — 是否返回所有路由器的 logits。它們對於計算路由器損失很有用,在推理期間不應返回。 - cache_position (
torch.LongTensor
,形狀為(sequence_length)
, optional) — 描述輸入序列標記在序列中位置的索引。與position_ids
不同,此張量不受填充影響。它用於在正確的位置更新快取並推斷完整的序列長度。 - logits_to_keep (
Union[int, torch.Tensor]
, 預設為0
) — 如果是int
,則計算最後logits_to_keep
個標記的 logits。如果為0
,則計算所有input_ids
的 logits(特殊情況)。生成時只需要最後一個標記的 logits,僅為該標記計算它們可以節省記憶體,這對於長序列或大詞彙表大小非常重要。如果是一個torch.Tensor
,則必須是一維的,對應於要在序列長度維度中保留的索引。這在使用打包張量格式(批處理和序列長度的單個維度)時很有用。
返回
transformers.modeling_outputs.MoeCausalLMOutputWithPast
或 tuple(torch.FloatTensor)
一個 transformers.modeling_outputs.MoeCausalLMOutputWithPast
或一個 torch.FloatTensor
的元組(如果傳遞了 return_dict=False
或當 config.return_dict=False
時),根據配置(JambaConfig)和輸入,包含各種元素。
-
loss (
torch.FloatTensor
形狀為(1,)
,可選,當提供labels
時返回) — 語言建模損失(用於下一個 token 預測)。 -
logits (形狀為
(batch_size, sequence_length, config.vocab_size)
的torch.FloatTensor
) — 語言建模頭部的預測分數(SoftMax 之前的每個詞彙標記的分數)。 -
aux_loss (
torch.FloatTensor
,可選,當提供labels
時返回) — 稀疏模組的輔助損失。 -
router_logits (
tuple(torch.FloatTensor)
, optional, 當傳遞output_router_probs=True
和config.add_router_probs=True
或config.output_router_probs=True
時返回) —torch.FloatTensor
的元組(每層一個),形狀為(batch_size, sequence_length, num_experts)
。由 MoE 路由器計算的原始路由器對數(softmax 後),這些術語用於計算專家混合模型的輔助損失。
-
past_key_values (
Cache
, optional, 當傳遞use_cache=True
或config.use_cache=True
時返回) — 這是一個 Cache 例項。更多詳情,請參閱我們的 kv 快取指南。包含預計算的隱藏狀態(自注意力塊中的鍵和值),可用於(參見
past_key_values
輸入)加速順序解碼。 -
hidden_states (
tuple(torch.FloatTensor)
, optional, 當傳遞output_hidden_states=True
或config.output_hidden_states=True
時返回) —torch.FloatTensor
的元組(一個用於嵌入層的輸出,如果模型有嵌入層,+ 一個用於每層的輸出),形狀為(batch_size, sequence_length, hidden_size)
。模型在每個層輸出的隱藏狀態以及可選的初始嵌入輸出。
-
attentions (
tuple(torch.FloatTensor)
, optional, 當傳遞output_attentions=True
或config.output_attentions=True
時返回) —torch.FloatTensor
的元組(每層一個),形狀為(batch_size, num_heads, sequence_length, sequence_length)
。注意力 softmax 後的注意力權重,用於計算自注意力頭中的加權平均值。
The JambaForCausalLM forward method, overrides the __call__
special method.
儘管前向傳遞的邏輯需要在此函式內定義,但之後應呼叫 Module
例項而不是此函式,因為前者會處理執行前處理和後處理步驟,而後者會靜默地忽略它們。
示例
>>> 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
class transformers.JambaForSequenceClassification
< 原始檔 >( config )
引數
- config (JambaForSequenceClassification) — 包含模型所有引數的模型配置類。使用配置檔案初始化不會載入與模型相關的權重,只會載入配置。請檢視 from_pretrained() 方法來載入模型權重。
Jamba 模型頂部帶有一個序列分類頭(線性層)。
JambaForSequenceClassification 使用最後一個標記來進行分類,就像其他因果模型(例如 GPT-2)一樣。
由於它對最後一個標記進行分類,因此需要知道最後一個標記的位置。如果在配置中定義了 pad_token_id
,它會找到每行中不是填充標記的最後一個標記。如果沒有定義 pad_token_id
,它會簡單地取批次中每行的最後一個值。由於當傳遞 inputs_embeds
而不是 input_ids
時無法猜測填充標記,它會執行相同的操作(取批次中每行的最後一個值)。
該模型繼承自 PreTrainedModel。請檢視超類文件,瞭解該庫為其所有模型實現的通用方法(例如下載或儲存、調整輸入嵌入大小、修剪頭部等)。
該模型也是一個 PyTorch torch.nn.Module 的子類。可以像常規的 PyTorch 模組一樣使用它,並參考 PyTorch 文件瞭解所有與通用用法和行為相關的事項。
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: typing.Optional[transformers.cache_utils.Cache] = None inputs_embeds: typing.Optional[torch.FloatTensor] = None labels: typing.Optional[torch.LongTensor] = None use_cache: typing.Optional[bool] = None output_attentions: typing.Optional[bool] = None output_hidden_states: typing.Optional[bool] = None ) → transformers.modeling_outputs.SequenceClassifierOutputWithPast
或 tuple(torch.FloatTensor)
引數
- input_ids (
torch.LongTensor
,形狀為(batch_size, sequence_length)
,可選) — 詞彙表中輸入序列標記的索引。預設情況下,填充將被忽略。可以使用 AutoTokenizer 獲取索引。有關詳細資訊,請參閱 PreTrainedTokenizer.encode() 和 PreTrainedTokenizer.call()。
- attention_mask (
torch.Tensor
,形狀為(batch_size, sequence_length)
,可選) — 用於避免對填充標記索引執行注意力操作的掩碼。掩碼值在[0, 1]
中選擇:- 1 表示標記未被掩碼,
- 0 表示標記已被掩碼。
- position_ids (
torch.LongTensor
,形狀為(batch_size, sequence_length)
,可選) — 位置嵌入中每個輸入序列標記的位置索引。在[0, config.n_positions - 1]
範圍內選擇。 - past_key_values (
~cache_utils.Cache
,可選) — 預計算的隱藏狀態(自注意力和交叉注意力塊中的鍵和值),可用於加速序列解碼。這通常包括在解碼的前一個階段,當use_cache=True
或config.use_cache=True
時,模型返回的past_key_values
。允許兩種格式:
- 一個 Cache 例項,請參閱我們的 kv 快取指南;
- 一個長度為
config.n_layers
的tuple(torch.FloatTensor)
元組,其中每個元組包含 2 個形狀為(batch_size, num_heads, sequence_length, embed_size_per_head)
的張量。這也稱為舊版快取格式。
模型將輸出與輸入相同的快取格式。如果沒有傳遞
past_key_values
,將返回舊版快取格式。如果使用
past_key_values
,使用者可以選擇只輸入最後的input_ids
(那些沒有提供其過去鍵值狀態給此模型的標記),形狀為(batch_size, 1)
,而不是所有形狀為(batch_size, sequence_length)
的input_ids
。 - inputs_embeds (
torch.FloatTensor
,形狀為(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
)。 - output_attentions (
bool
,可選) — 是否返回所有注意力層的注意力張量。有關更多詳細資訊,請參閱返回張量下的attentions
。 - output_hidden_states (
bool
,可選) — 是否返回所有層的隱藏狀態。有關更多詳細資訊,請參閱返回張量下的hidden_states
。
返回
transformers.modeling_outputs.SequenceClassifierOutputWithPast
或 tuple(torch.FloatTensor)
一個 transformers.modeling_outputs.SequenceClassifierOutputWithPast
或一個 torch.FloatTensor
元組(如果傳遞了 return_dict=False
或當 config.return_dict=False
時),包含根據配置(JambaConfig)和輸入而變化的不同元素。
-
loss (形狀為
(1,)
的torch.FloatTensor
,可選,當提供labels
時返回) — 分類損失(如果 config.num_labels==1,則為迴歸損失)。 -
logits (形狀為
(batch_size, config.num_labels)
的torch.FloatTensor
) — 分類(如果 config.num_labels==1,則為迴歸)分數(SoftMax 之前)。 -
past_key_values (
Cache
, optional, 當傳遞use_cache=True
或config.use_cache=True
時返回) — 這是一個 Cache 例項。更多詳情,請參閱我們的 kv 快取指南。包含預計算的隱藏狀態(自注意力塊中的鍵和值),可用於(參見
past_key_values
輸入)加速順序解碼。 -
hidden_states (
tuple(torch.FloatTensor)
, optional, 當傳遞output_hidden_states=True
或config.output_hidden_states=True
時返回) —torch.FloatTensor
的元組(一個用於嵌入層的輸出,如果模型有嵌入層,+ 一個用於每層的輸出),形狀為(batch_size, sequence_length, hidden_size)
。模型在每個層輸出的隱藏狀態以及可選的初始嵌入輸出。
-
attentions (
tuple(torch.FloatTensor)
, optional, 當傳遞output_attentions=True
或config.output_attentions=True
時返回) —torch.FloatTensor
的元組(每層一個),形狀為(batch_size, num_heads, sequence_length, sequence_length)
。注意力 softmax 後的注意力權重,用於計算自注意力頭中的加權平均值。
JambaForSequenceClassification 的前向方法重寫了 __call__
特殊方法。
儘管前向傳遞的邏輯需要在此函式內定義,但之後應呼叫 Module
例項而不是此函式,因為前者會處理執行前處理和後處理步驟,而後者會靜默地忽略它們。
單標籤分類示例
>>> import torch
>>> from transformers import AutoTokenizer, JambaForSequenceClassification
>>> tokenizer = AutoTokenizer.from_pretrained("ai21labs/Jamba-v0.1")
>>> model = JambaForSequenceClassification.from_pretrained("ai21labs/Jamba-v0.1")
>>> inputs = tokenizer("Hello, my dog is cute", return_tensors="pt")
>>> with torch.no_grad():
... logits = model(**inputs).logits
>>> predicted_class_id = logits.argmax().item()
>>> model.config.id2label[predicted_class_id]
...
>>> # To train a model on `num_labels` classes, you can pass `num_labels=num_labels` to `.from_pretrained(...)`
>>> num_labels = len(model.config.id2label)
>>> model = JambaForSequenceClassification.from_pretrained("ai21labs/Jamba-v0.1", num_labels=num_labels)
>>> labels = torch.tensor([1])
>>> loss = model(**inputs, labels=labels).loss
>>> round(loss.item(), 2)
...
多標籤分類示例
>>> import torch
>>> from transformers import AutoTokenizer, JambaForSequenceClassification
>>> tokenizer = AutoTokenizer.from_pretrained("ai21labs/Jamba-v0.1")
>>> model = JambaForSequenceClassification.from_pretrained("ai21labs/Jamba-v0.1", problem_type="multi_label_classification")
>>> inputs = tokenizer("Hello, my dog is cute", return_tensors="pt")
>>> with torch.no_grad():
... logits = model(**inputs).logits
>>> predicted_class_ids = torch.arange(0, logits.shape[-1])[torch.sigmoid(logits).squeeze(dim=0) > 0.5]
>>> # To train a model on `num_labels` classes, you can pass `num_labels=num_labels` to `.from_pretrained(...)`
>>> num_labels = len(model.config.id2label)
>>> model = JambaForSequenceClassification.from_pretrained(
... "ai21labs/Jamba-v0.1", num_labels=num_labels, problem_type="multi_label_classification"
... )
>>> labels = torch.sum(
... torch.nn.functional.one_hot(predicted_class_ids[None, :].clone(), num_classes=num_labels), dim=1
... ).to(torch.float)
>>> loss = model(**inputs, labels=labels).loss