Struct GPT2LMHeadModel

pub struct GPT2LMHeadModel { /* private fields */ }

GPT2 Language Modeling head

GPT2 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

Implementations

impl GPT2LMHeadModel

pub fn new<'p, P>(p: P, config: &Gpt2Config) -> GPT2LMHeadModel

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

Build a new GPT2LMHeadModel

Arguments

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

Example

use rust_bert::gpt2::{GPT2LMHeadModel, Gpt2Config};
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 = Gpt2Config::from_file(config_path);
let gpt2: GPT2LMHeadModel = GPT2LMHeadModel::new(&p.root() / "gpt2", &config);

pub fn forward_t( &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>, train: bool, ) -> Result<LMModelOutput, RustBertError>