Struct rust_bert::bart::BartModel

pub struct BartModel { /* fields omitted */ }

BART Base model

Base architecture for BART model. Usually complemented with a task-specific head, such as a language model head. It is made of the following blocks:

• encoder: BartEncoder (transformer) made of a vector of encoding layers
• decoder: BartDecoder (transformer) made of a vector of decoding layers with self attention and encoder cross-attention. caching is implemented for the decoder to avoid recalculating static states (encoder key/values and previously calculated decoder key/values)
• generation_mode: flag indicating if the model should run in generation mode (a decoder start token must then be provided)
• pad_token_id: padding token id

Methods

impl BartModel

pub fn new(p: &Path, config: &BartConfig, generation_mode: bool) -> BartModel

Build a new BartModel

Arguments

• p - Variable store path for the root of the BART model
• config - BartConfig object defining the model architecture
• generation_mode - flag indicating if the model should run in generation mode (a decoder start token must then be provided)

Example

use tch::{nn, Device};
use rust_bert::Config;
use std::path::Path;
use rust_bert::bart::{BartConfig, BartModel};

let config_path = Path::new("path/to/config.json");
let device = Device::Cpu;
let p = nn::VarStore::new(device);
let config = BartConfig::from_file(config_path);
let generation_mode = true;
let bart: BartModel = BartModel::new(&(&p.root() / "bart"), &config, generation_mode);

pub fn forward_t(
    &mut self,
    input_ids: Option<&Tensor>,
    attention_mask: Option<&Tensor>,
    decoder_input_ids: Option<&Tensor>,
    encoder_outputs: Option<(Tensor, Option<Vec<Tensor>>, Option<Vec<Tensor>>)>,
    decoder_attention_mask: Option<&Tensor>,
    train: bool
) -> (Tensor, Tensor, (Option<Tensor>, Option<Vec<(&LayerState, &LayerState)>>), Option<Vec<Tensor>>, Option<Vec<Tensor>>, Option<Vec<Tensor>>, Option<Vec<Tensor>>)

Forward pass through the model

Arguments

• input_ids - Optional input tensor of shape (batch size, source_sequence_length). Must be provided when not running in generation mode
• attention_mask - Optional attention mask of shape (batch size, source_sequence_length) for the encoder positions. Positions with a mask with value 0 will be masked.
• decoder_input_ids - Optional input tensor of shape (batch size, target_sequence_length). Must be provided when running in generation mode (e.g. initialiazed with a BOS token)
• encoder_outputs - Optional tuple made of a tensor of shape (batch size, source_sequence_length, encoder_hidden_dim) and optional vectors of tensors of length num_encoder_layers with shape (batch size, source_sequence_length, hidden_size). These correspond to the encoder last hidden state and optional hidden states/attention weights for encoder layers. When provided, the encoder hidden state will not be recalculated. Useful for generation tasks.
• decoder_attention_mask - Optional attention mask of shape (batch size, target_sequence_length) for the decoder positions. Positions with a mask 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

• decoder_output - Tensor of shape (batch size, target_sequence_length, hidden_size) representing the activations of the last decoder hidden state
• encoder_hidden_states - Tensor of shape (batch size, source_sequence_length, hidden_size) representing the activations of the last encoder hidden state
• decoder_cache - (Option<Tensor>, Option<Vec<&LayerState, &LayerState>>) of length n_layer containing the encoder padding mask and past keys and values for both the self attentiona nd teh encoder cross attention of each layer of the decoder. Note that the decoder is saving its cached states internally within each attention layer, and does not require the cached states to be passed as an input again.
• all_encoder_hidden_states - Option<Vec<Tensor>> of length num_encoder_layers with shape (batch size, source_sequence_length, hidden_size)
• all_encoder_attentions - Option<Vec<Tensor>> of length num_encoder_layers with shape (batch size, source_sequence_length, hidden_size)
• all_decoder_hidden_states - Option<Vec<Tensor>> of length num_decoder_layers with shape (batch size, target_sequence_length, hidden_size)
• all_decoder_attentions - Option<Vec<Tensor>> of length num_decoder_layers with shape (batch size, target_sequence_length, hidden_size)

Example

use rust_bert::bart::{BartConfig, BartModel};
 let (batch_size, source_sequence_length, target_sequence_length) = (64, 128, 56);
 let input_tensor = Tensor::rand(&[batch_size, source_sequence_length], (Int64, device));
 let target_tensor = Tensor::rand(&[batch_size, target_sequence_length], (Int64, device));
 let encoder_attention_mask = Tensor::ones(&[batch_size, source_sequence_length], (Int64, device));
 let decoder_attention_mask = Tensor::ones(&[batch_size, source_sequence_length], (Int64, device));

 let (decoder_output, encoder_hidden_states, decoder_cache,
      all_encoder_hidden_states, all_encoder_attentions,
      all_decoder_hidden_states, all_decoder_attentions) = no_grad(|| {
   bart_model
        .forward_t(Some(&input_tensor),
                   Some(&encoder_attention_mask),
                   Some(&target_tensor),
                   None,
                   Some(&decoder_attention_mask),
                   false)
   });

pub fn reset_cache(&mut self)

Resets the decoder cached keys and values. Should be run for every new generation using the model.