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