[−][src]Struct rust_bert::marian::MarianForConditionalGeneration
Marian Model for conditional generation
Marian model with a vocabulary decoding head It is made of the following blocks:
base_model
:BartModel
Base BART modellinear
: Linear layer with bias tied to the weights of the token id embeddings
Implementations
impl MarianForConditionalGeneration
[src]
pub fn new(
p: &Path,
config: &BartConfig,
generation_mode: bool
) -> MarianForConditionalGeneration
[src]
p: &Path,
config: &BartConfig,
generation_mode: bool
) -> MarianForConditionalGeneration
Build a new MarianForConditionalGeneration
Arguments
p
- Variable store path for the root of the BART modelconfig
-BartConfig
object defining the model architecturegeneration_mode
- flag indicating if the model should run in generation mode (a decoder start token must then be provided)
Example
use tch::{nn, Device}; use rust_bert::Config; use std::path::Path; use rust_bert::bart::{BartConfig, BartForConditionalGeneration}; let config_path = Path::new("path/to/config.json"); let device = Device::Cpu; let p = nn::VarStore::new(device); let config = BartConfig::from_file(config_path); let generation_mode = true; let bart: BartForConditionalGeneration = BartForConditionalGeneration::new(&(&p.root() / "bart"), &config, generation_mode);
pub fn forward_t(
&mut self,
input_ids: Option<&Tensor>,
attention_mask: Option<&Tensor>,
encoder_outputs: Option<(Tensor, Option<Vec<Tensor>>, Option<Vec<Tensor>>)>,
decoder_input_ids: Option<&Tensor>,
decoder_attention_mask: Option<&Tensor>,
train: bool
) -> (Tensor, Tensor, Option<Vec<Tensor>>, Option<Vec<Tensor>>, Option<Vec<Tensor>>, Option<Vec<Tensor>>)
[src]
&mut self,
input_ids: Option<&Tensor>,
attention_mask: Option<&Tensor>,
encoder_outputs: Option<(Tensor, Option<Vec<Tensor>>, Option<Vec<Tensor>>)>,
decoder_input_ids: Option<&Tensor>,
decoder_attention_mask: Option<&Tensor>,
train: bool
) -> (Tensor, Tensor, Option<Vec<Tensor>>, Option<Vec<Tensor>>, Option<Vec<Tensor>>, Option<Vec<Tensor>>)
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. initialiazed 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
lm_logits
-Tensor
of shape (batch size, target_sequence_length, vocab_size) representing the logits for each vocab item and positionencoder_hidden_states
-Tensor
of shape (batch size, source_sequence_length, hidden_size) representing the activations of the last encoder hidden stateall_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
use rust_bert::bart::{BartConfig, BartForConditionalGeneration}; 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 (decoder_output, encoder_hidden_states, all_encoder_hidden_states, all_encoder_attentions, all_decoder_hidden_states, all_decoder_attentions) = no_grad(|| { bart_model .forward_t(Some(&input_tensor), Some(&encoder_attention_mask), None, Some(&target_tensor), Some(&decoder_attention_mask), false) });
pub fn encode(
&mut self,
input_ids: &Tensor,
attention_mask: Option<&Tensor>
) -> Tensor
[src]
&mut self,
input_ids: &Tensor,
attention_mask: Option<&Tensor>
) -> Tensor
pub fn reset_cache(&mut self)
[src]
Resets the decoder cached keys and values. Should be run for every new generation using the model.
Trait Implementations
impl LMHeadModel for MarianForConditionalGeneration
[src]
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>
[src]
&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 (seeinput_embeds
)layer_past
- Unused for BARTattention_mask
- Optional mask of shape (batch size, sequence_length). Masked position have value 0, non-masked value 1. If None set to 1input_embeds
- Unused for BARTtoken_type_ids
- Unused for BARTposition_ids
- Unused for BARTencoder_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. initialiazed with a BOS token)train
- boolean flag to turn on/off the dropout layers in the model. Should be set to false for inference.
Returns
lm_logits
-Tensor
of shape (batch size, sequence_length, vocab_size) representing the logits for each vocab item and positionpast
- Noneencoder_hidden_states
-Option<Tensor>
Hidden states for the encoderhidden_states
- Noneattentions
- None
Example
use rust_bert::gpt2::{Gpt2Config, GPT2LMHeadModel}; 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 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 (output, encoder_hidden_states, _, hidden_states, attentions) = no_grad(|| { gpt2_model .forward_t(&Some(input_tensor), &Some(past), &Some(attention_mask), &Some(token_type_ids), &Some(position_ids), &None, None, &None, false).unwrap() });
impl LanguageGenerator<MarianForConditionalGeneration, MarianVocab, MarianTokenizer> for MarianGenerator
[src]
Auto Trait Implementations
impl !RefUnwindSafe for MarianForConditionalGeneration
impl !Send for MarianForConditionalGeneration
impl !Sync for MarianForConditionalGeneration
impl Unpin for MarianForConditionalGeneration
impl !UnwindSafe for MarianForConditionalGeneration
Blanket Implementations
impl<T> Any for T where
T: 'static + ?Sized,
[src]
T: 'static + ?Sized,
impl<T> Borrow<T> for T where
T: ?Sized,
[src]
T: ?Sized,
impl<T> BorrowMut<T> for T where
T: ?Sized,
[src]
T: ?Sized,
fn borrow_mut(&mut self) -> &mut T
[src]
impl<T> From<T> for T
[src]
impl<T, U> Into<U> for T where
U: From<T>,
[src]
U: From<T>,
impl<T, U> TryFrom<U> for T where
U: Into<T>,
[src]
U: Into<T>,
type Error = Infallible
The type returned in the event of a conversion error.
fn try_from(value: U) -> Result<T, <T as TryFrom<U>>::Error>
[src]
impl<T, U> TryInto<U> for T where
U: TryFrom<T>,
[src]
U: TryFrom<T>,