pub struct Gpt2Model { /* private fields */ }
Expand description
§GPT2 Base model
Base architecture for GPT2 model. Usually complemented with a task-specific head, such as a language model head. It is made of the following blocks:
wte
:token
embeddingswpe
:position
embeddingsh
: Encoder (transformer) made of a vector of layers. Each layer is made of a multi-head attention layer, layer-normalization layers and a MLP made of linear layers.output_past
: flag indicating if the model should return a past state. This can be fed back to the model to improve the quality of text generated.output_hidden_states
: flag indicating if the model should return all hidden states (as opposed to only the last layer)output_attentions
: flag indicating if the model should return activation weights
Implementations§
Source§impl Gpt2Model
impl Gpt2Model
Sourcepub fn new<'p, P>(p: P, config: &Gpt2Config) -> Gpt2Model
pub fn new<'p, P>(p: P, config: &Gpt2Config) -> Gpt2Model
Build a new Gpt2Model
§Arguments
p
- Variable store path for the root of the GPT2 modelconfig
-Gpt2Config
object defining the model architecture
§Example
use rust_bert::gpt2::{Gpt2Config, Gpt2Model};
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: Gpt2Model = Gpt2Model::new(&p.root() / "gpt2", &config);
Sourcepub 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<Gpt2ModelOutput, RustBertError>
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<Gpt2ModelOutput, RustBertError>
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 (seeinput_embeds
)layer_past
- Optional vector 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). 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 1input_embeds
- Optional pre-computed input embeddings of shape (batch size, sequence_length, hidden_size). If None, input ids must be provided (seeinput_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
Gpt2ModelOutput
containing:output
-Tensor
of shape (batch size, sequence_length, vocab_size) representing the activations of the last hidden statecache
-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)all_hidden_states
-Option<Vec<Tensor>>
of length num_hidden_layers with shape (batch size, sequence_length, hidden_size)all_attentions
-Option<Vec<Tensor>>
of length num_hidden_layers with shape (batch size, sequence_length, hidden_size)
§Example
use rust_bert::gpt2::{Gpt2Config, Gpt2Model};
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),
Some(&past),
Some(&attention_mask),
Some(&token_type_ids),
Some(&position_ids),
None,
false,
)
.unwrap()
});
Auto Trait Implementations§
impl Freeze for Gpt2Model
impl RefUnwindSafe for Gpt2Model
impl Send for Gpt2Model
impl !Sync for Gpt2Model
impl Unpin for Gpt2Model
impl UnwindSafe for Gpt2Model
Blanket Implementations§
Source§impl<T> BorrowMut<T> for Twhere
T: ?Sized,
impl<T> BorrowMut<T> for Twhere
T: ?Sized,
Source§fn borrow_mut(&mut self) -> &mut T
fn borrow_mut(&mut self) -> &mut T
Mutably borrows from an owned value. Read more
Source§impl<T> Instrument for T
impl<T> Instrument for T
Source§fn instrument(self, span: Span) -> Instrumented<Self>
fn instrument(self, span: Span) -> Instrumented<Self>
Source§fn in_current_span(self) -> Instrumented<Self>
fn in_current_span(self) -> Instrumented<Self>
Source§impl<T> IntoEither for T
impl<T> IntoEither for T
Source§fn into_either(self, into_left: bool) -> Either<Self, Self>
fn into_either(self, into_left: bool) -> Either<Self, Self>
Converts
self
into a Left
variant of Either<Self, Self>
if into_left
is true
.
Converts self
into a Right
variant of Either<Self, Self>
otherwise. Read moreSource§fn into_either_with<F>(self, into_left: F) -> Either<Self, Self>
fn into_either_with<F>(self, into_left: F) -> Either<Self, Self>
Converts
self
into a Left
variant of Either<Self, Self>
if into_left(&self)
returns true
.
Converts self
into a Right
variant of Either<Self, Self>
otherwise. Read more