AWS Trainium & Inferentia 文档
BERT
并获得增强的文档体验
开始使用
BERT
概述
BERT 是一种双向 Transformer 模型,通过在未标注文本上预训练,以预测句子中被遮蔽的词元以及预测一个句子是否跟在另一个句子之后。其主要思想是,通过随机遮蔽一些词元,模型可以同时利用左侧和右侧的文本进行训练,从而获得更全面的理解。BERT 也非常通用,因为其学习到的语言表示可以通过微调额外的层或头部来适应其他 NLP 任务。
您可以在 BERT 集合中找到所有原始 BERT 检查点。
导出到 Neuron
要在 Neuron 设备上部署 🤗 Transformers 模型,您首先需要编译模型并将其导出为序列化格式以进行推理。下面有两种编译模型的方法,您可以选择最适合您需求的一种。这里我们以 feature-extraction
为例。
选项 1:CLI
您可以使用 Optimum 命令行界面导出模型,如下所示:
optimum-cli export neuron --model google-bert/bert-base-uncased --task feature-extraction --batch_size 1 --sequence_length 128 bert_feature_extraction_neuronx/
执行 optimum-cli export neuron --help
以显示所有命令行选项及其描述。
选项 2:Python API
from optimum.neuron import NeuronModelForFeatureExtraction
input_shapes = {"batch_size": 1, "sequence_length": 128}
compiler_args = {"auto_cast": "matmul", "auto_cast_type": "bf16"}
neuron_model = NeuronModelForFeatureExtraction.from_pretrained(
"google-bert/bert-base-uncased",
export=True,
**input_shapes,
**compiler_args,
)
# Save locally
neuron_model.save_pretrained("bert_feature_extraction_neuronx")
# Upload to the HuggingFace Hub
neuron_model.push_to_hub(
"bert_feature_extraction_neuronx", repository_id="my-neuron-repo" # Replace with your HF Hub repo id
)
NeuronBertModel
class optimum.neuron.NeuronBertModel
< 源 >( model: ScriptModule config: PretrainedConfig model_save_dir: str | pathlib.Path | tempfile.TemporaryDirectory | None = None model_file_name: str | None = None preprocessors: list | None = None neuron_config: NeuronDefaultConfig | None = None **kwargs )
参数
- config (
transformers.PretrainedConfig
) — PretrainedConfig 是模型配置类,包含模型的所有参数。使用配置文件初始化不会加载与模型相关的权重,只加载配置。请查看optimum.neuron.modeling.NeuronTracedModel.from_pretrained
方法以加载模型权重。 - model (
torch.jit._script.ScriptModule
) — torch.jit._script.ScriptModule 是嵌入了由 neuron(x) 编译器编译的 NEFF(Neuron 可执行文件格式)的 TorchScript 模块。
裸 Bert 模型 Transformer 输出原始隐藏状态,顶部没有任何特定头部,用于“特征提取”任务。
此模型继承自 ~neuron.modeling.NeuronTracedModel
。请查看超类文档以了解库为其所有模型实现的通用方法(例如下载或保存)。
前向
< 源 >( input_ids: Tensor attention_mask: Tensor token_type_ids: torch.Tensor | None = None **kwargs )
参数
- input_ids (
torch.Tensor
of shape(batch_size, sequence_length)
) — 词汇表中输入序列词元的索引。索引可以通过AutoTokenizer
获得。有关详细信息,请参阅PreTrainedTokenizer.encode
和PreTrainedTokenizer.__call__
。什么是输入 ID? - attention_mask (
torch.Tensor | None
of shape(batch_size, sequence_length)
, 默认为None
) — 避免对填充词元索引执行注意力操作的掩码。掩码值选择范围为[0, 1]
:- 1 表示 未遮蔽 的词元,
- 0 表示 被遮蔽 的词元。什么是注意力掩码?
- token_type_ids (
torch.Tensor | None
of shape(batch_size, sequence_length)
, 默认为None
) — 用于指示输入的第一部分和第二部分的词元索引。索引选择范围为[0, 1]
:- 1 表示 句子 A 的词元,
- 0 表示 句子 B 的词元。什么是词元类型 ID?
NeuronBertModel
的前向方法,覆盖了 __call__
特殊方法。只接受编译步骤中跟踪的输入。在推理过程中提供的任何额外输入都将被忽略。要包含额外输入,请使用指定的输入重新编译模型。
示例
>>> from transformers import AutoTokenizer
>>> from optimum.neuron import NeuronBertModel
>>> tokenizer = AutoTokenizer.from_pretrained("optimum/bert-base-uncased-neuronx-bs1-sq128")
>>> model = NeuronBertModel.from_pretrained("optimum/bert-base-uncased-neuronx-bs1-sq128")
>>> inputs = tokenizer("Dear Evan Hansen is the winner of six Tony Awards.", return_tensors="pt")
>>> outputs = model(**inputs)
>>> last_hidden_state = outputs.last_hidden_state
>>> list(last_hidden_state.shape)
[1, 13, 384]
NeuronBertForMaskedLM
class optimum.neuron.NeuronBertForMaskedLM
< 源 >( model: ScriptModule config: PretrainedConfig model_save_dir: str | pathlib.Path | tempfile.TemporaryDirectory | None = None model_file_name: str | None = None preprocessors: list | None = None neuron_config: NeuronDefaultConfig | None = None **kwargs )
参数
- config (
transformers.PretrainedConfig
) — PretrainedConfig 是模型配置类,包含模型的所有参数。使用配置文件初始化不会加载与模型相关的权重,只加载配置。请查看optimum.neuron.modeling.NeuronTracedModel.from_pretrained
方法以加载模型权重。 - model (
torch.jit._script.ScriptModule
) — torch.jit._script.ScriptModule 是嵌入了由 neuron(x) 编译器编译的 NEFF(Neuron 可执行文件格式)的 TorchScript 模块。
带有 language modeling
头部用于 Neuron 设备上的掩码语言建模任务的掩码语言 Bert 模型。
此模型继承自 ~neuron.modeling.NeuronTracedModel
。请查看超类文档以了解库为其所有模型实现的通用方法(例如下载或保存)。
前向
< 源 >( input_ids: Tensor attention_mask: Tensor token_type_ids: torch.Tensor | None = None **kwargs )
参数
- input_ids (
torch.Tensor
of shape(batch_size, sequence_length)
) — 词汇表中输入序列词元的索引。索引可以通过AutoTokenizer
获得。有关详细信息,请参阅PreTrainedTokenizer.encode
和PreTrainedTokenizer.__call__
。什么是输入 ID? - attention_mask (
torch.Tensor | None
of shape(batch_size, sequence_length)
, 默认为None
) — 避免对填充词元索引执行注意力操作的掩码。掩码值选择范围为[0, 1]
:- 1 表示 未遮蔽 的词元,
- 0 表示 被遮蔽 的词元。什么是注意力掩码?
- token_type_ids (
torch.Tensor | None
of shape(batch_size, sequence_length)
, 默认为None
) — 用于指示输入的第一部分和第二部分的词元索引。索引选择范围为[0, 1]
:- 1 表示 句子 A 的词元,
- 0 表示 句子 B 的词元。什么是词元类型 ID?
NeuronBertForMaskedLM
的前向方法,覆盖了 __call__
特殊方法。只接受编译步骤中跟踪的输入。在推理过程中提供的任何额外输入都将被忽略。要包含额外输入,请使用指定的输入重新编译模型。
示例
>>> from transformers import AutoTokenizer
>>> from optimum.neuron import NeuronBertForMaskedLM
>>> tokenizer = AutoTokenizer.from_pretrained("optimum/legal-bert-base-uncased-neuronx")
>>> model = NeuronBertForMaskedLM.from_pretrained("optimum/legal-bert-base-uncased-neuronx")
>>> inputs = tokenizer("This [MASK] Agreement is between General Motors and John Murray.", return_tensors="pt")
>>> outputs = model(**inputs)
>>> logits = outputs.logits
>>> list(logits.shape)
[1, 13, 30522]
NeuronBertForSequenceClassification
class optimum.neuron.NeuronBertForSequenceClassification
< 源 >( model: ScriptModule config: PretrainedConfig model_save_dir: str | pathlib.Path | tempfile.TemporaryDirectory | None = None model_file_name: str | None = None preprocessors: list | None = None neuron_config: NeuronDefaultConfig | None = None **kwargs )
参数
- config (
transformers.PretrainedConfig
) — PretrainedConfig 是模型配置类,包含模型的所有参数。使用配置文件初始化不会加载与模型相关的权重,只加载配置。请查看optimum.neuron.modeling.NeuronTracedModel.from_pretrained
方法以加载模型权重。 - model (
torch.jit._script.ScriptModule
) — torch.jit._script.ScriptModule 是嵌入了由 neuron(x) 编译器编译的 NEFF(Neuron 可执行文件格式)的 TorchScript 模块。
Neuron 模型,顶部带有一个序列分类/回归头(在池化输出顶部的一个线性层),例如用于 GLUE 任务。
此模型继承自 ~neuron.modeling.NeuronTracedModel
。请查看超类文档以了解库为其所有模型实现的通用方法(例如下载或保存)。
前向
< 源 >( input_ids: Tensor attention_mask: Tensor token_type_ids: torch.Tensor | None = None **kwargs )
参数
- input_ids (
torch.Tensor
of shape(batch_size, sequence_length)
) — 词汇表中输入序列词元的索引。索引可以通过AutoTokenizer
获得。有关详细信息,请参阅PreTrainedTokenizer.encode
和PreTrainedTokenizer.__call__
。什么是输入 ID? - attention_mask (
torch.Tensor | None
of shape(batch_size, sequence_length)
, 默认为None
) — 避免对填充词元索引执行注意力操作的掩码。掩码值选择范围为[0, 1]
:- 1 表示 未遮蔽 的词元,
- 0 表示 被遮蔽 的词元。什么是注意力掩码?
- token_type_ids (
torch.Tensor | None
of shape(batch_size, sequence_length)
, 默认为None
) — 用于指示输入的第一部分和第二部分的词元索引。索引选择范围为[0, 1]
:- 1 表示 句子 A 的词元,
- 0 表示 句子 B 的词元。什么是词元类型 ID?
NeuronBertForSequenceClassification
的前向方法,覆盖了 __call__
特殊方法。只接受编译步骤中跟踪的输入。在推理过程中提供的任何额外输入都将被忽略。要包含额外输入,请使用指定的输入重新编译模型。
示例
>>> from transformers import AutoTokenizer
>>> from optimum.neuron import NeuronBertForSequenceClassification
>>> tokenizer = AutoTokenizer.from_pretrained("optimum/bert-base-multilingual-uncased-sentiment-neuronx")
>>> model = NeuronBertForSequenceClassification.from_pretrained("optimum/bert-base-multilingual-uncased-sentiment-neuronx")
>>> inputs = tokenizer("Hamilton is considered to be the best musical of human history.", return_tensors="pt")
>>> outputs = model(**inputs)
>>> logits = outputs.logits
>>> list(logits.shape)
[1, 2]
NeuronBertForTokenClassification
class optimum.neuron.NeuronBertForTokenClassification
< 源 >( model: ScriptModule config: PretrainedConfig model_save_dir: str | pathlib.Path | tempfile.TemporaryDirectory | None = None model_file_name: str | None = None preprocessors: list | None = None neuron_config: NeuronDefaultConfig | None = None **kwargs )
参数
- config (
transformers.PretrainedConfig
) — PretrainedConfig 是模型配置类,包含模型的所有参数。使用配置文件初始化不会加载与模型相关的权重,只加载配置。请查看optimum.neuron.modeling.NeuronTracedModel.from_pretrained
方法以加载模型权重。 - model (
torch.jit._script.ScriptModule
) — torch.jit._script.ScriptModule 是嵌入了由 neuron(x) 编译器编译的 NEFF(Neuron 可执行文件格式)的 TorchScript 模块。
Neuron 模型,顶部带有一个词元分类头(在隐藏状态输出顶部的一个线性层),例如用于命名实体识别 (NER) 任务。
此模型继承自 ~neuron.modeling.NeuronTracedModel
。请查看超类文档以了解库为其所有模型实现的通用方法(例如下载或保存)。
前向
< 源 >( input_ids: Tensor attention_mask: Tensor token_type_ids: torch.Tensor | None = None **kwargs )
参数
- input_ids (
torch.Tensor
of shape(batch_size, sequence_length)
) — 词汇表中输入序列词元的索引。索引可以通过AutoTokenizer
获得。有关详细信息,请参阅PreTrainedTokenizer.encode
和PreTrainedTokenizer.__call__
。什么是输入 ID? - attention_mask (
torch.Tensor | None
of shape(batch_size, sequence_length)
, 默认为None
) — 避免对填充词元索引执行注意力操作的掩码。掩码值选择范围为[0, 1]
:- 1 表示 未遮蔽 的词元,
- 0 表示 被遮蔽 的词元。什么是注意力掩码?
- token_type_ids (
torch.Tensor | None
of shape(batch_size, sequence_length)
, 默认为None
) — 用于指示输入的第一部分和第二部分的词元索引。索引选择范围为[0, 1]
:- 1 表示 句子 A 的词元,
- 0 表示 句子 B 的词元。什么是词元类型 ID?
NeuronBertForTokenClassification
的前向方法,覆盖了 __call__
特殊方法。只接受编译步骤中跟踪的输入。在推理过程中提供的任何额外输入都将被忽略。要包含额外输入,请使用指定的输入重新编译模型。
示例
>>> from transformers import AutoTokenizer
>>> from optimum.neuron import NeuronBertForTokenClassification
>>> tokenizer = AutoTokenizer.from_pretrained("optimum/bert-base-NER-neuronx")
>>> model = NeuronBertForTokenClassification.from_pretrained("optimum/bert-base-NER-neuronx")
>>> inputs = tokenizer("Lin-Manuel Miranda is an American songwriter, actor, singer, filmmaker, and playwright.", return_tensors="pt")
>>> outputs = model(**inputs)
>>> logits = outputs.logits
>>> list(logits.shape)
[1, 20, 9]
NeuronBertForQuestionAnswering
class optimum.neuron.NeuronBertForQuestionAnswering
< 源 >( model: ScriptModule config: PretrainedConfig model_save_dir: str | pathlib.Path | tempfile.TemporaryDirectory | None = None model_file_name: str | None = None preprocessors: list | None = None neuron_config: NeuronDefaultConfig | None = None **kwargs )
参数
- config (
transformers.PretrainedConfig
) — PretrainedConfig 是模型配置类,包含模型的所有参数。使用配置文件初始化不会加载与模型相关的权重,只加载配置。请查看optimum.neuron.modeling.NeuronTracedModel.from_pretrained
方法以加载模型权重。 - model (
torch.jit._script.ScriptModule
) — torch.jit._script.ScriptModule 是嵌入了由 neuron(x) 编译器编译的 NEFF(Neuron 可执行文件格式)的 TorchScript 模块。
Bert 顶部带有一个 span 分类头,用于抽取式问答任务(例如 SQuAD)(在隐藏状态输出顶部有一个线性层,用于计算 span start logits
和 span end logits
)。
此模型继承自 ~neuron.modeling.NeuronTracedModel
。请查看超类文档以了解库为其所有模型实现的通用方法(例如下载或保存)。
前向
< 源 >( input_ids: Tensor attention_mask: Tensor token_type_ids: torch.Tensor | None = None **kwargs )
参数
- input_ids (
torch.Tensor
of shape(batch_size, sequence_length)
) — 词汇表中输入序列词元的索引。索引可以通过AutoTokenizer
获得。有关详细信息,请参阅PreTrainedTokenizer.encode
和PreTrainedTokenizer.__call__
。什么是输入 ID? - attention_mask (
torch.Tensor | None
of shape(batch_size, sequence_length)
, 默认为None
) — 避免对填充词元索引执行注意力操作的掩码。掩码值选择范围为[0, 1]
:- 1 表示 未遮蔽 的词元,
- 0 表示 被遮蔽 的词元。什么是注意力掩码?
- token_type_ids (
torch.Tensor | None
of shape(batch_size, sequence_length)
, 默认为None
) — 用于指示输入的第一部分和第二部分的词元索引。索引选择范围为[0, 1]
:- 1 表示 句子 A 的词元,
- 0 表示 句子 B 的词元。什么是词元类型 ID?
NeuronBertForQuestionAnswering
的前向方法,覆盖了 __call__
特殊方法。只接受编译步骤中跟踪的输入。在推理过程中提供的任何额外输入都将被忽略。要包含额外输入,请使用指定的输入重新编译模型。
示例
>>> import torch
>>> from transformers import AutoTokenizer
>>> from optimum.neuron import NeuronBertForQuestionAnswering
>>> tokenizer = AutoTokenizer.from_pretrained("optimum/bert-base-cased-squad2-neuronx")
>>> model = NeuronBertForQuestionAnswering.from_pretrained("optimum/bert-base-cased-squad2-neuronx")
>>> question, text = "Are there wheelchair spaces in the theatres?", "Yes, we have reserved wheelchair spaces with a good view."
>>> inputs = tokenizer(question, text, return_tensors="pt")
>>> start_positions = torch.tensor([1])
>>> end_positions = torch.tensor([12])
>>> outputs = model(**inputs, start_positions=start_positions, end_positions=end_positions)
>>> start_scores = outputs.start_logits
>>> end_scores = outputs.end_logits
NeuronBertForMultipleChoice
class optimum.neuron.NeuronBertForMultipleChoice
< 源 >( model: ScriptModule config: PretrainedConfig model_save_dir: str | pathlib.Path | tempfile.TemporaryDirectory | None = None model_file_name: str | None = None preprocessors: list | None = None neuron_config: NeuronDefaultConfig | None = None **kwargs )
参数
- config (
transformers.PretrainedConfig
) — PretrainedConfig 是模型配置类,包含模型的所有参数。使用配置文件初始化不会加载与模型相关的权重,只加载配置。请查看optimum.neuron.modeling.NeuronTracedModel.from_pretrained
方法以加载模型权重。 - model (
torch.jit._script.ScriptModule
) — torch.jit._script.ScriptModule 是包含由 neuron(x) 编译器编译的 NEFF(神经元可执行文件格式)的 TorchScript 模块。
带有多项选择分类头的 Neuron 模型(在池化输出之上有一个线性层和一个 softmax),例如用于 RocStories/SWAG 任务。
此模型继承自 ~neuron.modeling.NeuronTracedModel
。请查看超类文档以了解库为其所有模型实现的通用方法(例如下载或保存)。
前向
< 来源 >( input_ids: Tensor attention_mask: Tensor token_type_ids: torch.Tensor | None = None **kwargs )
参数
- input_ids (
torch.Tensor
,形状为(batch_size, num_choices, sequence_length)
) — 词汇表中输入序列 token 的索引。索引可以通过AutoTokenizer
获取。有关详细信息,请参阅PreTrainedTokenizer.encode
和PreTrainedTokenizer.__call__
。什么是输入 ID? - attention_mask (
torch.Tensor | None
,形状为(batch_size, num_choices, sequence_length)
,默认为None
) — 掩码,用于避免在填充 token 索引上执行注意力。掩码值选择范围为[0, 1]
:- 对于未掩码的 token 为 1,
- 对于已掩码的 token 为 0。什么是注意力掩码?
- token_type_ids (
torch.Tensor | None
,形状为(batch_size, num_choices, sequence_length)
,默认为None
) — 段落 token 索引,用于指示输入的第一个和第二个部分。索引选择范围为[0, 1]
:- 对于句子 A 的 token 为 1,
- 对于句子 B 的 token 为 0。什么是 token 类型 ID?
NeuronBertForMultipleChoice
的 forward 方法覆盖了 __call__
特殊方法。它只接受在编译步骤中追踪到的输入。在推理过程中提供的任何额外输入都将被忽略。要包含额外的输入,请使用指定的输入重新编译模型。
示例
>>> from transformers import AutoTokenizer
>>> from optimum.neuron import NeuronBertForMultipleChoice
>>> tokenizer = AutoTokenizer.from_pretrained("optimum/bert-base-cased-swag-neuronx")
>>> model = NeuronBertForMultipleChoice.from_pretrained("optimum/bert-base-cased-swag-neuronx")
>>> num_choices = 4
>>> first_sentence = ["Members of the procession walk down the street holding small horn brass instruments."] * num_choices
>>> second_sentence = [
... "A drum line passes by walking down the street playing their instruments.",
... "A drum line has heard approaching them.",
... "A drum line arrives and they're outside dancing and asleep.",
... "A drum line turns the lead singer watches the performance."
... ]
>>> inputs = tokenizer(first_sentence, second_sentence, truncation=True, padding=True)
# Unflatten the inputs values expanding it to the shape [batch_size, num_choices, seq_length]
>>> for k, v in inputs.items():
... inputs[k] = [v[i: i + num_choices] for i in range(0, len(v), num_choices)]
>>> inputs = dict(inputs.convert_to_tensors(tensor_type="pt"))
>>> outputs = model(**inputs)
>>> logits = outputs.logits
>>> logits.shape
[1, 4]