pub struct M2M100ForConditionalGeneration { /* private fields */ }
Expand description
§M2M100 Model for conditional generation
M2M100 model with a vocabulary decoding head It is made of the following blocks:
base_model
:M2M100Model
Base M2M100 modellinear
: Linear layer without bias tied to the weights of the token id embeddings
Implementations§
Source§impl M2M100ForConditionalGeneration
impl M2M100ForConditionalGeneration
Sourcepub fn new<'p, P>(p: P, config: &M2M100Config) -> M2M100ForConditionalGeneration
pub fn new<'p, P>(p: P, config: &M2M100Config) -> M2M100ForConditionalGeneration
Build a new M2M100ForConditionalGeneration
§Arguments
p
- Variable store path for the root of the M2M100 modelconfig
-M2M100Config
object defining the model architecture
§Example
use rust_bert::m2m_100::{M2M100Config, M2M100ForConditionalGeneration};
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 = M2M100Config::from_file(config_path);
let m2m100: M2M100ForConditionalGeneration =
M2M100ForConditionalGeneration::new(&p.root(), &config);
Sourcepub fn forward_t(
&self,
input_ids: Option<&Tensor>,
attention_mask: Option<&Tensor>,
encoder_output: Option<&Tensor>,
decoder_input_ids: Option<&Tensor>,
decoder_attention_mask: Option<&Tensor>,
old_layer_states: Option<Vec<(Option<LayerState>, Option<LayerState>)>>,
train: bool,
) -> MBartModelOutput
pub fn forward_t( &self, input_ids: Option<&Tensor>, attention_mask: Option<&Tensor>, encoder_output: Option<&Tensor>, decoder_input_ids: Option<&Tensor>, decoder_attention_mask: Option<&Tensor>, old_layer_states: Option<Vec<(Option<LayerState>, Option<LayerState>)>>, train: bool, ) -> MBartModelOutput
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 modeattention_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.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_input_ids
- Optional input tensor of shape (batch size, target_sequence_length). Must be provided when running in generation mode (e.g. initialized with a BOS token)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
M2M100ModelOutput
containing:decoder_output
-Tensor
of shape (batch size, target_sequence_length, vocab_size) representing the logits for each vocabulary item and positionencoder_hidden_states
-Tensor
of shape (batch size, source_sequence_length, hidden_size) representing the activations of the last encoder hidden statecache
-(Option<Tensor>, Option<Vec<&LayerState, &LayerState>>)
of length n_layer containing the encoder padding mask and past keys and values for both the self attention and the encoder cross attention of each layer of the decoder.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
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 model_output = no_grad(|| {
m2m100_model
.forward_t(Some(&input_tensor),
Some(&encoder_attention_mask),
None,
Some(&target_tensor),
Some(&decoder_attention_mask),
None,
false)
});
pub fn encode( &self, input_ids: &Tensor, attention_mask: Option<&Tensor>, ) -> Tensor
Auto Trait Implementations§
impl !Freeze for M2M100ForConditionalGeneration
impl RefUnwindSafe for M2M100ForConditionalGeneration
impl Send for M2M100ForConditionalGeneration
impl !Sync for M2M100ForConditionalGeneration
impl Unpin for M2M100ForConditionalGeneration
impl UnwindSafe for M2M100ForConditionalGeneration
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