Transformers documentation

GTE

Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

该模型于 2024-07-29 在 HF papers 上发布,并于 2026-09-30 贡献至 Hugging Face Transformers。

GTE

概述

GTE 由 Xin Zhang、Yanzhao Zhang、Dingkun Long、Wen Xie、Ziqi Dai、Jialong Tang、Huan Lin、Baosong Yang、Pengjun Xie、Fei Huang、Meishan Zhang、Wenjie Li 和 Min Zhang 在 mGTE: Generalized Long-Context Text Representation and Reranking Models for Multilingual Text Retrieval 中提出。

该论文的摘要如下:

我们介绍了从零开始构建长上下文多语言文本表示模型(TRM)和重排序器以用于文本检索的系统性工作。首先,我们引入了一个基于 RoPE 和无填充(unpadding)技术增强的文本编码器(基础规模),并在原生的 8192 个 token 上下文中进行预训练(长于以往多语言编码器的 512 个 token)。随后,我们通过对比学习构建了一个混合 TRM 和一个交叉编码器重排序器。评估结果表明,我们的文本编码器在性能上超越了同等规模的先前最先进模型 XLM-R。同时,我们的 TRM 和重排序器在性能上比肩大规模的最先进模型 BGE-M3,并在长上下文检索基准测试中取得了更好的结果。进一步的分析表明,我们提出的模型在训练和推理阶段均展现出更高的效率。我们相信,它们的效率和有效性将有助于各类研究和工业应用。

GTE 是一种 BERT 风格的双向编码器,它用 RoPE 取代了绝对位置嵌入,采用门控 MLP,并在每个残差连接之后应用层归一化。相同的架构支撑着阿里巴巴的 gte-*-v1.5、gte-multilingual-* 和 gte-en-mlm-* 检查点,以及 Snowflake 的 snowflake-arctic-embed-m-v2.0。

该模型由 Harshal Janjani 贡献。 原始代码可在此处找到。

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

以下示例演示了如何使用 Pipeline 和 AutoModel 提取特征(嵌入)。

Pipeline
AutoModel
from transformers import pipeline


pipeline = pipeline(
    task="feature-extraction",
    model="alibaba-nlp-community/gte-multilingual-base",
)
pipeline("Plants create oxygen through a process known as photosynthesis.")

备注

  • GTE 使用 RoPE,因此为了获得正确的位置编码,请使用右侧填充(默认设置),或者使用左侧填充并相应地准备 position_ids。
  • type_vocab_size 在不同检查点中有所不同。alibaba-nlp-community/gte-base-en-v1.5 将其设置为 0,此时不会创建 token 类型嵌入,并且会忽略 token_type_ids。
  • gte-*-v1.5 和 gte-multilingual-* 检查点在 RoPE 之上应用了静态 NTK 缩放。它表示为一个 linear RopeParameters 条目,其 rope_theta 为经过 NTK 因子缩放后的基础值。

GteConfig

class transformers.GteConfig

< >

( 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 = 250048hidden_size: int = 768num_hidden_layers: int = 12num_attention_heads: int = 12intermediate_size: int = 3072hidden_act: str = 'gelu'hidden_dropout_prob: float | int = 0.1attention_probs_dropout_prob: float | int = 0.0max_position_embeddings: int = 8192type_vocab_size: int = 1initializer_range: float = 0.02layer_norm_eps: float = 1e-12pad_token_id: int | None = 1classifier_dropout: float | int | None = Nonebos_token_id: int | None = Noneeos_token_id: int | list[int] | None = Nonetie_word_embeddings: bool = Truerope_parameters: transformers.modeling_rope_utils.RopeParameters | dict | None = None )

参数

  • vocab_size (int,可选,默认为 250048) — 模型的词汇表大小。定义了 input_ids 可以表示的不同 token 数量。
  • hidden_size (int,可选,默认为 768) — 隐藏表示的维度。
  • num_hidden_layers (int, 可选, 默认为 12) — Transformer 解码器中的隐藏层数量。
  • num_attention_heads (int,可选,默认为 12) — Transformer 解码器中每个注意力层的注意力头数量。
  • intermediate_size (int,可选,默认为 3072) — MLP 表示的维度。
  • hidden_act (str,可选,默认为 gelu) — 解码器中的非线性激活函数(函数或字符串)。例如,"gelu"、 "relu"、"silu" 等。
  • hidden_dropout_prob (Union[float, int],可选,默认为 0.1) — 嵌入层、编码器和池化层中所有全连接层的 dropout 概率。
  • attention_probs_dropout_prob (Union[float, int],可选,默认为 0.0) — 注意力概率的 dropout 比率。
  • max_position_embeddings (int,可选,默认为 8192) — 该模型可能使用的最大序列长度。
  • type_vocab_size (int,可选,默认为 1) — token_type_ids 的词汇表大小。
  • initializer_range (float,可选,默认为 0.02) — 用于初始化所有权重矩阵的 truncated_normal_initializer 的标准差。
  • layer_norm_eps (float,可选,默认为 1e-12) — 层归一化层所使用的 epsilon 值。
  • pad_token_id (int,可选,默认为 1) — 词表中用于填充的 token id。
  • classifier_dropout (Union[float, int], 可选) — 分类器的 dropout 比率。
  • bos_token_id (int, 可选) — 词汇表中用于表示流起始的 Token id。
  • eos_token_id (Union[int, list[int]], 可选) — 词汇表中用于表示流结束(end-of-stream)的 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,则可选地包含用于缩放的参数。

这是用于存储 GteModel 配置的类。它根据指定的参数实例化 Gte 模型,并定义模型架构。使用默认值实例化配置,将得到与 alibaba-nlp-community/gte-multilingual-base 类似的配置。

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

示例:

>>> from transformers import GteConfig, GteModel

>>> # Initializing a GTE alibaba-nlp-community/gte-multilingual-base style configuration
>>> configuration = GteConfig()

>>> # Initializing a model (with random weights) from the alibaba-nlp-community/gte-multilingual-base style configuration
>>> model = GteModel(configuration)

>>> # Accessing the model configuration
>>> configuration = model.config

GteModel

class transformers.GteModel

< >

( config: GteConfigadd_pooling_layer = False )

参数

  • config (GteConfig) — 模型配置类,包含模型的所有参数。使用配置文件进行初始化时,不会加载与模型关联的权重,仅加载配置。请查看 from_pretrained() 方法以加载模型权重。
  • add_pooling_layer (bool,可选,默认为 False) — 是否添加池化层。

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

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

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

forward

< >

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

参数

  • input_ids (torch.Tensor 形状为 (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。

    什么是注意力掩码?

  • token_type_ids (torch.Tensor 形状为 (batch_size, sequence_length), ,可选) — 用于指示输入中第一和第二部分的片段标记索引。索引取自 [0, 1]:
    • 0 对应 sentence A 标记,
    • 1 对应于 sentence B 标记。

    什么是 token type ID?

  • position_ids (torch.Tensor 形状为 (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 索引如何转换为相应向量拥有比模型内部嵌入查找矩阵更多的控制权,这将非常有用。

返回值

BaseModelOutputWithPooling 或 tuple(torch.Tensor)

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

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

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

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

  • pooler_output (torch.Tensor,形状为 (batch_size, hidden_size)) — 序列中第一个 token(分类 token)的最后一层隐藏状态,经过用于辅助预训练任务的层进一步处理后得到。例如,对于 BERT 系列模型,这返回的是分类 token 经过线性层和 tanh 激活函数处理后的结果。线性层的权重在预训练期间通过下一句预测(分类)目标进行训练。

  • 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 之后的权重,用于在自注意力头中计算加权平均值。

GteForMaskedLM

class transformers.GteForMaskedLM

< >

( config: GteConfig )

参数

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

带有 language modeling 头的 Gte 模型。

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

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

forward

< >

( input_ids: typing.Optional[torch.LongTensor] = Noneattention_mask: typing.Optional[torch.Tensor] = Nonetoken_type_ids: typing.Optional[torch.LongTensor] = Noneposition_ids: typing.Optional[torch.LongTensor] = Noneinputs_embeds: typing.Optional[torch.Tensor] = Nonelabels: typing.Optional[torch.LongTensor] = None**kwargs: *<class 'transformers.utils.generic.TransformersKwargs'> ) → MaskedLMOutput 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。

    什么是注意力掩码?

  • token_type_ids (torch.LongTensor 形状为 (batch_size, sequence_length), ,可选) — 用于指示输入中第一和第二部分的片段 token 索引。索引取值范围为 [0, 1]:
    • 0 对应 sentence A 标记,
    • 1 对应于 sentence B 标记。

    什么是 token type ID?

  • 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 索引如何转换为相应向量拥有比模型内部嵌入查找矩阵更多的控制权,这将非常有用。
  • labels (torch.LongTensor,形状为 (batch_size, sequence_length),可选) — 用于计算掩码语言建模损失的标签。索引应在 [-100, 0, ..., config.vocab_size] 范围内(参见 input_ids 的文档字符串)。索引设置为 -100 的 token 将被忽略(即被掩码),损失仅针对标签在 [0, ..., config.vocab_size] 范围内的 token 进行计算。

返回值

MaskedLMOutput 或 tuple(torch.Tensor)

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

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

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

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

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

  • 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, GteForMaskedLM
>>> import torch

>>> tokenizer = AutoTokenizer.from_pretrained("alibaba-nlp-community/gte-multilingual-base")
>>> model = GteForMaskedLM.from_pretrained("alibaba-nlp-community/gte-multilingual-base")

>>> inputs = tokenizer("The capital of France is <mask>.", return_tensors="pt")

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

>>> # retrieve index of <mask>
>>> mask_token_index = (inputs.input_ids == tokenizer.mask_token_id)[0].nonzero(as_tuple=True)[0]

>>> predicted_token_id = logits[0, mask_token_index].argmax(axis=-1)
>>> tokenizer.decode(predicted_token_id)
...

>>> labels = tokenizer("The capital of France is Paris.", return_tensors="pt")["input_ids"]
>>> # mask labels of non-<mask> tokens
>>> labels = torch.where(inputs.input_ids == tokenizer.mask_token_id, labels, -100)

>>> outputs = model(**inputs, labels=labels)
>>> round(outputs.loss.item(), 2)
...

GteForSequenceClassification

class transformers.GteForSequenceClassification

< >

( config: GteConfig )

参数

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

Gte 模型,顶部带有序列分类/回归头(在池化输出之上添加一个线性层),例如用于 GLUE 任务。

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

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

forward

< >

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

参数

  • input_ids (torch.Tensor 形状为 (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。

    什么是注意力掩码?

  • token_type_ids (torch.Tensor 形状为 (batch_size, sequence_length), ,可选) — 用于指示输入中第一和第二部分的片段标记索引。索引取自 [0, 1]:
    • 0 对应 sentence A 标记,
    • 1 对应于 sentence B 标记。

    什么是 token type ID?

  • position_ids (torch.Tensor 形状为 (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 索引如何转换为相应向量拥有比模型内部嵌入查找矩阵更多的控制权,这将非常有用。
  • 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),其包含的元素取决于配置(GteConfig)和输入。

GteForSequenceClassification 的 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, GteForSequenceClassification

>>> tokenizer = AutoTokenizer.from_pretrained("alibaba-nlp-community/gte-multilingual-base")
>>> model = GteForSequenceClassification.from_pretrained("alibaba-nlp-community/gte-multilingual-base")

>>> 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 = GteForSequenceClassification.from_pretrained("alibaba-nlp-community/gte-multilingual-base", num_labels=num_labels)

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

多标签分类示例:

>>> import torch
>>> from transformers import AutoTokenizer, GteForSequenceClassification

>>> tokenizer = AutoTokenizer.from_pretrained("alibaba-nlp-community/gte-multilingual-base")
>>> model = GteForSequenceClassification.from_pretrained("alibaba-nlp-community/gte-multilingual-base", 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 = GteForSequenceClassification.from_pretrained(
...     "alibaba-nlp-community/gte-multilingual-base", 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

GteForTokenClassification

class transformers.GteForTokenClassification

< >

( config: GteConfig )

参数

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

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

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

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

forward

< >

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

参数

  • input_ids (torch.Tensor 形状为 (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。

    什么是注意力掩码?

  • token_type_ids (torch.Tensor 形状为 (batch_size, sequence_length), ,可选) — 用于指示输入中第一和第二部分的片段标记索引。索引取自 [0, 1]:
    • 0 对应 sentence A 标记,
    • 1 对应于 sentence B 标记。

    什么是 token type ID?

  • position_ids (torch.Tensor 形状为 (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 索引如何转换为相应向量拥有比模型内部嵌入查找矩阵更多的控制权,这将非常有用。
  • labels(形状为 (batch_size, sequence_length) 的 torch.LongTensor,可选) — 用于计算 token 分类损失的标签。索引应在 [0, ..., config.num_labels - 1] 范围内。

返回值

TokenClassifierOutput 或 tuple(torch.Tensor)

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

GteForTokenClassification 的 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, GteForTokenClassification
>>> import torch

>>> tokenizer = AutoTokenizer.from_pretrained("alibaba-nlp-community/gte-multilingual-base")
>>> model = GteForTokenClassification.from_pretrained("alibaba-nlp-community/gte-multilingual-base")

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