Struct BertModel

pub struct BertModel<T: BertEmbedding> { /* private fields */ }

BERT Base model

Base architecture for BERT models. Task-specific models will be built from this common base model It is made of the following blocks:

• embeddings: token, position and segment_id embeddings
• encoder: Encoder (transformer) made of a vector of layers. Each layer is made of a self-attention layer, an intermediate (linear) and output (linear + layer norm) layers
• pooler: linear layer applied to the first element of the sequence (MASK token)
• is_decoder: Flag indicating if the model is used as a decoder. If set to true, a causal mask will be applied to hide future positions that should not be attended to.

Implementations

impl<T: BertEmbedding> BertModel<T>

Defines the implementation of the BertModel. The BERT model shares many similarities with RoBERTa, main difference being the embeddings. Therefore the forward pass of the model is shared and the type of embedding used is abstracted away. This allows to create BertModel<RobertaEmbeddings> or BertModel<BertEmbeddings> for each model type.

pub fn new<'p, P>(p: P, config: &BertConfig) -> BertModel<T>

where P: Borrow<Path<'p>>,

Build a new BertModel

Arguments

• p - Variable store path for the root of the BERT model
• config - BertConfig object defining the model architecture and decoder status

Example

use rust_bert::bert::{BertConfig, BertEmbeddings, BertModel};
use rust_bert::Config;
use std::path::Path;
use tch::{nn, Device};

let config_path = Path::new("path/to/config.json");
let device = Device::Cpu;
let p = nn::VarStore::new(device);
let config = BertConfig::from_file(config_path);
let bert: BertModel<BertEmbeddings> = BertModel::new(&p.root() / "bert", &config);

pub fn new_with_optional_pooler<'p, P>( p: P, config: &BertConfig, add_pooling_layer: bool, ) -> BertModel<T>

where P: Borrow<Path<'p>>,

Build a new BertModel with an optional Pooling layer

Arguments

• p - Variable store path for the root of the BERT model
• config - BertConfig object defining the model architecture and decoder status
• add_pooling_layer - Enable/Disable an optional pooling layer at the end of the model

Example

use rust_bert::bert::{BertConfig, BertEmbeddings, BertModel};
use rust_bert::Config;
use std::path::Path;
use tch::{nn, Device};

let config_path = Path::new("path/to/config.json");
let device = Device::Cpu;
let p = nn::VarStore::new(device);
let config = BertConfig::from_file(config_path);
let bert: BertModel<BertEmbeddings> =
    BertModel::new_with_optional_pooler(&p.root() / "bert", &config, false);

pub fn forward_t( &self, input_ids: Option<&Tensor>, mask: Option<&Tensor>, token_type_ids: Option<&Tensor>, position_ids: Option<&Tensor>, input_embeds: Option<&Tensor>, encoder_hidden_states: Option<&Tensor>, encoder_mask: Option<&Tensor>, train: bool, ) -> Result<BertModelOutput, RustBertError>

Forward pass through the model

Arguments

• input_ids - Optional input tensor of shape (batch size, sequence_length). If None, pre-computed embeddings must be provided (see input_embeds)
• mask - Optional mask of shape (batch size, sequence_length). Masked position have value 0, non-masked value 1. If None set to 1
• token_type_ids - Optional segment id of shape (batch size, sequence_length). Convention is value of 0 for the first sentence (incl. SEP) and 1 for the second sentence. If None set to 0.
• position_ids - Optional position ids of shape (batch size, sequence_length). If None, will be incremented from 0.
• input_embeds - Optional pre-computed input embeddings of shape (batch size, sequence_length, hidden_size). If None, input ids must be provided (see input_ids)
• encoder_hidden_states - Optional encoder hidden state of shape (batch size, encoder_sequence_length, hidden_size). If the model is defined as a decoder and the encoder_hidden_states is not None, used in the cross-attention layer as keys and values (query from the decoder).
• encoder_mask - Optional encoder attention mask of shape (batch size, encoder_sequence_length). If the model is defined as a decoder and the encoder_hidden_states is not None, used to mask encoder values. Positions with value 0 will be masked.
• train - boolean flag to turn on/off the dropout layers in the model. Should be set to false for inference.

Returns

• BertOutput containing:
  • hidden_state - Tensor of shape (batch size, sequence_length, hidden_size)
  • pooled_output - Tensor of shape (batch size, hidden_size)
  • all_hidden_states - Option<Vec<Tensor>> of length num_hidden_layers with shape (batch size, sequence_length, hidden_size)
  • all_attentions - Option<Vec<Tensor>> of length num_hidden_layers with shape (batch size, sequence_length, hidden_size)

Example

let (batch_size, sequence_length) = (64, 128);
let input_tensor = Tensor::rand(&[batch_size, sequence_length], (Kind::Int64, device));
let mask = Tensor::zeros(&[batch_size, sequence_length], (Kind::Int64, device));
let token_type_ids = Tensor::zeros(&[batch_size, sequence_length], (Kind::Int64, device));
let position_ids = Tensor::arange(sequence_length, (Kind::Int64, device))
    .expand(&[batch_size, sequence_length], true);

let model_output = no_grad(|| {
    bert_model
        .forward_t(
            Some(&input_tensor),
            Some(&mask),
            Some(&token_type_ids),
            Some(&position_ids),
            None,
            None,
            None,
            false,
        )
        .unwrap()
});