Struct rust_bert::openai_gpt::OpenAIGPTLMHeadModel

pub struct OpenAIGPTLMHeadModel { /* fields omitted */ }

GPT Language Modeling head

GPT model with a decoding head (linear layer without bias). The weights of the linear layer are tied to the word embeddings It is made of the following blocks:

• transformer: Base Gpt2Model
• lm_head: Linear layer without bias tied to the weights of the token id embeddings

Implementations

impl OpenAIGPTLMHeadModel

pub fn new(p: &Path, config: &Gpt2Config) -> OpenAIGPTLMHeadModel

Build a new OpenAIGPTLMHeadModel

Arguments

• p - Variable store path for the root of the GPT model
• config - Gpt2Config object defining the model architecture

Example

use tch::{nn, Device};
use rust_bert::Config;
use std::path::Path;
use rust_bert::gpt2::Gpt2Config;
use rust_bert::openai_gpt::OpenAIGPTLMHeadModel;

let config_path = Path::new("path/to/config.json");
let device = Device::Cpu;
let p = nn::VarStore::new(device);
let config = Gpt2Config::from_file(config_path);
let gpt2: OpenAIGPTLMHeadModel = OpenAIGPTLMHeadModel::new(&(&p.root() / "gpt"), &config);

Trait Implementations

impl LMHeadModel for OpenAIGPTLMHeadModel

fn forward_t(
    &mut self,
    input_ids: &Option<Tensor>,
    _layer_past: &Option<Vec<Tensor>>,
    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<(Tensor, Option<Tensor>, Option<Vec<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)
• _layer_past - Unused for GPT
• 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.
• _encoder_outputs - Unused for GPT
• _decoder_input_ids - Unused for GPT
• 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
• encoder_hidden_states - None
• past - None
• 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::Gpt2Config;
use rust_bert::openai_gpt::OpenAIGPTLMHeadModel;
use rust_bert::pipelines::generation::LMHeadModel;
 let (batch_size, sequence_length, past_sequence_length) = (64, 128, 56);
 let input_tensor = Tensor::rand(&[batch_size, sequence_length], (Int64, 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 (output, _, _, hidden_states, attentions) = no_grad(|| {
   gpt_model
        .forward_t(&Some(input_tensor),
                   &None,
                   &Some(attention_mask),
                   &Some(token_type_ids),
                   &Some(position_ids),
                   &None,
                   None,
                   &None,
                   false).unwrap()
   });

impl LanguageGenerator<OpenAIGPTLMHeadModel, OpenAiGptVocab, OpenAiGptTokenizer> for OpenAIGenerator

fn generate(
    &mut self,
    prompt_texts: Option<Vec<&str>>,
    attention_mask: Option<Tensor>
) -> Vec<String>

Generate text based on a vector of promp texts.