Trait rust_bert::pipelines::generation_utils::LMHeadModel

source · [−]
pub trait LMHeadModel {
    fn forward_t(
        &self, 
        input_ids: Option<&Tensor>, 
        layer_past: Cache, 
        attention_mask: Option<&Tensor>, 
        token_type_ids: Option<&Tensor>, 
        position_ids: Option<&Tensor>, 
        input_embeds: Option<&Tensor>, 
        encoder_outputs: Option<&Tensor>, 
        decoder_input_ids: Option<&Tensor>, 
        train: bool
    ) -> Result<LMModelOutput, RustBertError>;
}
Expand description

Language Model trait

Shared trait between language generation models (e.g. GPT2, GPT, BART) used in language generation pipelines.

Required Methods

Forward pass through the model. Example provided for GPT2.

Arguments
  • input_ids - Optional input tensor of shape (batch size, sequence_length). If None, pre-computed embeddings must be provided (see input_embeds)
  • layer_past - Optional vector of size n_layer containing the past keys and values of each layer of shape (2, batch size, number of heads, past_sequence_length, hidden size per head). When provided, these are concatenated with the current input keys and values.
  • attention_mask - Optional mask of shape (batch size, sequence_length). Masked position have value 0, non-masked value 1. If None set to 1
  • 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)
  • token_type_ids - Optional token type ids used to indicate the portion of the input the token belongs to. If not None, token type embeddings will be added to the token and position embeddings.
  • position_ids - Optional position ids of shape (batch size, sequence_length). If None, will be incremented starting from the length of the past input.
  • 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, vocab_size) representing the logits for each vocab item and position
  • past - Option<Vec<Tensor>> of length n_layer containing the past keys and values of each layer of shape (2, batch size, number of heads, past_sequence_length, hidden size per head)
  • 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
use rust_bert::gpt2::{GPT2LMHeadModel, Gpt2Config};
use rust_bert::pipelines::generation_utils::{Cache, LMHeadModel};
let (batch_size, sequence_length, past_sequence_length) = (64, 128, 56);
let input_tensor = Tensor::rand(&[batch_size, sequence_length], (Int64, device));
let mut past: Vec<Tensor> = Vec::with_capacity(config.n_layer as usize);
for _ in 0..config.n_layer as usize {
    past.push(Tensor::rand(
        &[
            2,
            batch_size,
            config.n_head,
            past_sequence_length,
            config.n_embd / config.n_head,
        ],
        (Double, device),
    ))
}
let attention_mask = Tensor::zeros(&[batch_size, sequence_length], (Int64, device));
let token_type_ids = Tensor::ones(&[batch_size, sequence_length], (Int64, device));
let position_ids = Tensor::arange(sequence_length, (Int64, device))
    .expand(&[batch_size, sequence_length], true);

let model_output = no_grad(|| {
    gpt2_model
        .forward_t(
            Some(&input_tensor),
            Cache::GPT2Cache(Some(past)),
            Some(&attention_mask),
            Some(&token_type_ids),
            Some(&position_ids),
            None,
            None,
            None,
            false,
        )
        .unwrap()
});

Implementors