Transformers documentation

T5Gemma

You are viewing main version, which requires installation from source. If you'd like regular pip install, checkout the latest stable version (v5.22.1).
Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

该模型于 2025-04-08 在 HF papers 上发布,并于 2025-06-25 贡献至 Hugging Face Transformers。

FlashAttention SDPA

T5Gemma

T5Gemma(又称 encoder-decoder Gemma)由 Google 在一篇研究论文中提出。它是一个 encoder-decoder 大语言模型系列,通过将预训练的 decoder-only 模型改造为 encoder-decoder 架构而开发。T5Gemma 包含预训练版本和指令微调版本。其架构基于遵循 T5 的 transformer encoder-decoder 设计,并引入了来自 Gemma 2 的改进:GQA、RoPE、GeGLU 激活函数、RMSNorm,以及交错的局部/全局注意力机制。

T5Gemma 包含两组模型规模:1) Gemma 2 规模(2B-2B、9B-2B 和 9B-9B),基于官方的 Gemma 2 模型(2B 和 9B);2) T5 规模(Small、Base、Large 和 XL),这些模型遵循 T5 配置,在 Gemma 2 框架下进行预训练。此外,我们还提供了一个 ML 规模(medium large,总计约 2B)的模型,其大小介于 T5 Large 和 T5 XL 之间。

预训练变体分别使用两个目标进行训练:带有知识蒸馏的前缀语言建模(PrefixLM)以及 UL2。我们为每个模型规模发布了这两种变体。指令化变体则通过监督微调和强化学习进行了后训练。

点击右侧边栏中的 T5Gemma 模型,以查看将 T5Gemma 应用于不同语言任务的更多示例。

以下示例演示了如何使用 Pipeline 或 AutoModel 类与模型进行对话,以及如何从命令行进行对话。

AutoModel
# pip install accelerate
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer


tokenizer = AutoTokenizer.from_pretrained("google/t5gemma-2b-2b-prefixlm-it")
model = AutoModelForSeq2SeqLM.from_pretrained(
    "google/t5gemma-2b-2b-prefixlm-it",
    device_map="auto",
)

messages = [
    {"role": "user", "content": "Tell me an unknown interesting biology fact about the brain."},
]
input_ids = tokenizer.apply_chat_template(messages, return_tensors="pt", return_dict=True, add_generation_prompt=True).to(model.device)

outputs = model.generate(**input_ids, max_new_tokens=32)
print(tokenizer.decode(outputs[0]))

T5GemmaConfig

类 transformers.T5GemmaConfig

< >

( transformers_version: str | None = Nonearchitectures: list[str] | None = Noneoutput_hidden_states: bool | None = Falsereturn_dict: bool | None = Truedtype: str | torch.dtype | None = Nonechunk_size_feed_forward: int = 0id2label: dict[int, str] | dict[str, str] | None = Nonelabel2id: dict[str, int] | dict[str, str] | None = Noneproblem_type: Literal['regression', 'single_label_classification', 'multi_label_classification'] | None = Noneis_encoder_decoder: bool = Trueencoder: transformers.models.t5gemma.configuration_t5gemma.T5GemmaModuleConfig | dict[typing.Any, typing.Any] | None = Nonedecoder: transformers.models.t5gemma.configuration_t5gemma.T5GemmaModuleConfig | dict[typing.Any, typing.Any] | None = Nonedropout_rate: int | float = 0.0classifier_dropout_rate: int | float = 0.0attention_dropout: float | int = 0.0tie_word_embeddings: bool = Truevocab_size: int = 256000 )

参数

  • is_encoder_decoder (bool,可选,默认为 True) — 指示该模型是否用作编码器/解码器。
  • encoder (Union[T5GemmaModuleConfig, dict],可选,optional) — 编码器的配置。
  • decoder (Union[T5GemmaModuleConfig, dict],可选,optional) — 解码器的配置。
  • dropout_rate (Union[int, float],可选,默认为 0.0) — 所有 dropout 层的比率。
  • classifier_dropout_rate (Union[int, float],可选,默认为 0.0) — 分类器的 dropout 比率。
  • attention_dropout (Union[float, int],可选,默认为 0.0) — 注意力概率的 dropout 比率。
  • tie_word_embeddings (bool,可选,默认为 True) — 是否根据模型的 tied_weights_keys 映射来绑定权重嵌入。
  • vocab_size (int,可选,默认为 256000) — 模型的词表大小。定义了 input_ids 可以表示的不同 token 数量。

这是用于存储 T5GemmaModel 配置的类。它根据指定的参数实例化 T5Gemma 模型,并定义模型架构。使用默认值实例化配置,将得到与 google/t5_gemma_module-7b 类似的配置。

配置对象继承自 PreTrainedConfig,可用于控制模型输出。有关更多信息,请参阅 PreTrainedConfig 的文档。

示例:

>>> from transformers import T5GemmaConfig, T5GemmaModel
>>> t5gemma_config = T5GemmaConfig.from_pretrained("google/t5gemma-2b-2b-prefixlm-it")
>>> model = T5GemmaModel(t5gemma_config)

T5GemmaModuleConfig

类 transformers.T5GemmaModuleConfig

< >

( transformers_version: str | None = Nonearchitectures: list[str] | None = Noneoutput_hidden_states: bool | None = Falsereturn_dict: bool | None = Truedtype: str | torch.dtype | None = Nonechunk_size_feed_forward: int = 0is_encoder_decoder: bool = Falseid2label: dict[int, str] | dict[str, str] | None = Nonelabel2id: dict[str, int] | dict[str, str] | None = Noneproblem_type: Literal['regression', 'single_label_classification', 'multi_label_classification'] | None = Nonevocab_size: int = 256000hidden_size: int = 2304intermediate_size: int = 9216num_hidden_layers: int = 26num_attention_heads: int = 8num_key_value_heads: int = 4head_dim: int = 256hidden_activation: str = 'gelu_pytorch_tanh'max_position_embeddings: int = 8192initializer_range: float = 0.02rms_norm_eps: float = 1e-06use_cache: bool = Truepad_token_id: int | None = 0eos_token_id: int | list[int] | None = 1bos_token_id: int | None = 2tie_word_embeddings: bool = Truerope_parameters: transformers.modeling_rope_utils.RopeParameters | dict | None = Noneattention_bias: bool = Falseattention_dropout: int | float | None = 0.0query_pre_attn_scalar: int = 256sliding_window: int | None = 4096layer_types: list[str] | None = Nonefinal_logit_softcapping: float | None = 30.0attn_logit_softcapping: float | None = 50.0is_decoder: bool = False )

参数

  • vocab_size (int,可选,默认为 256000) — 模型的词表大小。定义了 input_ids 可以表示的不同 token 数量。
  • hidden_size (int,可选,默认为 2304) — 隐藏表示的维度。
  • intermediate_size (int,可选,默认为 9216) — MLP 表示的维度。
  • num_hidden_layers (int,可选,默认为 26) — Transformer 解码器中的隐藏层数量。
  • num_attention_heads (int,可选,默认为 8) — Transformer 解码器中每个注意力层的注意力头数量。
  • num_key_value_heads (int, 可选,默认为 4) — 这是用于实现 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 头应通过对该组内所有原始头进行 meanpooling 来构建。更多详情请参阅 这篇论文。如果未指定,则默认为 num_attention_heads。
  • head_dim (int,可选,默认为 256) — 注意力头维度。如果为 None,则默认为 hidden_size // num_attention_heads
  • hidden_activation (str, 可选, 默认为 gelu_pytorch_tanh) — 解码器中的非线性激活函数(函数或字符串)。例如,"gelu"、 "relu"、"silu" 等。
  • max_position_embeddings (int,可选,默认为 8192) — 该模型可能使用的最大序列长度。
  • initializer_range (float,可选,默认为 0.02) — 用于初始化所有权重矩阵的 truncated_normal_initializer 的标准差。
  • rms_norm_eps (float,可选,默认为 1e-06) — RMS 归一化层使用的 epsilon 值。
  • use_cache (bool,可选,默认为 True) — 模型是否应返回最后一次键/值注意力(并非所有模型都使用)。仅当 config.is_decoder=True 或模型为仅解码器生成式模型时相关。
  • pad_token_id (int, 可选,默认为 0) — 词表中用于填充的 token id。
  • eos_token_id (Union[int, list[int]],可选,默认为 1) — 词汇表中用于表示流结束(end-of-stream)的 token id。
  • bos_token_id (int,可选,默认为 2) — 词汇表中用于表示流起始的 Token id。
  • tie_word_embeddings (bool,可选,默认为 True) — 是否根据模型的 tied_weights_keys 映射来绑定权重嵌入。
  • rope_parameters (Union[~modeling_rope_utils.RopeParameters, dict], 可选) — 包含 RoPE 嵌入配置参数的字典。该字典应包含 rope_theta 的值,并且如果希望使用 RoPE 处理更长的 max_position_embeddings,则可选地包含用于缩放的参数。
  • attention_bias (bool,可选,默认为 False) — 在自注意力机制中,是否在 query、key、value 和 output 投影层中使用偏置。
  • attention_dropout (Union[int, float],可选,默认为 0.0) — 注意力概率的 dropout 比率。
  • query_pre_attn_scalar(float,可选,默认为 256)— 用于缩放注意力得分的系数
  • sliding_window (int,可选,默认为 4096) — 滑动窗口注意力的窗口大小。如果为 None,则不应用滑动窗口。
  • layer_types (list[str], 可选) — 一个列表,用于显式地将每个层索引映射到其层类型。如果未提供,将根据 config 值自动生成。
  • final_logit_softcapping (float,可选,默认值为 30.0) — 在对 logits 应用 tanh softcapping 时使用的缩放因子。
  • attn_logit_softcapping(float,可选,默认值为 50.0)—— 在对注意力分数应用 tanh 软截断时的缩放因子。
  • is_decoder (bool,可选,默认为 False) — 指示该模型是否用作解码器。如果为 False,则模型用作编码器。

这是用于存储 T5GemmaModel 配置的类。它根据指定的参数实例化 T5Gemma 模型,并定义模型架构。使用默认值实例化配置,将得到与 google/t5_gemma_module-7b 类似的配置。

配置对象继承自 PreTrainedConfig,可用于控制模型输出。有关更多信息,请参阅 PreTrainedConfig 的文档。

>>> from transformers import T5GemmaModuleModel, T5GemmaModuleConfig
>>> # Initializing a T5GemmaModule t5_gemma_module-7b style configuration
>>> configuration = T5GemmaModuleConfig()
>>> # Initializing a model from the t5_gemma_module-7b style configuration
>>> model = T5GemmaModuleModel(configuration)
>>> # Accessing the model configuration
>>> configuration = model.config

T5GemmaModel

class transformers.T5GemmaModel

< >

( config: T5GemmaConfig )

参数

  • config (T5GemmaConfig) — 模型的配置类,包含模型的所有参数。使用配置文件进行初始化时,不会加载与模型关联的权重,仅加载配置。请查看 from_pretrained() 方法以加载模型权重。

未添加任何特定头部的裸 T5Gemma 模型,直接输出原始隐藏状态。

该模型继承自 PreTrainedModel。请查阅父类文档,了解该库为所有模型实现的通用方法(例如下载或保存、调整输入嵌入大小、剪枝头等)。

该模型也是 PyTorch torch.nn.Module 的子类。 请将其作为常规的 PyTorch Module 使用,有关通用用法和行为的所有事项,请参阅 PyTorch 文档。

forward

< >

( input_ids: typing.Optional[torch.LongTensor] = Noneattention_mask: typing.Optional[torch.Tensor] = Noneposition_ids: typing.Optional[torch.LongTensor] = Nonedecoder_input_ids: typing.Optional[torch.LongTensor] = Nonedecoder_attention_mask: typing.Optional[torch.BoolTensor] = Nonedecoder_position_ids: typing.Optional[torch.LongTensor] = Noneencoder_outputs: transformers.modeling_outputs.BaseModelOutput | None = Nonepast_key_values: transformers.cache_utils.EncoderDecoderCache | None = Noneinputs_embeds: typing.Optional[torch.Tensor] = Nonedecoder_inputs_embeds: typing.Optional[torch.Tensor] = Noneuse_cache: bool | None = None**kwargs: *<class 'transformers.utils.generic.TransformersKwargs'> ) → Seq2SeqModelOutput or tuple(torch.Tensor)

参数

  • input_ids (torch.LongTensor 形状为 (batch_size, sequence_length), ,可选) — 输入序列中各 token 在词汇表中的索引。默认情况下,填充(padding)将被忽略。

    可以使用 AutoTokenizer 获取索引。有关详细信息,请参阅 PreTrainedTokenizer.encode() 和 PreTrainedTokenizer.call()。

    什么是输入 ID?

  • attention_mask (torch.Tensor 形状为 (batch_size, sequence_length), ,可选) — 用于避免对填充 token 索引执行注意力机制的掩码。掩码值取自 [0, 1]:
    • 1 表示未被屏蔽的 token,
    • 对于被掩码的 token,值为 0。

    什么是注意力掩码?

  • position_ids (torch.LongTensor 形状为 (batch_size, sequence_length), 可选) — 输入序列中每个 token 在位置嵌入中的位置索引。取值范围为 [0, config.n_positions - 1].

    什么是位置 ID?

  • decoder_input_ids (torch.LongTensor 形状为 (batch_size, target_sequence_length), ,可选) — 解码器输入序列中各 token 在词汇表中的索引。可通过 AutoTokenizer 获取索引。详情请参阅 PreTrainedTokenizer.encode() 和 PreTrainedTokenizer.call()。什么是解码器输入 ID?

    可以使用 AutoTokenizer 获取索引。有关详细信息,请参阅 PreTrainedTokenizer.encode() 和 PreTrainedTokenizer.call()。

    什么是解码器输入 ID?

  • decoder_attention_mask(形状为 (batch_size, target_sequence_length) 的 torch.BoolTensor,可选) — 用于避免对特定 token 索引执行注意力的掩码。默认情况下,会使用因果掩码,以确保模型只能查看之前的输入来预测未来。
  • decoder_position_ids (torch.LongTensor,形状为 (batch_size, decoder_sequence_length),可选) — 解码器输入序列中每个 token 在位置嵌入中的位置索引。取值范围为 [0, config.decoder.n_positions - 1]。什么是位置 ID?
  • encoder_outputs (~modeling_outputs.BaseModelOutput, 可选) — 元组由 (last_hidden_state, 可选: hidden_states, 可选: attentions) last_hidden_state 的形状为 (batch_size, sequence_length, hidden_size), 可选),是编码器最后一层输出的隐藏状态序列。用于解码器的交叉注意力。
  • past_key_values (~cache_utils.EncoderDecoderCache, 可选) — 预计算的隐藏状态(自注意力块和交叉注意力块中的键和值),可用于加速顺序解码。这通常由 past_key_values 在解码的先前阶段返回,当 use_cache=True 或 config.use_cache=True.

    仅允许将 Cache 实例作为输入,请参阅我们的 kv cache 指南。 如果未传入 past_key_values,则默认会初始化 DynamicCache。

    模型将输出与输入相同的缓存格式。

    如果使用了 past_key_values,用户应仅输入未处理的 input_ids(即未向该模型提供其过去键值状态的那些),其形状为 (batch_size, unprocessed_length),而不是形状为 (batch_size, sequence_length) 的全部 input_ids。

  • inputs_embeds(形状为 (batch_size, sequence_length, hidden_size) 的 torch.Tensor,可选) — 可选地,您可以选择直接传入嵌入表示,而不是传入 input_ids。如果您希望对 input_ids 索引如何转换为相应向量拥有比模型内部嵌入查找矩阵更多的控制权,这将非常有用。
  • decoder_inputs_embeds (torch.Tensor 形状为 (batch_size, target_sequence_length, hidden_size), ,可选) — 可选地,除了传递 decoder_input_ids ,还可以选择直接传递嵌入表示。如果使用了 past_key_values ,则只需输入最后一个 decoder_inputs_embeds (参见 past_key_values). 。这在希望比模型内部的嵌入查找矩阵更精细地控制如何将 decoder_input_ids 索引转换为对应向量时非常有用。

    如果 decoder_input_ids 和 decoder_inputs_embeds 均未设置,decoder_inputs_embeds 将取 inputs_embeds 的值。

  • use_cache (bool,可选) — 如果设置为 True,则会返回 past_key_values 键值状态,可用于加速解码(参见 past_key_values)。

返回值

Seq2SeqModelOutput 或 tuple(torch.Tensor)

一个 Seq2SeqModelOutput 或 torch.Tensor 元组(如果传入了 return_dict=False 或 config.return_dict=False),其包含的元素取决于配置(T5GemmaConfig)和输入。

T5GemmaModel 的 forward 方法重写了 __call__ 特殊方法。

尽管前向传播的实现需要在此函数中定义,但之后应调用 Module 实例,而不是直接调用此函数,因为前者会负责执行预处理和后处理步骤,而后者则会静默地忽略这些步骤。

  • last_hidden_state(形状为 (batch_size, sequence_length, hidden_size) 的 torch.Tensor)— 模型解码器最后一层输出的隐藏状态序列。

    如果使用了 past_key_values,则仅输出形状为 (batch_size, 1, hidden_size) 的序列的最后一个隐藏状态。

  • past_key_values(EncoderDecoderCache,可选,在传入 use_cache=True 或 config.use_cache=True 时返回)—— 它是一个 EncoderDecoderCache 实例。更多详情,请参阅我们的 kv cache 指南。

    包含预计算的隐藏状态(自注意力块和交叉注意力块中的键和值),可用于(参见 past_key_values 输入)加速顺序解码。

  • decoder_hidden_states (tuple(torch.Tensor),可选,在传入 output_hidden_states=True 或 config.output_hidden_states=True 时返回) — 一个 torch.Tensor 元组(如果模型具有嵌入层,则包含嵌入层输出的一个张量,以及每个层输出的一个张量),形状为 (batch_size, sequence_length, hidden_size)。

    解码器在每一层输出处的隐藏状态,以及可选的初始嵌入输出。

  • decoder_attentions (tuple(torch.Tensor),可选,在传入 output_attentions=True 或 config.output_attentions=True 时返回) — 一个 torch.Tensor 元组(每层一个),形状为 (batch_size, num_heads, sequence_length, sequence_length)。

    解码器在注意力 softmax 之后的注意力权重,用于在自注意力头中计算加权平均值。

  • cross_attentions (tuple(torch.Tensor),可选,在传入 output_attentions=True 或 config.output_attentions=True 时返回) — 一个 torch.Tensor 元组(每层一个),形状为 (batch_size, num_heads, sequence_length, sequence_length)。

    解码器交叉注意力层在注意力 softmax 之后的注意力权重,用于在交叉注意力头中计算加权平均值。

  • encoder_last_hidden_state(形状为 (batch_size, sequence_length, hidden_size) 的 torch.Tensor,可选)— 模型编码器最后一层输出的隐藏状态序列。

  • encoder_hidden_states (tuple(torch.Tensor),可选,在传入 output_hidden_states=True 或 config.output_hidden_states=True 时返回) — 一个 torch.Tensor 元组(如果模型具有嵌入层,则包含嵌入层输出的一个张量,以及每个层输出的一个张量),形状为 (batch_size, sequence_length, hidden_size)。

    编码器每层输出处的隐藏状态,以及可选的初始嵌入输出。

  • encoder_attentions(tuple(torch.Tensor),可选,在传入 output_attentions=True 或 config.output_attentions=True 时返回)—— 一个 torch.Tensor 元组(每层一个),形状为 (batch_size, num_heads, sequence_length, sequence_length)。

    编码器在注意力 softmax 之后的注意力权重,用于在自注意力头中计算加权平均值。

T5GemmaEncoderModel

类 transformers.T5GemmaEncoderModel

< >

( config: T5GemmaConfig )

参数

  • config (T5GemmaConfig) — 模型的配置类,包含模型的所有参数。使用配置文件进行初始化时,不会加载与模型关联的权重,仅加载配置。请查看 from_pretrained() 方法以加载模型权重。

未添加任何特定头部的裸 T5Gemma 模型,直接输出原始隐藏状态。

该模型继承自 PreTrainedModel。请查阅父类文档,了解该库为所有模型实现的通用方法(例如下载或保存、调整输入嵌入大小、剪枝头等)。

该模型也是 PyTorch torch.nn.Module 的子类。 请将其作为常规的 PyTorch Module 使用,有关通用用法和行为的所有事项,请参阅 PyTorch 文档。

forward

< >

( input_ids: typing.Optional[torch.LongTensor] = Noneattention_mask: typing.Optional[torch.Tensor] = Noneposition_ids: typing.Optional[torch.LongTensor] = Noneinputs_embeds: typing.Optional[torch.Tensor] = None**kwargs: *<class 'transformers.utils.generic.TransformersKwargs'> ) → BaseModelOutput or tuple(torch.Tensor)

参数

  • input_ids (torch.LongTensor 形状为 (batch_size, sequence_length), ,可选) — 输入序列中各 token 在词汇表中的索引。默认情况下,填充(padding)将被忽略。

    可以使用 AutoTokenizer 获取索引。有关详细信息,请参阅 PreTrainedTokenizer.encode() 和 PreTrainedTokenizer.call()。

    什么是输入 ID?

  • attention_mask (torch.Tensor 形状为 (batch_size, sequence_length), ,可选) — 用于避免对填充 token 索引执行注意力机制的掩码。掩码值取自 [0, 1]:
    • 1 表示未被屏蔽的 token,
    • 对于被掩码的 token,值为 0。

    什么是注意力掩码?

  • position_ids (torch.LongTensor 形状为 (batch_size, sequence_length), 可选) — 输入序列中每个 token 在位置嵌入中的位置索引。取值范围为 [0, config.n_positions - 1].

    什么是位置 ID?

  • inputs_embeds(形状为 (batch_size, sequence_length, hidden_size) 的 torch.Tensor,可选) — 可选地,您可以选择直接传入嵌入表示,而不是传入 input_ids。如果您希望对 input_ids 索引如何转换为相应向量拥有比模型内部嵌入查找矩阵更多的控制权,这将非常有用。

返回值

BaseModelOutput 或 tuple(torch.Tensor)

一个 BaseModelOutput 或 torch.Tensor 元组(如果传入了 return_dict=False 或 config.return_dict=False),其包含的元素取决于配置(T5GemmaConfig)和输入。

T5GemmaEncoderModel 的 forward 方法重写了 __call__ 特殊方法。

尽管前向传播的实现需要在此函数中定义,但之后应调用 Module 实例,而不是直接调用此函数,因为前者会负责执行预处理和后处理步骤,而后者则会静默地忽略这些步骤。

  • last_hidden_state(形状为 (batch_size, sequence_length, hidden_size) 的 torch.Tensor)— 模型最后一层输出的隐藏状态序列。

  • hidden_states (tuple(torch.Tensor),可选,当传入 output_hidden_states=True 或 config.output_hidden_states=True 时返回) — 一个 torch.Tensor 元组(如果模型具有嵌入层,则包含嵌入层的输出,以及每一层的输出),形状为 (batch_size, sequence_length, hidden_size)。

    模型在每一层输出处的隐藏状态,以及可选的初始嵌入输出。

  • attentions (tuple(torch.Tensor),可选,当传入 output_attentions=True 或 config.output_attentions=True 时返回) — 一个 torch.Tensor 元组(每层一个),形状为 (batch_size, num_heads, sequence_length, sequence_length)。

    注意力权重,即经过注意力 softmax 之后的权重,用于在自注意力头中计算加权平均值。

T5GemmaForConditionalGeneration

类 transformers.T5GemmaForConditionalGeneration

< >

( config: T5GemmaConfig )

forward

< >

( input_ids: typing.Optional[torch.LongTensor] = Noneattention_mask: typing.Optional[torch.Tensor] = Noneposition_ids: typing.Optional[torch.LongTensor] = Nonedecoder_input_ids: typing.Optional[torch.LongTensor] = Nonedecoder_attention_mask: typing.Optional[torch.BoolTensor] = Nonedecoder_position_ids: typing.Optional[torch.LongTensor] = Noneencoder_outputs: transformers.modeling_outputs.BaseModelOutput | None = Nonepast_key_values: transformers.cache_utils.EncoderDecoderCache | None = Noneinputs_embeds: typing.Optional[torch.Tensor] = Nonedecoder_inputs_embeds: typing.Optional[torch.Tensor] = Nonelabels: typing.Optional[torch.LongTensor] = Noneuse_cache: bool | None = Nonelogits_to_keep: typing.Union[int, torch.Tensor] = 0**kwargs: *<class 'transformers.utils.generic.TransformersKwargs'> ) → Seq2SeqLMOutput or tuple(torch.Tensor)

参数

  • input_ids (torch.LongTensor 形状为 (batch_size, sequence_length), ,可选) — 输入序列中各 token 在词汇表中的索引。默认情况下,填充(padding)将被忽略。

    可以使用 AutoTokenizer 获取索引。有关详细信息,请参阅 PreTrainedTokenizer.encode() 和 PreTrainedTokenizer.call()。

    什么是输入 ID?

  • attention_mask (torch.Tensor 形状为 (batch_size, sequence_length), ,可选) — 用于避免对填充 token 索引执行注意力机制的掩码。掩码值取自 [0, 1]:
    • 1 表示未被屏蔽的 token,
    • 对于被掩码的 token,值为 0。

    什么是注意力掩码?

  • position_ids (torch.LongTensor 形状为 (batch_size, sequence_length), 可选) — 输入序列中每个 token 在位置嵌入中的位置索引。取值范围为 [0, config.n_positions - 1].

    什么是位置 ID?

  • decoder_input_ids (torch.LongTensor 形状为 (batch_size, target_sequence_length), ,可选) — 解码器输入序列中各 token 在词汇表中的索引。可通过 AutoTokenizer 获取索引。详情请参阅 PreTrainedTokenizer.encode() 和 PreTrainedTokenizer.call()。什么是解码器输入 ID?

    可以使用 AutoTokenizer 获取索引。有关详细信息,请参阅 PreTrainedTokenizer.encode() 和 PreTrainedTokenizer.call()。

    什么是解码器输入 ID?

  • decoder_attention_mask(形状为 (batch_size, target_sequence_length) 的 torch.BoolTensor,可选) — 用于避免对特定 token 索引执行注意力的掩码。默认情况下,会使用因果掩码,以确保模型只能查看之前的输入来预测未来。
  • decoder_position_ids (torch.LongTensor,形状为 (batch_size, decoder_sequence_length),可选) — 解码器输入序列中每个 token 在位置嵌入中的位置索引。取值范围为 [0, config.decoder.n_positions - 1]。什么是位置 ID?
  • encoder_outputs (~modeling_outputs.BaseModelOutput, 可选) — 元组由 (last_hidden_state, 可选: hidden_states, 可选: attentions) last_hidden_state 的形状为 (batch_size, sequence_length, hidden_size), 可选),是编码器最后一层输出的隐藏状态序列。用于解码器的交叉注意力。
  • past_key_values (~cache_utils.EncoderDecoderCache, 可选) — 预计算的隐藏状态(自注意力块和交叉注意力块中的键和值),可用于加速顺序解码。这通常由 past_key_values 在解码的先前阶段返回,当 use_cache=True 或 config.use_cache=True.

    仅允许将 Cache 实例作为输入,请参阅我们的 kv cache 指南。 如果未传入 past_key_values,则默认会初始化 DynamicCache。

    模型将输出与输入相同的缓存格式。

    如果使用了 past_key_values,用户应仅输入未处理的 input_ids(即未向该模型提供其过去键值状态的那些),其形状为 (batch_size, unprocessed_length),而不是形状为 (batch_size, sequence_length) 的全部 input_ids。

  • inputs_embeds(形状为 (batch_size, sequence_length, hidden_size) 的 torch.Tensor,可选) — 可选地,您可以选择直接传入嵌入表示,而不是传入 input_ids。如果您希望对 input_ids 索引如何转换为相应向量拥有比模型内部嵌入查找矩阵更多的控制权,这将非常有用。
  • decoder_inputs_embeds (torch.Tensor 形状为 (batch_size, target_sequence_length, hidden_size), ,可选) — 可选地,除了传递 decoder_input_ids ,还可以选择直接传递嵌入表示。如果使用了 past_key_values ,则只需输入最后一个 decoder_inputs_embeds (参见 past_key_values). 。这在希望比模型内部的嵌入查找矩阵更精细地控制如何将 decoder_input_ids 索引转换为对应向量时非常有用。

    如果 decoder_input_ids 和 decoder_inputs_embeds 均未设置,decoder_inputs_embeds 将取 inputs_embeds 的值。

  • labels (torch.LongTensor,形状为 (batch_size, sequence_length),可选) — 用于计算掩码语言建模损失的标签。索引值应在 [0, ..., config.vocab_size] 范围内,或为 -100(参见 input_ids 的文档字符串)。索引值设为 -100 的 token 将被忽略(即被掩码),损失仅针对标签位于 [0, ..., config.vocab_size] 范围内的 token 进行计算。
  • use_cache (bool,可选) — 如果设置为 True,则会返回 past_key_values 键值状态,可用于加速解码(参见 past_key_values)。
  • logits_to_keep (Union[int, torch.Tensor], 可选,默认为 0) — 如果为 int,则计算最后 logits_to_keep 个 token 的 logits。如果为 0,则计算所有 input_ids 的 logits(特殊情况)。生成过程只需要最后一个 token 的 logits,仅针对该 token 进行计算可以节省内存,对于长序列或大词汇量而言,这种节省相当显著。 如果为 torch.Tensor,则必须是一维张量,对应于序列长度维度中需要保留的索引。 这在使用打包张量格式(将批次和序列长度合并为单一维度)时非常有用。

返回值

Seq2SeqLMOutput 或 tuple(torch.Tensor)

一个 Seq2SeqLMOutput 或 torch.Tensor 元组(如果传入了 return_dict=False 或 config.return_dict=False),其包含的元素取决于配置(T5GemmaConfig)和输入。

T5GemmaForConditionalGeneration 的 forward 方法重写了 __call__ 特殊方法。

尽管前向传播的实现需要在此函数中定义,但之后应调用 Module 实例,而不是直接调用此函数,因为前者会负责执行预处理和后处理步骤,而后者则会静默地忽略这些步骤。

  • loss(形状为 (1,) 的 torch.Tensor,可选,在提供 labels 时返回)— 语言建模损失。

  • logits(形状为 (batch_size, sequence_length, config.vocab_size) 的 torch.Tensor)— 语言建模头的预测分数(SoftMax 之前每个词汇表 token 的分数)。

  • past_key_values(EncoderDecoderCache,可选,在传入 use_cache=True 或 config.use_cache=True 时返回)—— 它是一个 EncoderDecoderCache 实例。更多详情,请参阅我们的 kv cache 指南。

    包含预计算的隐藏状态(自注意力块和交叉注意力块中的键和值),可用于(参见 past_key_values 输入)加速顺序解码。

  • decoder_hidden_states (tuple(torch.Tensor),可选,在传入 output_hidden_states=True 或 config.output_hidden_states=True 时返回) — 一个 torch.Tensor 元组(如果模型具有嵌入层,则包含嵌入层输出的一个张量,以及每个层输出的一个张量),形状为 (batch_size, sequence_length, hidden_size)。

    解码器在每一层输出处的隐藏状态,以及初始嵌入输出。

  • decoder_attentions (tuple(torch.Tensor),可选,在传入 output_attentions=True 或 config.output_attentions=True 时返回) — 一个 torch.Tensor 元组(每层一个),形状为 (batch_size, num_heads, sequence_length, sequence_length)。

    解码器在注意力 softmax 之后的注意力权重,用于在自注意力头中计算加权平均值。

  • cross_attentions (tuple(torch.Tensor),可选,在传入 output_attentions=True 或 config.output_attentions=True 时返回) — 一个 torch.Tensor 元组(每层一个),形状为 (batch_size, num_heads, sequence_length, sequence_length)。

    解码器交叉注意力层在注意力 softmax 之后的注意力权重,用于在交叉注意力头中计算加权平均值。

  • encoder_last_hidden_state(形状为 (batch_size, sequence_length, hidden_size) 的 torch.Tensor,可选)— 模型编码器最后一层输出的隐藏状态序列。

  • encoder_hidden_states (tuple(torch.Tensor),可选,在传入 output_hidden_states=True 或 config.output_hidden_states=True 时返回) — 一个 torch.Tensor 元组(如果模型具有嵌入层,则包含嵌入层输出的一个张量,以及每个层输出的一个张量),形状为 (batch_size, sequence_length, hidden_size)。

    编码器在每一层输出处的隐藏状态,以及初始嵌入输出。

  • encoder_attentions(tuple(torch.Tensor),可选,在传入 output_attentions=True 或 config.output_attentions=True 时返回)—— 一个 torch.Tensor 元组(每层一个),形状为 (batch_size, num_heads, sequence_length, sequence_length)。

    编码器在注意力 softmax 之后的注意力权重,用于在自注意力头中计算加权平均值。

示例:

T5GemmaForSequenceClassification

class transformers.T5GemmaForSequenceClassification

< >

( config: T5GemmaConfigis_encoder_decoder: bool | None = None )

参数

  • config (T5GemmaConfig) — 模型的配置类,包含模型的所有参数。使用配置文件进行初始化时,不会加载与模型关联的权重,仅加载配置。请查看 from_pretrained() 方法以加载模型权重。
  • is_encoder_decoder (Optional,可选) — 是否使用 encoder_decoder 进行序列分类。当设置为 False 时,仅使用 encoder。

带有序列分类/回归头的 T5Gemma 模型,例如用于 GLUE 任务。

该模型继承自 PreTrainedModel。请查阅父类文档,了解该库为所有模型实现的通用方法(例如下载或保存、调整输入嵌入大小、剪枝头等)。

该模型也是 PyTorch torch.nn.Module 的子类。 请将其作为常规的 PyTorch Module 使用,有关通用用法和行为的所有事项,请参阅 PyTorch 文档。

forward

< >

( input_ids: typing.Optional[torch.LongTensor] = Noneattention_mask: typing.Optional[torch.Tensor] = Noneposition_ids: typing.Optional[torch.LongTensor] = Nonedecoder_input_ids: typing.Optional[torch.LongTensor] = Nonedecoder_attention_mask: typing.Optional[torch.Tensor] = Nonedecoder_position_ids: typing.Optional[torch.LongTensor] = Noneencoder_outputs: transformers.modeling_outputs.BaseModelOutput | None = Noneinputs_embeds: typing.Optional[torch.Tensor] = Nonedecoder_inputs_embeds: typing.Optional[torch.Tensor] = Nonelabels: typing.Optional[torch.LongTensor] = None**kwargs: *<class 'transformers.utils.generic.TransformersKwargs'> ) → SequenceClassifierOutput or tuple(torch.Tensor)

参数

  • input_ids (torch.LongTensor 形状为 (batch_size, sequence_length), ,可选) — 输入序列中各 token 在词汇表中的索引。默认情况下,填充(padding)将被忽略。

    可以使用 AutoTokenizer 获取索引。有关详细信息,请参阅 PreTrainedTokenizer.encode() 和 PreTrainedTokenizer.call()。

    什么是输入 ID?

  • attention_mask (torch.Tensor 形状为 (batch_size, sequence_length), ,可选) — 用于避免对填充 token 索引执行注意力机制的掩码。掩码值取自 [0, 1]:
    • 1 表示未被屏蔽的 token,
    • 对于被掩码的 token,值为 0。

    什么是注意力掩码?

  • position_ids (torch.LongTensor 形状为 (batch_size, sequence_length), 可选) — 输入序列中每个 token 在位置嵌入中的位置索引。取值范围为 [0, config.n_positions - 1].

    什么是位置 ID?

  • decoder_input_ids (torch.LongTensor 形状为 (batch_size, target_sequence_length), ,可选) — 解码器输入序列中各 token 在词汇表中的索引。可通过 AutoTokenizer 获取索引。详情请参阅 PreTrainedTokenizer.encode() 和 PreTrainedTokenizer.call()。什么是解码器输入 ID?

    可以使用 AutoTokenizer 获取索引。有关详细信息,请参阅 PreTrainedTokenizer.encode() 和 PreTrainedTokenizer.call()。

    什么是解码器输入 ID?

  • decoder_attention_mask(形状为 (batch_size, target_sequence_length) 的 torch.Tensor,可选) — 用于避免对特定 token 索引执行注意力机制的掩码。默认情况下,会使用因果掩码,以确保模型只能查看之前的输入来预测未来。
  • decoder_position_ids (torch.LongTensor,形状为 (batch_size, decoder_sequence_length),可选) — 解码器输入序列中每个 token 在位置嵌入中的位置索引。取值范围为 [0, config.decoder.n_positions - 1]。什么是位置 ID?
  • encoder_outputs (~modeling_outputs.BaseModelOutput, 可选) — 元组由 (last_hidden_state, 可选: hidden_states, 可选: attentions) last_hidden_state 的形状为 (batch_size, sequence_length, hidden_size), 可选),是编码器最后一层输出的隐藏状态序列。用于解码器的交叉注意力。
  • inputs_embeds(形状为 (batch_size, sequence_length, hidden_size) 的 torch.Tensor,可选) — 可选地,您可以选择直接传入嵌入表示,而不是传入 input_ids。如果您希望对 input_ids 索引如何转换为相应向量拥有比模型内部嵌入查找矩阵更多的控制权,这将非常有用。
  • decoder_inputs_embeds (torch.Tensor 形状为 (batch_size, target_sequence_length, hidden_size), ,可选) — 可选地,除了传递 decoder_input_ids ,还可以选择直接传递嵌入表示。如果使用了 past_key_values ,则只需输入最后一个 decoder_inputs_embeds (参见 past_key_values). 。这在希望比模型内部的嵌入查找矩阵更精细地控制如何将 decoder_input_ids 索引转换为对应向量时非常有用。

    如果 decoder_input_ids 和 decoder_inputs_embeds 均未设置,decoder_inputs_embeds 将取 inputs_embeds 的值。

  • labels (torch.LongTensor,形状为 (batch_size,),可选) — 用于计算序列分类/回归损失的标签。索引应在 [0, ..., config.num_labels - 1] 范围内。如果 config.num_labels == 1,则计算回归损失(均方损失);如果 config.num_labels > 1,则计算分类损失(交叉熵)。

返回值

SequenceClassifierOutput 或 tuple(torch.Tensor)

一个 SequenceClassifierOutput 或 torch.Tensor 元组(如果传入了 return_dict=False 或 config.return_dict=False),其包含的元素取决于配置(T5GemmaConfig)和输入。

T5GemmaForSequenceClassification 的 forward 方法重写了 __call__ 特殊方法。

尽管前向传播的实现需要在此函数中定义,但之后应调用 Module 实例,而不是直接调用此函数,因为前者会负责执行预处理和后处理步骤,而后者则会静默地忽略这些步骤。

  • loss (torch.Tensor,形状为 (1,),可选,当提供 labels 时返回) — 分类损失(如果 config.num_labels==1,则为回归损失)。

  • logits(形状为 (batch_size, config.num_labels) 的 torch.Tensor)— 分类(或当 config.num_labels==1 时为回归)得分(SoftMax 之前)。

  • hidden_states (tuple(torch.Tensor),可选,当传入 output_hidden_states=True 或 config.output_hidden_states=True 时返回) — 一个 torch.Tensor 元组(如果模型具有嵌入层,则包含嵌入层的输出,以及每一层的输出),形状为 (batch_size, sequence_length, hidden_size)。

    模型在每一层输出处的隐藏状态,以及可选的初始嵌入输出。

  • attentions (tuple(torch.Tensor),可选,当传入 output_attentions=True 或 config.output_attentions=True 时返回) — 一个 torch.Tensor 元组(每层一个),形状为 (batch_size, num_heads, sequence_length, sequence_length)。

    注意力权重,即经过注意力 softmax 之后的权重,用于在自注意力头中计算加权平均值。

单标签分类示例:

>>> import torch
>>> from transformers import AutoTokenizer, T5GemmaForSequenceClassification

>>> tokenizer = AutoTokenizer.from_pretrained("google/t5_gemma_module-7b")
>>> model = T5GemmaForSequenceClassification.from_pretrained("google/t5_gemma_module-7b")

>>> 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 = T5GemmaForSequenceClassification.from_pretrained("google/t5_gemma_module-7b", num_labels=num_labels)

>>> labels = torch.tensor([1])
>>> loss = model(**inputs, labels=labels).loss
>>> round(loss.item(), 2)
...

多标签分类示例:

>>> import torch
>>> from transformers import AutoTokenizer, T5GemmaForSequenceClassification

>>> tokenizer = AutoTokenizer.from_pretrained("google/t5_gemma_module-7b")
>>> model = T5GemmaForSequenceClassification.from_pretrained("google/t5_gemma_module-7b", 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 = T5GemmaForSequenceClassification.from_pretrained(
...     "google/t5_gemma_module-7b", 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

T5GemmaForTokenClassification

class transformers.T5GemmaForTokenClassification

< >

( config: T5GemmaConfigis_encoder_decoder: bool | None = None )

参数

  • config (T5GemmaConfig) — 模型的配置类,包含模型的所有参数。使用配置文件进行初始化时,不会加载与模型关联的权重,仅加载配置。请查看 from_pretrained() 方法以加载模型权重。
  • is_encoder_decoder (Optional,可选) — 是否使用 encoder_decoder 进行 token 分类。设置为 False 时,仅使用 encoder。

带有 token 分类头(即在隐藏状态输出之上添加一个线性层)的 T5Gemma 变换器,例如用于命名实体识别(NER)任务。

该模型继承自 PreTrainedModel。请查阅父类文档,了解该库为所有模型实现的通用方法(例如下载或保存、调整输入嵌入大小、剪枝头等)。

该模型也是 PyTorch torch.nn.Module 的子类。 请将其作为常规的 PyTorch Module 使用,有关通用用法和行为的所有事项,请参阅 PyTorch 文档。

forward

< >

( input_ids: typing.Optional[torch.LongTensor] = Noneattention_mask: typing.Optional[torch.Tensor] = Noneposition_ids: typing.Optional[torch.LongTensor] = Nonedecoder_input_ids: typing.Optional[torch.LongTensor] = Nonedecoder_attention_mask: typing.Optional[torch.Tensor] = Nonedecoder_position_ids: typing.Optional[torch.LongTensor] = Noneencoder_outputs: transformers.modeling_outputs.BaseModelOutput | None = Noneinputs_embeds: typing.Optional[torch.Tensor] = Nonedecoder_inputs_embeds: typing.Optional[torch.Tensor] = Nonelabels: typing.Optional[torch.LongTensor] = None**kwargs: *<class 'transformers.utils.generic.TransformersKwargs'> ) → TokenClassifierOutput or tuple(torch.Tensor)

参数

  • input_ids (torch.LongTensor 形状为 (batch_size, sequence_length), ,可选) — 输入序列中各 token 在词汇表中的索引。默认情况下,填充(padding)将被忽略。

    可以使用 AutoTokenizer 获取索引。有关详细信息,请参阅 PreTrainedTokenizer.encode() 和 PreTrainedTokenizer.call()。

    什么是输入 ID?

  • attention_mask (torch.Tensor 形状为 (batch_size, sequence_length), ,可选) — 用于避免对填充 token 索引执行注意力机制的掩码。掩码值取自 [0, 1]:
    • 1 表示未被屏蔽的 token,
    • 对于被掩码的 token,值为 0。

    什么是注意力掩码?

  • position_ids (torch.LongTensor 形状为 (batch_size, sequence_length), 可选) — 输入序列中每个 token 在位置嵌入中的位置索引。取值范围为 [0, config.n_positions - 1].

    什么是位置 ID?

  • decoder_input_ids (torch.LongTensor 形状为 (batch_size, target_sequence_length), ,可选) — 解码器输入序列中各 token 在词汇表中的索引。可通过 AutoTokenizer 获取索引。详情请参阅 PreTrainedTokenizer.encode() 和 PreTrainedTokenizer.call()。什么是解码器输入 ID?

    可以使用 AutoTokenizer 获取索引。有关详细信息,请参阅 PreTrainedTokenizer.encode() 和 PreTrainedTokenizer.call()。

    什么是解码器输入 ID?

  • decoder_attention_mask(形状为 (batch_size, target_sequence_length) 的 torch.Tensor,可选) — 用于避免对特定 token 索引执行注意力机制的掩码。默认情况下,会使用因果掩码,以确保模型只能查看之前的输入来预测未来。
  • decoder_position_ids (torch.LongTensor,形状为 (batch_size, decoder_sequence_length),可选) — 解码器输入序列中每个 token 在位置嵌入中的位置索引。取值范围为 [0, config.decoder.n_positions - 1]。什么是位置 ID?
  • encoder_outputs (~modeling_outputs.BaseModelOutput, 可选) — 元组由 (last_hidden_state, 可选: hidden_states, 可选: attentions) last_hidden_state 的形状为 (batch_size, sequence_length, hidden_size), 可选),是编码器最后一层输出的隐藏状态序列。用于解码器的交叉注意力。
  • inputs_embeds(形状为 (batch_size, sequence_length, hidden_size) 的 torch.Tensor,可选) — 可选地,您可以选择直接传入嵌入表示,而不是传入 input_ids。如果您希望对 input_ids 索引如何转换为相应向量拥有比模型内部嵌入查找矩阵更多的控制权,这将非常有用。
  • decoder_inputs_embeds (torch.Tensor 形状为 (batch_size, target_sequence_length, hidden_size), ,可选) — 可选地,除了传递 decoder_input_ids ,还可以选择直接传递嵌入表示。如果使用了 past_key_values ,则只需输入最后一个 decoder_inputs_embeds (参见 past_key_values). 。这在希望比模型内部的嵌入查找矩阵更精细地控制如何将 decoder_input_ids 索引转换为对应向量时非常有用。

    如果 decoder_input_ids 和 decoder_inputs_embeds 均未设置,decoder_inputs_embeds 将取 inputs_embeds 的值。

  • labels (torch.LongTensor,形状为 (batch_size,),可选) — 用于计算序列分类/回归损失的标签。索引应在 [0, ..., config.num_labels - 1] 范围内。如果 config.num_labels == 1,则计算回归损失(均方损失);如果 config.num_labels > 1,则计算分类损失(交叉熵)。

返回值

TokenClassifierOutput 或 tuple(torch.Tensor)

一个 TokenClassifierOutput 或 torch.Tensor 元组(如果传入了 return_dict=False 或 config.return_dict=False),其包含的元素取决于配置(T5GemmaConfig)和输入。

T5GemmaForTokenClassification 的 forward 方法重写了 __call__ 特殊方法。

尽管前向传播的实现需要在此函数中定义,但之后应调用 Module 实例,而不是直接调用此函数,因为前者会负责执行预处理和后处理步骤,而后者则会静默地忽略这些步骤。

  • loss(形状为 (1,) 的 torch.Tensor,可选,在提供 labels 时返回) — 分类损失。

  • logits(形状为 (batch_size, sequence_length, config.num_labels) 的 torch.Tensor)— 分类分数(SoftMax 之前)。

  • hidden_states (tuple(torch.Tensor),可选,当传入 output_hidden_states=True 或 config.output_hidden_states=True 时返回) — 一个 torch.Tensor 元组(如果模型具有嵌入层,则包含嵌入层的输出,以及每一层的输出),形状为 (batch_size, sequence_length, hidden_size)。

    模型在每一层输出处的隐藏状态,以及可选的初始嵌入输出。

  • attentions (tuple(torch.Tensor),可选,当传入 output_attentions=True 或 config.output_attentions=True 时返回) — 一个 torch.Tensor 元组(每层一个),形状为 (batch_size, num_heads, sequence_length, sequence_length)。

    注意力权重,即经过注意力 softmax 之后的权重,用于在自注意力头中计算加权平均值。

示例:

>>> from transformers import AutoTokenizer, T5GemmaForTokenClassification
>>> import torch

>>> tokenizer = AutoTokenizer.from_pretrained("google/t5_gemma_module-7b")
>>> model = T5GemmaForTokenClassification.from_pretrained("google/t5_gemma_module-7b")

>>> inputs = tokenizer(
...     "HuggingFace is a company based in Paris and New York", add_special_tokens=False, return_tensors="pt"
... )

>>> with torch.no_grad():
...     logits = model(**inputs).logits

>>> predicted_token_class_ids = logits.argmax(-1)

>>> # Note that tokens are classified rather then input words which means that
>>> # there might be more predicted token classes than words.
>>> # Multiple token classes might account for the same word
>>> predicted_tokens_classes = [model.config.id2label[t.item()] for t in predicted_token_class_ids[0]]
>>> predicted_tokens_classes
...

>>> labels = predicted_token_class_ids
>>> loss = model(**inputs, labels=labels).loss
>>> round(loss.item(), 2)
...
Update on GitHub