rust_bert/models/mbart/mbart_model.rs
1// Copyright 2021, The Facebook AI Research Team and The HuggingFace Inc. team. All rights reserved.
2// Copyright 2020 Guillaume Becquin
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6// http://www.apache.org/licenses/LICENSE-2.0
7// Unless required by applicable law or agreed to in writing, software
8// distributed under the License is distributed on an "AS IS" BASIS,
9// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
10// See the License for the specific language governing permissions and
11// limitations under the License.
12
13use crate::bart::BartModelOutput;
14use crate::common::dropout::Dropout;
15use crate::mbart::decoder::MBartDecoder;
16use crate::mbart::encoder::MBartEncoder;
17use crate::mbart::LayerState;
18use crate::pipelines::common::{ModelType, TokenizerOption};
19use crate::pipelines::generation_utils::private_generation_utils::{
20 PreparedInput, PrivateLanguageGenerator,
21};
22use crate::pipelines::generation_utils::{Cache, GenerateConfig, LMModelOutput, LanguageGenerator};
23use crate::pipelines::translation::Language;
24use crate::{Activation, Config, RustBertError};
25use serde::{Deserialize, Serialize};
26use std::borrow::Borrow;
27use std::collections::HashMap;
28use tch::kind::Kind::Int64;
29use tch::nn::{embedding, EmbeddingConfig, Init};
30use tch::{nn, Device, Tensor};
31
32/// # MBART Pretrained model weight files
33pub struct MBartModelResources;
34
35/// # MBART Pretrained model config files
36pub struct MBartConfigResources;
37
38/// # MBART Pretrained model vocab files
39pub struct MBartVocabResources;
40
41/// # MBART source languages pre-sets
42pub struct MBartSourceLanguages;
43
44/// # MBART target languages pre-sets
45pub type MBartTargetLanguages = MBartSourceLanguages;
46
47impl MBartModelResources {
48 /// Shared under MIT license by the Facebook AI Research Fairseq team at <https://github.com/pytorch/fairseq>. Modified with conversion to C-array format.
49 pub const MBART50_MANY_TO_MANY: (&'static str, &'static str) = (
50 "mbart-50-many-to-many-mmt/model",
51 "https://huggingface.co/facebook/mbart-large-50-many-to-many-mmt/resolve/main/rust_model.ot",
52 );
53}
54
55impl MBartConfigResources {
56 /// Shared under MIT license by the Facebook AI Research Fairseq team at <https://github.com/pytorch/fairseq>. Modified with conversion to C-array format.
57 pub const MBART50_MANY_TO_MANY: (&'static str, &'static str) = (
58 "mbart-50-many-to-many-mmt/config",
59 "https://huggingface.co/facebook/mbart-large-50-many-to-many-mmt/resolve/main/config.json",
60 );
61}
62
63impl MBartVocabResources {
64 /// Shared under MIT license by the Facebook AI Research Fairseq team at <https://github.com/pytorch/fairseq>. Modified with conversion to C-array format.
65 pub const MBART50_MANY_TO_MANY: (&'static str, &'static str) = (
66 "mbart-50-many-to-many-mmt/vocab",
67 "https://huggingface.co/facebook/mbart-large-50-many-to-many-mmt/resolve/main/sentencepiece.bpe.model",
68 );
69}
70
71#[rustfmt::skip]
72impl MBartSourceLanguages {
73 pub const MBART50_MANY_TO_MANY: [Language; 51] = [Language::Arabic, Language::Czech, Language::German, Language::English, Language::Spanish, Language::Estonian, Language::Finnish, Language::French, Language::Gujarati, Language::Hindi, Language::Italian, Language::Japanese, Language::Kazakh, Language::Korean, Language::Lithuanian, Language::Latvian, Language::Burmese, Language::Nepali, Language::Dutch, Language::Romanian, Language::Russian, Language::Sinhala, Language::Turkish, Language::Vietnamese, Language::ChineseMandarin, Language::Afrikaans, Language::Azerbaijani, Language::Bengali, Language::Farsi, Language::Hebrew, Language::Croatian, Language::Indonesian, Language::Georgian, Language::CentralKhmer, Language::Macedonian, Language::Malayalam, Language::Mongolian, Language::Marathi, Language::Polish, Language::Pashto, Language::Portuguese, Language::Swedish, Language::Swahili, Language::Tamil, Language::Thai, Language::Tagalog, Language::Ukrainian, Language::Urdu, Language::Xhosa, Language::Galician, Language::Slovenian];
74}
75
76#[derive(Debug, Serialize, Deserialize, Clone)]
77/// # MBART model configuration
78/// Defines the MBART model architecture (e.g. number of layers, hidden layer size, label mapping...)
79pub struct MBartConfig {
80 pub vocab_size: i64,
81 pub max_position_embeddings: i64,
82 pub encoder_layers: i64,
83 pub encoder_attention_heads: i64,
84 pub encoder_ffn_dim: i64,
85 pub encoder_layerdrop: f64,
86 pub decoder_layers: i64,
87 pub decoder_ffn_dim: i64,
88 pub decoder_attention_heads: i64,
89 pub decoder_layerdrop: f64,
90 pub is_encoder_decoder: Option<bool>,
91 pub activation_function: Option<Activation>,
92 pub d_model: i64,
93 pub dropout: f64,
94 pub activation_dropout: f64,
95 pub attention_dropout: f64,
96 pub classifier_dropout: Option<f64>,
97 pub scale_embedding: Option<bool>,
98 pub bos_token_id: Option<i64>,
99 pub eos_token_id: Option<i64>,
100 pub pad_token_id: Option<i64>,
101 pub forced_bos_token_id: Option<i64>,
102 pub forced_eos_token_id: Option<i64>,
103 pub decoder_start_token_id: Option<i64>,
104 pub id2label: Option<HashMap<i64, String>>,
105 pub label2id: Option<HashMap<String, i64>>,
106 pub init_std: f64,
107 pub min_length: Option<i64>,
108 pub no_repeat_ngram_size: Option<i64>,
109 pub normalize_embedding: Option<bool>,
110 pub output_attentions: Option<bool>,
111 pub output_hidden_states: Option<bool>,
112 pub output_past: Option<bool>,
113}
114
115impl Config for MBartConfig {}
116
117impl Default for MBartConfig {
118 fn default() -> Self {
119 MBartConfig {
120 vocab_size: 50265,
121 max_position_embeddings: 1024,
122 encoder_layers: 12,
123 encoder_attention_heads: 16,
124 encoder_ffn_dim: 4096,
125 encoder_layerdrop: 0.0,
126 decoder_layers: 12,
127 decoder_ffn_dim: 4096,
128 decoder_attention_heads: 16,
129 decoder_layerdrop: 0.0,
130 is_encoder_decoder: Some(true),
131 activation_function: Some(Activation::gelu),
132 d_model: 1024,
133 dropout: 0.1,
134 activation_dropout: 0.0,
135 attention_dropout: 0.0,
136 classifier_dropout: None,
137 scale_embedding: Some(false),
138 bos_token_id: Some(0),
139 eos_token_id: Some(2),
140 pad_token_id: Some(1),
141 forced_bos_token_id: None,
142 forced_eos_token_id: Some(2),
143 decoder_start_token_id: None,
144 id2label: None,
145 label2id: None,
146 init_std: 0.02,
147 min_length: None,
148 no_repeat_ngram_size: None,
149 normalize_embedding: None,
150 output_attentions: None,
151 output_hidden_states: None,
152 output_past: None,
153 }
154 }
155}
156
157fn _shift_tokens_right(input_ids: &Tensor, pad_token_id: i64) -> Tensor {
158 let output = input_ids.masked_fill(&input_ids.eq(-100), pad_token_id);
159 let index_eos: Tensor = input_ids
160 .ne(pad_token_id)
161 .sum_dim_intlist([1].as_slice(), true, Int64)
162 - 1;
163 output
164 .select(1, 0)
165 .copy_(&input_ids.gather(1, &index_eos, false).squeeze());
166 output
167 .slice(1, 1, *output.size().last().unwrap(), 1)
168 .copy_(&input_ids.slice(1, 0, *output.size().last().unwrap() - 1, 1));
169 output
170}
171
172pub struct MBartClassificationHead {
173 dense: nn::Linear,
174 dropout: Dropout,
175 out_proj: nn::Linear,
176}
177
178impl MBartClassificationHead {
179 pub fn new<'p, P>(p: P, config: &MBartConfig) -> Result<MBartClassificationHead, RustBertError>
180 where
181 P: Borrow<nn::Path<'p>>,
182 {
183 let p = p.borrow();
184
185 let dense = nn::linear(
186 p / "dense",
187 config.d_model,
188 config.d_model,
189 Default::default(),
190 );
191
192 let num_labels = config
193 .id2label
194 .as_ref()
195 .ok_or_else(|| {
196 RustBertError::InvalidConfigurationError(
197 "num_labels not provided in configuration".to_string(),
198 )
199 })?
200 .len() as i64;
201 let out_proj = nn::linear(
202 p / "out_proj",
203 config.d_model,
204 num_labels,
205 Default::default(),
206 );
207
208 let dropout = Dropout::new(config.classifier_dropout.unwrap_or(0.0));
209
210 Ok(MBartClassificationHead {
211 dense,
212 dropout,
213 out_proj,
214 })
215 }
216
217 pub fn forward_t(&self, hidden_states: &Tensor, train: bool) -> Tensor {
218 hidden_states
219 .apply_t(&self.dropout, train)
220 .apply(&self.dense)
221 .tanh()
222 .apply_t(&self.dropout, train)
223 .apply(&self.out_proj)
224 }
225}
226
227/// # MBart Base model
228/// Base architecture for MBart model. Usually complemented with a task-specific head, such as a language model head.
229/// It is made of the following blocks:
230/// - `encoder`: `MBartEncoder` (transformer) made of a vector of encoding layers
231/// - `decoder`: `MBartDecoder` (transformer) made of a vector of decoding layers with self attention and encoder cross-attention.
232/// caching is implemented for the decoder to avoid recalculating static states (encoder key/values and previously calculated decoder key/values)
233/// - `pad_token_id`: padding token id
234pub struct MBartModel {
235 pub(crate) encoder: MBartEncoder,
236 decoder: MBartDecoder,
237 pub(crate) embeddings: nn::Embedding,
238 pad_token_id: i64,
239}
240
241impl MBartModel {
242 /// Build a new `MBartModel`
243 ///
244 /// # Arguments
245 ///
246 /// * `p` - Variable store path for the root of the MBart model
247 /// * `config` - `MBartConfig` object defining the model architecture
248 ///
249 /// # Example
250 ///
251 /// ```no_run
252 /// use rust_bert::mbart::{MBartConfig, MBartModel};
253 /// use rust_bert::Config;
254 /// use std::path::Path;
255 /// use tch::{nn, Device};
256 ///
257 /// let config_path = Path::new("path/to/config.json");
258 /// let device = Device::Cpu;
259 /// let p = nn::VarStore::new(device);
260 /// let config = MBartConfig::from_file(config_path);
261 /// let mbart: MBartModel = MBartModel::new(&p.root() / "bart", &config);
262 /// ```
263 pub fn new<'p, P>(p: P, config: &MBartConfig) -> MBartModel
264 where
265 P: Borrow<nn::Path<'p>>,
266 {
267 let p = p.borrow();
268
269 let pad_token_id = config.pad_token_id.unwrap_or(1);
270 let embedding_config = EmbeddingConfig {
271 padding_idx: pad_token_id,
272 ..Default::default()
273 };
274 let embeddings: nn::Embedding = embedding(
275 p / "shared",
276 config.vocab_size,
277 config.d_model,
278 embedding_config,
279 );
280
281 let encoder = MBartEncoder::new(p / "encoder", config);
282 let decoder = MBartDecoder::new(p / "decoder", config);
283
284 MBartModel {
285 encoder,
286 decoder,
287 embeddings,
288 pad_token_id,
289 }
290 }
291
292 /// Forward pass through the model
293 ///
294 /// # Arguments
295 ///
296 /// * `input_ids` - Optional input tensor of shape (*batch size*, *source_sequence_length*). Must be provided when not running in generation mode
297 /// * `attention_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.
298 /// * `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)
299 /// * `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*).
300 /// 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.
301 /// * `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.
302 /// * `train` - boolean flag to turn on/off the dropout layers in the model. Should be set to false for inference.
303 ///
304 /// # Returns
305 ///
306 /// * `MBartModelOutput` containing:
307 /// - `decoder_output` - `Tensor` of shape (*batch size*, *target_sequence_length*, *hidden_size*) representing the activations of the last decoder hidden state
308 /// - `encoder_hidden_states` - `Option<Tensor>` of shape (*batch size*, *source_sequence_length*, *hidden_size*) representing the activations of the last encoder hidden state if it was not provided, otherwise None
309 /// - `cache` - `(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.
310 /// - `all_encoder_hidden_states` - `Option<Vec<Tensor>>` of length *num_encoder_layers* with shape (*batch size*, *source_sequence_length*, *hidden_size*)
311 /// - `all_encoder_attentions` - `Option<Vec<Tensor>>` of length *num_encoder_layers* with shape (*batch size*, *source_sequence_length*, *hidden_size*)
312 /// - `all_decoder_hidden_states` - `Option<Vec<Tensor>>` of length *num_decoder_layers* with shape (*batch size*, *target_sequence_length*, *hidden_size*)
313 /// - `all_decoder_attentions` - `Option<Vec<Tensor>>` of length *num_decoder_layers* with shape (*batch size*, *target_sequence_length*, *hidden_size*)
314 ///
315 /// # Example
316 ///
317 /// ```no_run
318 /// # use tch::{nn, Device, Tensor, no_grad};
319 /// # use rust_bert::Config;
320 /// # use std::path::Path;
321 /// # use tch::kind::Kind::{Int64, Double};
322 /// use rust_bert::mbart::{MBartConfig, MBartModel};
323 /// # let config_path = Path::new("path/to/config.json");
324 /// # let vocab_path = Path::new("path/to/vocab.txt");
325 /// # let device = Device::Cpu;
326 /// # let vs = nn::VarStore::new(device);
327 /// # let config = MBartConfig::from_file(config_path);
328 /// # let mbart_model: MBartModel = MBartModel::new(&vs.root(), &config);
329 /// let (batch_size, source_sequence_length, target_sequence_length) = (64, 128, 56);
330 /// let input_tensor = Tensor::rand(&[batch_size, source_sequence_length], (Int64, device));
331 /// let target_tensor = Tensor::rand(&[batch_size, target_sequence_length], (Int64, device));
332 /// let encoder_attention_mask =
333 /// Tensor::ones(&[batch_size, source_sequence_length], (Int64, device));
334 /// let decoder_attention_mask =
335 /// Tensor::ones(&[batch_size, source_sequence_length], (Int64, device));
336 ///
337 /// let model_output = no_grad(|| {
338 /// mbart_model.forward_t(
339 /// Some(&input_tensor),
340 /// Some(&encoder_attention_mask),
341 /// Some(&target_tensor),
342 /// None,
343 /// Some(&decoder_attention_mask),
344 /// None,
345 /// false,
346 /// )
347 /// });
348 /// ```
349 pub fn forward_t(
350 &self,
351 input_ids: Option<&Tensor>,
352 attention_mask: Option<&Tensor>,
353 decoder_input_ids: Option<&Tensor>,
354 encoder_output: Option<&Tensor>,
355 decoder_attention_mask: Option<&Tensor>,
356 layer_states: Option<Vec<(Option<LayerState>, Option<LayerState>)>>,
357 train: bool,
358 ) -> MBartModelOutput {
359 let calc_decoder_input_ids = if decoder_input_ids.is_none() {
360 Some(_shift_tokens_right(input_ids.unwrap(), self.pad_token_id))
361 } else {
362 None
363 };
364
365 let decoder_input_ids =
366 decoder_input_ids.unwrap_or_else(|| calc_decoder_input_ids.as_ref().unwrap());
367
368 let calc_encoder_output = if encoder_output.is_none() {
369 Some(self.encoder.forward_t(
370 input_ids.unwrap(),
371 attention_mask,
372 &self.embeddings,
373 train,
374 ))
375 } else {
376 None
377 };
378
379 let (calc_hidden_states, all_encoder_hidden_states, all_encoder_attentions) =
380 if let Some(calc_encoder_output) = calc_encoder_output {
381 (
382 Some(calc_encoder_output.hidden_state),
383 calc_encoder_output.all_hidden_states,
384 calc_encoder_output.all_attentions,
385 )
386 } else {
387 (None, None, None)
388 };
389
390 let encoder_output = encoder_output.unwrap_or_else(|| calc_hidden_states.as_ref().unwrap());
391
392 let decoder_output = self.decoder.forward_t(
393 decoder_input_ids,
394 encoder_output,
395 attention_mask,
396 decoder_attention_mask,
397 &self.embeddings,
398 layer_states,
399 train,
400 );
401
402 MBartModelOutput {
403 decoder_output: decoder_output.hidden_state,
404 encoder_hidden_state: calc_hidden_states,
405 cache: decoder_output.next_decoder_cache,
406 all_decoder_hidden_states: decoder_output.all_hidden_states,
407 all_decoder_attentions: decoder_output.all_attentions,
408 all_encoder_hidden_states,
409 all_encoder_attentions,
410 }
411 }
412}
413
414/// # MBart Model for conditional generation
415/// MBart model with a vocabulary decoding head
416/// It is made of the following blocks:
417/// - `base_model`: `MBartModel` Base MBart model
418/// - `linear`: Linear layer without bias tied to the weights of the token id embeddings
419pub struct MBartForConditionalGeneration {
420 base_model: MBartModel,
421 final_logits_bias: Tensor,
422}
423
424impl MBartForConditionalGeneration {
425 /// Build a new `MBartForConditionalGeneration`
426 ///
427 /// # Arguments
428 ///
429 /// * `p` - Variable store path for the root of the MBart model
430 /// * `config` - `MBartConfig` object defining the model architecture
431 ///
432 /// # Example
433 ///
434 /// ```no_run
435 /// use rust_bert::mbart::{MBartConfig, MBartForConditionalGeneration};
436 /// use rust_bert::Config;
437 /// use std::path::Path;
438 /// use tch::{nn, Device};
439 ///
440 /// let config_path = Path::new("path/to/config.json");
441 /// let device = Device::Cpu;
442 /// let p = nn::VarStore::new(device);
443 /// let config = MBartConfig::from_file(config_path);
444 /// let mbart: MBartForConditionalGeneration =
445 /// MBartForConditionalGeneration::new(&p.root(), &config);
446 /// ```
447 pub fn new<'p, P>(p: P, config: &MBartConfig) -> MBartForConditionalGeneration
448 where
449 P: Borrow<nn::Path<'p>>,
450 {
451 let p = p.borrow();
452
453 let base_model = MBartModel::new(p / "model", config);
454 let final_logits_bias = p.var(
455 "final_logits_bias",
456 &[1, config.vocab_size],
457 Init::Const(0.0),
458 );
459
460 MBartForConditionalGeneration {
461 base_model,
462 final_logits_bias,
463 }
464 }
465
466 /// Forward pass through the model
467 ///
468 /// # Arguments
469 ///
470 /// * `input_ids` - Optional input tensor of shape (*batch size*, *source_sequence_length*). Must be provided when not running in generation mode
471 /// * `attention_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.
472 /// * `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*).
473 /// 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.
474 /// * `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)
475 /// * `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.
476 /// * `train` - boolean flag to turn on/off the dropout layers in the model. Should be set to false for inference.
477 ///
478 /// # Returns
479 ///
480 /// * `MBartModelOutput` containing:
481 /// - `decoder_output` - `Tensor` of shape (*batch size*, *target_sequence_length*, *vocab_size*) representing the logits for each vocabulary item and position
482 /// - `encoder_hidden_states` - `Tensor` of shape (*batch size*, *source_sequence_length*, *hidden_size*) representing the activations of the last encoder hidden state
483 /// - `cache` - `(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.
484 /// - `all_encoder_hidden_states` - `Option<Vec<Tensor>>` of length *num_encoder_layers* with shape (*batch size*, *source_sequence_length*, *hidden_size*)
485 /// - `all_encoder_attentions` - `Option<Vec<Tensor>>` of length *num_encoder_layers* with shape (*batch size*, *source_sequence_length*, *hidden_size*)
486 /// - `all_decoder_hidden_states` - `Option<Vec<Tensor>>` of length *num_decoder_layers* with shape (*batch size*, *target_sequence_length*, *hidden_size*)
487 /// - `all_decoder_attentions` - `Option<Vec<Tensor>>` of length *num_decoder_layers* with shape (*batch size*, *target_sequence_length*, *hidden_size*)
488 ///
489 /// # Example
490 ///
491 /// ```no_run
492 /// # use tch::{nn, Device, Tensor, no_grad};
493 /// # use rust_bert::Config;
494 /// # use std::path::Path;
495 /// # use tch::kind::Kind::{Int64, Double};
496 /// use rust_bert::mbart::{MBartConfig, MBartForConditionalGeneration};
497 /// # let config_path = Path::new("path/to/config.json");
498 /// # let vocab_path = Path::new("path/to/vocab.txt");
499 /// # let device = Device::Cpu;
500 /// # let vs = nn::VarStore::new(device);
501 /// # let config = MBartConfig::from_file(config_path);
502 /// # let mbart_model: MBartForConditionalGeneration = MBartForConditionalGeneration::new(&vs.root(), &config);
503 /// let (batch_size, source_sequence_length, target_sequence_length) = (64, 128, 56);
504 /// let input_tensor = Tensor::rand(&[batch_size, source_sequence_length], (Int64, device));
505 /// let target_tensor = Tensor::rand(&[batch_size, target_sequence_length], (Int64, device));
506 /// let encoder_attention_mask = Tensor::ones(&[batch_size, source_sequence_length], (Int64, device));
507 /// let decoder_attention_mask = Tensor::ones(&[batch_size, source_sequence_length], (Int64, device));
508 ///
509 /// let model_output = no_grad(|| {
510 /// mbart_model
511 /// .forward_t(Some(&input_tensor),
512 /// Some(&encoder_attention_mask),
513 /// None,
514 /// Some(&target_tensor),
515 /// Some(&decoder_attention_mask),
516 /// None,
517 /// false)
518 /// });
519 /// ```
520 pub fn forward_t(
521 &self,
522 input_ids: Option<&Tensor>,
523 attention_mask: Option<&Tensor>,
524 encoder_output: Option<&Tensor>,
525 decoder_input_ids: Option<&Tensor>,
526 decoder_attention_mask: Option<&Tensor>,
527 old_layer_states: Option<Vec<(Option<LayerState>, Option<LayerState>)>>,
528 train: bool,
529 ) -> MBartModelOutput {
530 let base_model_output = self.base_model.forward_t(
531 input_ids,
532 attention_mask,
533 decoder_input_ids,
534 encoder_output,
535 decoder_attention_mask,
536 old_layer_states,
537 train,
538 );
539
540 let lm_logits = base_model_output
541 .decoder_output
542 .linear::<Tensor>(&self.base_model.embeddings.ws, None)
543 + &self.final_logits_bias;
544 BartModelOutput {
545 decoder_output: lm_logits,
546 ..base_model_output
547 }
548 }
549
550 pub fn encode(&self, input_ids: &Tensor, attention_mask: Option<&Tensor>) -> Tensor {
551 self.base_model
552 .encoder
553 .forward_t(
554 input_ids,
555 attention_mask,
556 &self.base_model.embeddings,
557 false,
558 )
559 .hidden_state
560 }
561}
562
563/// # MBart Model for sequence classification
564/// MBart model with a classification head
565/// It is made of the following blocks:
566/// - `base_model`: `MBartModel` Base MBart model
567/// - `classification_head`: `BartClassificationHead` made of 2 linear layers mapping hidden states to a target class
568/// - `eos_token_id`: token id for the EOS token carrying the pooled representation for classification
569pub struct MBartForSequenceClassification {
570 base_model: MBartModel,
571 classification_head: MBartClassificationHead,
572 eos_token_id: i64,
573}
574
575impl MBartForSequenceClassification {
576 /// Build a new `MBartForSequenceClassification`
577 ///
578 /// # Arguments
579 ///
580 /// * `p` - Variable store path for the root of the MBart model
581 /// * `config` - `MBartConfig` object defining the model architecture
582 ///
583 /// # Example
584 ///
585 /// ```no_run
586 /// use rust_bert::mbart::{MBartConfig, MBartForSequenceClassification};
587 /// use rust_bert::Config;
588 /// use std::path::Path;
589 /// use tch::{nn, Device};
590 ///
591 /// let config_path = Path::new("path/to/config.json");
592 /// let device = Device::Cpu;
593 /// let p = nn::VarStore::new(device);
594 /// let config = MBartConfig::from_file(config_path);
595 /// let mbart: MBartForSequenceClassification =
596 /// MBartForSequenceClassification::new(&p.root(), &config).unwrap();
597 /// ```
598 pub fn new<'p, P>(
599 p: P,
600 config: &MBartConfig,
601 ) -> Result<MBartForSequenceClassification, RustBertError>
602 where
603 P: Borrow<nn::Path<'p>>,
604 {
605 let p = p.borrow();
606
607 let base_model = MBartModel::new(p / "model", config);
608 let classification_head = MBartClassificationHead::new(p / "classification_head", config)?;
609 let eos_token_id = config.eos_token_id.unwrap_or(3);
610 Ok(MBartForSequenceClassification {
611 base_model,
612 classification_head,
613 eos_token_id,
614 })
615 }
616
617 /// Forward pass through the model
618 ///
619 /// # Arguments
620 ///
621 /// * `input_ids` - Optional input tensor of shape (*batch size*, *source_sequence_length*). Must be provided when not running in generation mode
622 /// * `attention_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.
623 /// * `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*).
624 /// 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.
625 /// * `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)
626 /// * `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.
627 /// * `train` - boolean flag to turn on/off the dropout layers in the model. Should be set to false for inference.
628 ///
629 /// # Returns
630 ///
631 /// * `MBartModelOutput` containing:
632 /// - `decoder_output` - `Tensor` of shape (*batch size*, *num_classes*) representing the activations for each class and batch item
633 /// - `encoder_hidden_states` - `Option<Tensor>` of shape (*batch size*, *source_sequence_length*, *hidden_size*) representing the activations of the last encoder hidden state if it was not provided, otherwise None.
634 /// - `cache` - `(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.
635 /// - `all_encoder_hidden_states` - `Option<Vec<Tensor>>` of length *num_encoder_layers* with shape (*batch size*, *source_sequence_length*, *hidden_size*)
636 /// - `all_encoder_attentions` - `Option<Vec<Tensor>>` of length *num_encoder_layers* with shape (*batch size*, *source_sequence_length*, *hidden_size*)
637 /// - `all_decoder_hidden_states` - `Option<Vec<Tensor>>` of length *num_decoder_layers* with shape (*batch size*, *target_sequence_length*, *hidden_size*)
638 /// - `all_decoder_attentions` - `Option<Vec<Tensor>>` of length *num_decoder_layers* with shape (*batch size*, *target_sequence_length*, *hidden_size*)
639 ///
640 /// # Example
641 ///
642 /// ```no_run
643 /// # use tch::{nn, Device, Tensor, no_grad};
644 /// # use rust_bert::Config;
645 /// # use std::path::Path;
646 /// # use tch::kind::Kind::{Int64, Double};
647 /// use rust_bert::mbart::{MBartConfig, MBartForSequenceClassification};
648 /// # let config_path = Path::new("path/to/config.json");
649 /// # let vocab_path = Path::new("path/to/vocab.txt");
650 /// # let device = Device::Cpu;
651 /// # let vs = nn::VarStore::new(device);
652 /// # let config = MBartConfig::from_file(config_path);
653 /// # let mbart_model: MBartForSequenceClassification = MBartForSequenceClassification::new(&vs.root(), &config).unwrap();
654 /// let (batch_size, source_sequence_length, target_sequence_length) = (64, 128, 56);
655 /// let input_tensor = Tensor::rand(&[batch_size, source_sequence_length], (Int64, device));
656 /// let target_tensor = Tensor::rand(&[batch_size, target_sequence_length], (Int64, device));
657 /// let encoder_attention_mask = Tensor::ones(&[batch_size, source_sequence_length], (Int64, device));
658 /// let decoder_attention_mask = Tensor::ones(&[batch_size, source_sequence_length], (Int64, device));
659 ///
660 /// let model_output = no_grad(|| {
661 /// mbart_model
662 /// .forward_t(&input_tensor,
663 /// Some(&encoder_attention_mask),
664 /// None,
665 /// Some(&target_tensor),
666 /// Some(&decoder_attention_mask),
667 /// false)
668 /// });
669 /// ```
670 pub fn forward_t(
671 &self,
672 input_ids: &Tensor,
673 attention_mask: Option<&Tensor>,
674 encoder_output: Option<&Tensor>,
675 decoder_input_ids: Option<&Tensor>,
676 decoder_attention_mask: Option<&Tensor>,
677 train: bool,
678 ) -> MBartModelOutput {
679 let base_model_output = self.base_model.forward_t(
680 Some(input_ids),
681 attention_mask,
682 decoder_input_ids,
683 encoder_output,
684 decoder_attention_mask,
685 None,
686 train,
687 );
688 let eos_mask = input_ids.eq(self.eos_token_id);
689 let reshape = eos_mask.sum_dim_intlist([1].as_slice(), true, Int64);
690 let sentence_representation = base_model_output
691 .decoder_output
692 .permute([2, 0, 1])
693 .masked_select(&eos_mask)
694 .view((-1, reshape.size()[0] * reshape.int64_value(&[0, 0])))
695 .transpose(0, 1)
696 .view((
697 base_model_output.decoder_output.size()[0],
698 -1,
699 *base_model_output.decoder_output.size().last().unwrap(),
700 ))
701 .select(1, -1);
702
703 let logits = self
704 .classification_head
705 .forward_t(&sentence_representation, train);
706 MBartModelOutput {
707 decoder_output: logits,
708 encoder_hidden_state: base_model_output.encoder_hidden_state,
709 cache: None,
710 all_decoder_hidden_states: base_model_output.all_decoder_hidden_states,
711 all_decoder_attentions: base_model_output.all_decoder_attentions,
712 all_encoder_hidden_states: base_model_output.all_encoder_hidden_states,
713 all_encoder_attentions: base_model_output.all_encoder_attentions,
714 }
715 }
716}
717
718/// Container holding a MBART model output
719pub type MBartModelOutput = BartModelOutput;
720
721/// # Language generation model based on the MBart architecture
722pub struct MBartGenerator {
723 model: MBartForConditionalGeneration,
724 tokenizer: TokenizerOption,
725 var_store: nn::VarStore,
726 generate_config: GenerateConfig,
727 bos_token_id: Option<i64>,
728 eos_token_ids: Option<Vec<i64>>,
729 forced_eos_token_id: Option<i64>,
730 pad_token_id: Option<i64>,
731 is_encoder_decoder: bool,
732 vocab_size: i64,
733 decoder_start_id: Option<i64>,
734 max_position_embeddings: i64,
735}
736
737impl MBartGenerator {
738 /// Build a new `MBartGenerator`
739 ///
740 /// # Arguments
741 ///
742 /// * `vocab_path` - Path to the model vocabulary, expected to have a structure following the [Transformers library](https://github.com/huggingface/transformers) convention
743 /// * `merges_path` - Path to the bpe merges, expected to have a structure following the [Transformers library](https://github.com/huggingface/transformers) convention
744 /// * `config_path` - Path to the model configuration, expected to have a structure following the [Transformers library](https://github.com/huggingface/transformers) convention
745 /// * `weights_path` - Path to the model weight files. These need to be converted form the `.bin` to `.ot` format using the utility script provided.
746 /// * `device` - Device to run the model on, e.g. `Device::Cpu` or `Device::Cuda(0)`
747 ///
748 /// # Example
749 ///
750 /// ```no_run
751 /// # use std::path::PathBuf;
752 /// # use tch::Device;
753 /// # fn main() -> anyhow::Result<()> {
754 /// use rust_bert::mbart::MBartGenerator;
755 /// use rust_bert::pipelines::generation_utils::GenerateConfig;
756 /// # let mut home: PathBuf = dirs::home_dir().unwrap();
757 /// # home.push("rustbert");
758 /// # home.push("openai-gpt");
759 /// # let config_path = &home.as_path().join("config.json");
760 /// # let vocab_path = &home.as_path().join("vocab.txt");
761 /// # let merges_path = &home.as_path().join("merges.txt");
762 /// # let weights_path = &home.as_path().join("model.ot");
763 /// let device = Device::cuda_if_available();
764 /// let generate_config = GenerateConfig {
765 /// max_length: Some(30),
766 /// do_sample: true,
767 /// num_beams: 5,
768 /// temperature: 1.1,
769 /// num_return_sequences: 3,
770 /// ..Default::default()
771 /// };
772 /// let mbart_generator = MBartGenerator::new(generate_config)?;
773 /// # Ok(())
774 /// # }
775 /// ```
776 pub fn new(generate_config: GenerateConfig) -> Result<MBartGenerator, RustBertError> {
777 let vocab_path = generate_config.vocab_resource.get_local_path()?;
778
779 let tokenizer = TokenizerOption::from_file(
780 ModelType::MBart,
781 vocab_path.to_str().unwrap(),
782 None,
783 false,
784 None,
785 None,
786 )?;
787
788 Self::new_with_tokenizer(generate_config, tokenizer)
789 }
790
791 pub fn new_with_tokenizer(
792 generate_config: GenerateConfig,
793 tokenizer: TokenizerOption,
794 ) -> Result<MBartGenerator, RustBertError> {
795 let config_path = generate_config.config_resource.get_local_path()?;
796 let device = generate_config.device;
797
798 generate_config.validate();
799 let mut var_store = nn::VarStore::new(device);
800
801 let config = MBartConfig::from_file(config_path);
802 let model = MBartForConditionalGeneration::new(var_store.root(), &config);
803 crate::resources::load_weights(
804 &generate_config.model_resource,
805 &mut var_store,
806 generate_config.kind,
807 device,
808 )?;
809
810 let bos_token_id = Some(config.bos_token_id.unwrap_or(0));
811 let eos_token_ids = Some(match config.eos_token_id {
812 Some(value) => vec![value],
813 None => vec![2],
814 });
815 let forced_eos_token_id = config.forced_eos_token_id;
816 let pad_token_id = Some(config.pad_token_id.unwrap_or(1));
817 let vocab_size = config.vocab_size;
818 let is_encoder_decoder = true;
819 let decoder_start_id = config.decoder_start_token_id;
820 let max_position_embeddings = config.max_position_embeddings;
821
822 Ok(MBartGenerator {
823 model,
824 tokenizer,
825 var_store,
826 generate_config,
827 bos_token_id,
828 eos_token_ids,
829 forced_eos_token_id,
830 pad_token_id,
831 is_encoder_decoder,
832 vocab_size,
833 decoder_start_id,
834 max_position_embeddings,
835 })
836 }
837}
838
839impl PrivateLanguageGenerator for MBartGenerator {
840 fn _get_tokenizer(&self) -> &TokenizerOption {
841 &self.tokenizer
842 }
843 fn _get_tokenizer_mut(&mut self) -> &mut TokenizerOption {
844 &mut self.tokenizer
845 }
846 fn get_device(&self) -> Device {
847 self.var_store.device()
848 }
849 fn get_var_store_mut(&mut self) -> Result<&mut nn::VarStore, RustBertError> {
850 Ok(&mut self.var_store)
851 }
852 fn get_config(&self) -> &GenerateConfig {
853 &self.generate_config
854 }
855 fn get_bos_id(&self) -> Option<i64> {
856 self.bos_token_id
857 }
858 fn get_eos_ids(&self) -> Option<&Vec<i64>> {
859 self.eos_token_ids.as_ref()
860 }
861 fn get_forced_eos_token_id(&self) -> Option<i64> {
862 self.forced_eos_token_id
863 }
864 fn get_pad_id(&self) -> Option<i64> {
865 self.pad_token_id
866 }
867 fn is_encoder_decoder(&self) -> bool {
868 self.is_encoder_decoder
869 }
870 fn get_vocab_size(&self) -> i64 {
871 self.vocab_size
872 }
873 fn get_decoder_start_id(&self) -> Option<i64> {
874 self.decoder_start_id
875 }
876
877 fn forward_t(
878 &self,
879 input_ids: Option<&Tensor>,
880 cache: Cache,
881 attention_mask: Option<&Tensor>,
882 _token_type_ids: Option<&Tensor>,
883 _position_ids: Option<&Tensor>,
884 _input_embeds: Option<&Tensor>,
885 encoder_outputs: Option<&Tensor>,
886 decoder_input_ids: Option<&Tensor>,
887 train: bool,
888 ) -> Result<LMModelOutput, RustBertError> {
889 let base_model_output = match cache {
890 Cache::BARTCache(cached_layer_states) => self.model.forward_t(
891 input_ids,
892 attention_mask,
893 encoder_outputs,
894 decoder_input_ids,
895 None,
896 cached_layer_states,
897 train,
898 ),
899
900 Cache::None => self.model.forward_t(
901 input_ids,
902 attention_mask,
903 encoder_outputs,
904 decoder_input_ids,
905 None,
906 None,
907 train,
908 ),
909 _ => {
910 return Err(RustBertError::ValueError(
911 "Cache not compatible with MBART Model".into(),
912 ));
913 }
914 };
915
916 Ok(LMModelOutput {
917 lm_logits: base_model_output.decoder_output,
918 cache: Cache::BARTCache(base_model_output.cache),
919 })
920 }
921
922 fn get_max_positions_embeddings(&self) -> Option<i64> {
923 Some(self.max_position_embeddings)
924 }
925
926 fn encode(&self, input_ids: &Tensor, attention_mask: Option<&Tensor>) -> Option<Tensor> {
927 Some(self.model.encode(input_ids, attention_mask))
928 }
929
930 fn prepare_inputs_for_generation<'a>(
931 &self,
932 input_ids: Tensor,
933 encoder_outputs: Option<&'a Tensor>,
934 past: Cache,
935 attention_mask: Tensor,
936 ) -> PreparedInput<'a> {
937 match past {
938 Cache::BARTCache(past) => PreparedInput {
939 prepared_input: None,
940 prepared_attention_mask: Some(attention_mask),
941 prepared_encoder_output: encoder_outputs,
942 prepared_decoder_input: Some(input_ids.narrow(1, -1, 1)),
943 prepared_position_ids: None,
944 prepared_past: Cache::BARTCache(past),
945 },
946 Cache::None => PreparedInput {
947 prepared_input: None,
948 prepared_attention_mask: Some(attention_mask),
949 prepared_encoder_output: encoder_outputs,
950 prepared_decoder_input: Some(input_ids),
951 prepared_position_ids: None,
952 prepared_past: Cache::BARTCache(None),
953 },
954 _ => panic!("Cache type incompatible with MBart"),
955 }
956 }
957
958 fn reorder_cache(
959 &self,
960 past: &mut Cache,
961 encoder_outputs: Option<Tensor>,
962 beam_indices: &Tensor,
963 ) -> Option<Tensor> {
964 let encoder_outputs = encoder_outputs.map(|value| value.index_select(0, beam_indices));
965 match past {
966 Cache::BARTCache(old_cache_option) => match old_cache_option {
967 Some(old_cache) => {
968 for (self_layer_state, encoder_layer_state) in old_cache.iter_mut() {
969 if self_layer_state.is_some() {
970 self_layer_state
971 .as_mut()
972 .unwrap()
973 .reorder_cache(beam_indices)
974 };
975 if encoder_layer_state.is_some() {
976 encoder_layer_state
977 .as_mut()
978 .unwrap()
979 .reorder_cache(beam_indices)
980 };
981 }
982 }
983 None => {}
984 },
985 Cache::None => {}
986 _ => {
987 panic!("Invalid cache for MBart model");
988 }
989 };
990 encoder_outputs
991 }
992}
993
994impl LanguageGenerator for MBartGenerator {}
995
996#[cfg(test)]
997mod test {
998 use tch::Device;
999
1000 use crate::{
1001 resources::{RemoteResource, ResourceProvider},
1002 Config,
1003 };
1004
1005 use super::*;
1006
1007 #[test]
1008 #[ignore] // compilation is enough, no need to run
1009 fn mbart_model_send() {
1010 let config_resource = Box::new(RemoteResource::from_pretrained(
1011 MBartConfigResources::MBART50_MANY_TO_MANY,
1012 ));
1013 let config_path = config_resource.get_local_path().expect("");
1014
1015 // Set-up masked LM model
1016 let device = Device::cuda_if_available();
1017 let vs = tch::nn::VarStore::new(device);
1018 let config = MBartConfig::from_file(config_path);
1019
1020 let _: Box<dyn Send> = Box::new(MBartModel::new(vs.root(), &config));
1021 }
1022}