Struct rust_bert::bert::BertModel

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

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.

Methods

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: &Path, config: &BertConfig) -> BertModel<T>

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::{BertModel, BertConfig, BertEmbeddings};
use tch::{nn, Device};
use rust_bert::Config;
use std::path::Path;

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 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<(Tensor, Tensor, Option<Vec<Tensor>>, Option<Vec<Tensor>>), &'static str>

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

• output - Tensor of shape (batch size, sequence_length, hidden_size)
• pooled_output - Tensor of shape (batch size, hidden_size)
• hidden_states - Option<Vec<Tensor>> of length num_hidden_layers with shape (batch size, sequence_length, hidden_size)
• 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], (Int64, device));
 let mask = Tensor::zeros(&[batch_size, sequence_length], (Int64, device));
 let token_type_ids = Tensor::zeros(&[batch_size, sequence_length], (Int64, device));
 let position_ids = Tensor::arange(sequence_length, (Int64, device)).expand(&[batch_size, sequence_length], true);

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