mirror of
https://github.com/guillaume-be/rust-bert.git
synced 2024-09-11 12:55:34 +03:00
Added MBart50 and M2M100 to supported translation models
This commit is contained in:
parent
3b72b7cc9b
commit
5dc7f33c39
@ -21,15 +21,16 @@ fn main() -> anyhow::Result<()> {
|
|||||||
let model = TranslationModelBuilder::new()
|
let model = TranslationModelBuilder::new()
|
||||||
.with_device(Device::cuda_if_available())
|
.with_device(Device::cuda_if_available())
|
||||||
// .with_model_type(ModelType::Marian)
|
// .with_model_type(ModelType::Marian)
|
||||||
|
.with_model_type(ModelType::M2M100)
|
||||||
.with_large_model()
|
.with_large_model()
|
||||||
.with_source_languages(vec![Language::English])
|
.with_source_languages(vec![Language::English])
|
||||||
.with_target_languages(vec![Language::French])
|
.with_target_languages(vec![Language::French])
|
||||||
.create_model()?;
|
.create_model()?;
|
||||||
|
|
||||||
let input_context_1 = "The quick brown fox jumps over the lazy dog";
|
// let input_context_1 = "The quick brown fox jumps over the lazy dog.";
|
||||||
let input_context_2 = "The dog did not wake up";
|
let input_context_2 = "The dog did not wake up.";
|
||||||
|
|
||||||
let output = model.translate(&[input_context_1, input_context_2], None, Language::French)?;
|
let output = model.translate(&[input_context_2], Language::English, Language::French)?;
|
||||||
|
|
||||||
for sentence in output {
|
for sentence in output {
|
||||||
println!("{}", sentence);
|
println!("{}", sentence);
|
||||||
|
@ -21,7 +21,7 @@ use rust_bert::resources::{RemoteResource, Resource};
|
|||||||
|
|
||||||
fn main() -> anyhow::Result<()> {
|
fn main() -> anyhow::Result<()> {
|
||||||
let generate_config = GenerateConfig {
|
let generate_config = GenerateConfig {
|
||||||
max_length: 142,
|
max_length: 512,
|
||||||
min_length: 0,
|
min_length: 0,
|
||||||
model_resource: Resource::Remote(RemoteResource::from_pretrained(
|
model_resource: Resource::Remote(RemoteResource::from_pretrained(
|
||||||
M2M100ModelResources::M2M100_418M,
|
M2M100ModelResources::M2M100_418M,
|
||||||
@ -38,14 +38,16 @@ fn main() -> anyhow::Result<()> {
|
|||||||
do_sample: false,
|
do_sample: false,
|
||||||
early_stopping: true,
|
early_stopping: true,
|
||||||
num_beams: 3,
|
num_beams: 3,
|
||||||
|
no_repeat_ngram_size: 0,
|
||||||
..Default::default()
|
..Default::default()
|
||||||
};
|
};
|
||||||
|
|
||||||
let model = M2M100Generator::new(generate_config)?;
|
let model = M2M100Generator::new(generate_config)?;
|
||||||
|
|
||||||
let input_context_1 = ">>en.<< The dog did not wake up.";
|
let input_context_1 = ">>en.<< The dog did not wake up.";
|
||||||
let target_language = model.get_tokenizer().convert_tokens_to_ids([">>es.<<"])[0];
|
let target_language = model.get_tokenizer().convert_tokens_to_ids([">>fr.<<"])[0];
|
||||||
|
|
||||||
|
println!("{:?} - {:?}", input_context_1, target_language);
|
||||||
let output = model.generate(
|
let output = model.generate(
|
||||||
Some(&[input_context_1]),
|
Some(&[input_context_1]),
|
||||||
None,
|
None,
|
||||||
|
@ -50,7 +50,7 @@ use serde::{Deserialize, Serialize};
|
|||||||
use std::collections::HashMap;
|
use std::collections::HashMap;
|
||||||
use std::path::Path;
|
use std::path::Path;
|
||||||
|
|
||||||
#[derive(Clone, Copy, Serialize, Deserialize, Debug)]
|
#[derive(Clone, Copy, Serialize, Deserialize, Debug, PartialEq)]
|
||||||
/// # Identifies the type of model
|
/// # Identifies the type of model
|
||||||
pub enum ModelType {
|
pub enum ModelType {
|
||||||
Bart,
|
Bart,
|
||||||
|
@ -16,18 +16,19 @@ use tch::{Device, Tensor};
|
|||||||
use crate::common::error::RustBertError;
|
use crate::common::error::RustBertError;
|
||||||
use crate::common::resources::Resource;
|
use crate::common::resources::Resource;
|
||||||
use crate::m2m_100::{
|
use crate::m2m_100::{
|
||||||
M2M100ConfigResources, M2M100MergesResources, M2M100ModelResources, M2M100SourceLanguages,
|
M2M100ConfigResources, M2M100Generator, M2M100MergesResources, M2M100ModelResources,
|
||||||
M2M100TargetLanguages, M2M100VocabResources,
|
M2M100SourceLanguages, M2M100TargetLanguages, M2M100VocabResources,
|
||||||
};
|
};
|
||||||
use crate::marian::{
|
use crate::marian::{
|
||||||
MarianConfigResources, MarianGenerator, MarianModelResources, MarianSourceLanguages,
|
MarianConfigResources, MarianGenerator, MarianModelResources, MarianSourceLanguages,
|
||||||
MarianSpmResources, MarianTargetLanguages, MarianVocabResources,
|
MarianSpmResources, MarianTargetLanguages, MarianVocabResources,
|
||||||
};
|
};
|
||||||
use crate::mbart::{
|
use crate::mbart::{
|
||||||
MBartConfigResources, MBartModelResources, MBartSourceLanguages, MBartTargetLanguages,
|
MBartConfigResources, MBartGenerator, MBartModelResources, MBartSourceLanguages,
|
||||||
MBartVocabResources,
|
MBartTargetLanguages, MBartVocabResources,
|
||||||
};
|
};
|
||||||
use crate::pipelines::common::ModelType;
|
use crate::pipelines::common::ModelType;
|
||||||
|
use crate::pipelines::generation_utils::private_generation_utils::PrivateLanguageGenerator;
|
||||||
use crate::pipelines::generation_utils::{GenerateConfig, LanguageGenerator};
|
use crate::pipelines::generation_utils::{GenerateConfig, LanguageGenerator};
|
||||||
use crate::resources::RemoteResource;
|
use crate::resources::RemoteResource;
|
||||||
use crate::t5::T5Generator;
|
use crate::t5::T5Generator;
|
||||||
@ -501,7 +502,7 @@ impl TranslationConfig {
|
|||||||
max_length: 512,
|
max_length: 512,
|
||||||
do_sample: false,
|
do_sample: false,
|
||||||
early_stopping: true,
|
early_stopping: true,
|
||||||
num_beams: 4,
|
num_beams: 3,
|
||||||
temperature: 1.0,
|
temperature: 1.0,
|
||||||
top_k: 50,
|
top_k: 50,
|
||||||
top_p: 1.0,
|
top_p: 1.0,
|
||||||
@ -547,6 +548,10 @@ pub enum TranslationOption {
|
|||||||
Marian(MarianGenerator),
|
Marian(MarianGenerator),
|
||||||
/// Translator based on T5 model
|
/// Translator based on T5 model
|
||||||
T5(T5Generator),
|
T5(T5Generator),
|
||||||
|
/// Translator based on MBart50 model
|
||||||
|
MBart(MBartGenerator),
|
||||||
|
/// Translator based on M2M100 model
|
||||||
|
M2M100(M2M100Generator),
|
||||||
}
|
}
|
||||||
|
|
||||||
impl TranslationOption {
|
impl TranslationOption {
|
||||||
@ -556,6 +561,12 @@ impl TranslationOption {
|
|||||||
config.into(),
|
config.into(),
|
||||||
)?)),
|
)?)),
|
||||||
ModelType::T5 => Ok(TranslationOption::T5(T5Generator::new(config.into())?)),
|
ModelType::T5 => Ok(TranslationOption::T5(T5Generator::new(config.into())?)),
|
||||||
|
ModelType::MBart => Ok(TranslationOption::MBart(MBartGenerator::new(
|
||||||
|
config.into(),
|
||||||
|
)?)),
|
||||||
|
ModelType::M2M100 => Ok(TranslationOption::M2M100(M2M100Generator::new(
|
||||||
|
config.into(),
|
||||||
|
)?)),
|
||||||
_ => Err(RustBertError::InvalidConfigurationError(format!(
|
_ => Err(RustBertError::InvalidConfigurationError(format!(
|
||||||
"Translation not implemented for {:?}!",
|
"Translation not implemented for {:?}!",
|
||||||
config.model_type
|
config.model_type
|
||||||
@ -568,16 +579,18 @@ impl TranslationOption {
|
|||||||
match *self {
|
match *self {
|
||||||
Self::Marian(_) => ModelType::Marian,
|
Self::Marian(_) => ModelType::Marian,
|
||||||
Self::T5(_) => ModelType::T5,
|
Self::T5(_) => ModelType::T5,
|
||||||
|
Self::MBart(_) => ModelType::MBart,
|
||||||
|
Self::M2M100(_) => ModelType::M2M100,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
fn validate_and_get_prefix(
|
fn validate_and_get_prefix_and_forced_bos_id(
|
||||||
&self,
|
&self,
|
||||||
source_language: Option<&Language>,
|
source_language: Option<&Language>,
|
||||||
target_language: Option<&Language>,
|
target_language: Option<&Language>,
|
||||||
supported_source_languages: &HashSet<Language>,
|
supported_source_languages: &HashSet<Language>,
|
||||||
supported_target_languages: &HashSet<Language>,
|
supported_target_languages: &HashSet<Language>,
|
||||||
) -> Result<Option<String>, RustBertError> {
|
) -> Result<(Option<String>, Option<i64>), RustBertError> {
|
||||||
if let Some(source_language) = source_language {
|
if let Some(source_language) = source_language {
|
||||||
if !supported_source_languages.contains(source_language) {
|
if !supported_source_languages.contains(source_language) {
|
||||||
return Err(RustBertError::ValueError(format!(
|
return Err(RustBertError::ValueError(format!(
|
||||||
@ -601,30 +614,112 @@ impl TranslationOption {
|
|||||||
Ok(match *self {
|
Ok(match *self {
|
||||||
Self::Marian(_) => {
|
Self::Marian(_) => {
|
||||||
if supported_target_languages.len() > 1 {
|
if supported_target_languages.len() > 1 {
|
||||||
Some(format!(
|
(
|
||||||
">>{}<< ",
|
Some(format!(
|
||||||
match target_language {
|
">>{}<< ",
|
||||||
Some(value) => value.get_iso_639_1_code(),
|
match target_language {
|
||||||
None => {
|
Some(value) => value.get_iso_639_1_code(),
|
||||||
|
None => {
|
||||||
|
return Err(RustBertError::ValueError(
|
||||||
|
"Missing target language for Marian".to_string(),
|
||||||
|
));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
)),
|
||||||
|
None,
|
||||||
|
)
|
||||||
|
} else {
|
||||||
|
(None, None)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Self::T5(_) => (
|
||||||
|
Some(format!(
|
||||||
|
"translate {} to {}:",
|
||||||
|
match source_language {
|
||||||
|
Some(value) => value.to_string(),
|
||||||
|
None => {
|
||||||
|
return Err(RustBertError::ValueError(
|
||||||
|
"Missing source language for T5".to_string(),
|
||||||
|
));
|
||||||
|
}
|
||||||
|
},
|
||||||
|
match target_language {
|
||||||
|
Some(value) => value.to_string(),
|
||||||
|
None => {
|
||||||
|
return Err(RustBertError::ValueError(
|
||||||
|
"Missing target language for T5".to_string(),
|
||||||
|
));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
)),
|
||||||
|
None,
|
||||||
|
),
|
||||||
|
Self::MBart(ref model) => (
|
||||||
|
Some(format!(
|
||||||
|
">>{}<< ",
|
||||||
|
match source_language {
|
||||||
|
Some(value) => value.get_iso_639_1_code(),
|
||||||
|
None => {
|
||||||
|
return Err(RustBertError::ValueError(
|
||||||
|
"Missing source language for MBart".to_string(),
|
||||||
|
));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
)),
|
||||||
|
if let Some(target_language) = target_language {
|
||||||
|
Some(
|
||||||
|
model._get_tokenizer().convert_tokens_to_ids([format!(
|
||||||
|
">>{}<<",
|
||||||
|
target_language.get_iso_639_1_code()
|
||||||
|
)])[0],
|
||||||
|
)
|
||||||
|
} else {
|
||||||
|
return Err(RustBertError::ValueError(
|
||||||
|
"Missing target language for MBart".to_string(),
|
||||||
|
));
|
||||||
|
},
|
||||||
|
),
|
||||||
|
Self::M2M100(ref model) => (
|
||||||
|
Some(match source_language {
|
||||||
|
Some(value) => {
|
||||||
|
let language_code = value.get_iso_639_1_code();
|
||||||
|
match language_code.len() {
|
||||||
|
2 => format!(">>{}.<< ", language_code),
|
||||||
|
3 => format!(">>{}<< ", language_code),
|
||||||
|
_ => {
|
||||||
return Err(RustBertError::ValueError(
|
return Err(RustBertError::ValueError(
|
||||||
"Missing target language for Marian".to_string(),
|
"Invalid ISO 639-3 code".to_string(),
|
||||||
));
|
));
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
))
|
}
|
||||||
|
None => {
|
||||||
|
return Err(RustBertError::ValueError(
|
||||||
|
"Missing source language for M2M100".to_string(),
|
||||||
|
));
|
||||||
|
}
|
||||||
|
}),
|
||||||
|
if let Some(target_language) = target_language {
|
||||||
|
let language_code = target_language.get_iso_639_1_code();
|
||||||
|
Some(
|
||||||
|
model
|
||||||
|
._get_tokenizer()
|
||||||
|
.convert_tokens_to_ids([match language_code.len() {
|
||||||
|
2 => format!(">>{}.<<", language_code),
|
||||||
|
3 => format!(">>{}<<", language_code),
|
||||||
|
_ => {
|
||||||
|
return Err(RustBertError::ValueError(
|
||||||
|
"Invalid ISO 639-3 code".to_string(),
|
||||||
|
));
|
||||||
|
}
|
||||||
|
}])[0],
|
||||||
|
)
|
||||||
} else {
|
} else {
|
||||||
None
|
return Err(RustBertError::ValueError(
|
||||||
}
|
"Missing target language for MBart".to_string(),
|
||||||
}
|
));
|
||||||
Self::T5(_) => Some(format!(
|
},
|
||||||
"translate {} to {}:",
|
),
|
||||||
source_language
|
|
||||||
.expect("Missing source language for T5")
|
|
||||||
.to_string(),
|
|
||||||
target_language
|
|
||||||
.expect("Missing target language for T5")
|
|
||||||
.to_string()
|
|
||||||
)),
|
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
@ -633,6 +728,7 @@ impl TranslationOption {
|
|||||||
&self,
|
&self,
|
||||||
prompt_texts: Option<S>,
|
prompt_texts: Option<S>,
|
||||||
attention_mask: Option<Tensor>,
|
attention_mask: Option<Tensor>,
|
||||||
|
forced_bos_token_id: Option<i64>,
|
||||||
) -> Vec<String>
|
) -> Vec<String>
|
||||||
where
|
where
|
||||||
S: AsRef<[&'a str]>,
|
S: AsRef<[&'a str]>,
|
||||||
@ -666,6 +762,34 @@ impl TranslationOption {
|
|||||||
.into_iter()
|
.into_iter()
|
||||||
.map(|output| output.text)
|
.map(|output| output.text)
|
||||||
.collect(),
|
.collect(),
|
||||||
|
Self::MBart(ref model) => model
|
||||||
|
.generate(
|
||||||
|
prompt_texts,
|
||||||
|
attention_mask,
|
||||||
|
None,
|
||||||
|
None,
|
||||||
|
None,
|
||||||
|
forced_bos_token_id,
|
||||||
|
None,
|
||||||
|
false,
|
||||||
|
)
|
||||||
|
.into_iter()
|
||||||
|
.map(|output| output.text)
|
||||||
|
.collect(),
|
||||||
|
Self::M2M100(ref model) => model
|
||||||
|
.generate(
|
||||||
|
prompt_texts,
|
||||||
|
attention_mask,
|
||||||
|
None,
|
||||||
|
None,
|
||||||
|
None,
|
||||||
|
forced_bos_token_id,
|
||||||
|
None,
|
||||||
|
false,
|
||||||
|
)
|
||||||
|
.into_iter()
|
||||||
|
.map(|output| output.text)
|
||||||
|
.collect(),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@ -797,7 +921,7 @@ impl TranslationModel {
|
|||||||
where
|
where
|
||||||
S: AsRef<[&'a str]>,
|
S: AsRef<[&'a str]>,
|
||||||
{
|
{
|
||||||
let prefix = self.model.validate_and_get_prefix(
|
let (prefix, forced_bos_token_id) = self.model.validate_and_get_prefix_and_forced_bos_id(
|
||||||
source_language.into().as_ref(),
|
source_language.into().as_ref(),
|
||||||
target_language.into().as_ref(),
|
target_language.into().as_ref(),
|
||||||
&self.supported_source_languages,
|
&self.supported_source_languages,
|
||||||
@ -814,9 +938,10 @@ impl TranslationModel {
|
|||||||
self.model.generate(
|
self.model.generate(
|
||||||
Some(texts.iter().map(AsRef::as_ref).collect::<Vec<&str>>()),
|
Some(texts.iter().map(AsRef::as_ref).collect::<Vec<&str>>()),
|
||||||
None,
|
None,
|
||||||
|
forced_bos_token_id,
|
||||||
)
|
)
|
||||||
}
|
}
|
||||||
None => self.model.generate(Some(texts), None),
|
None => self.model.generate(Some(texts), None, forced_bos_token_id),
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@ -876,12 +1001,14 @@ where
|
|||||||
}
|
}
|
||||||
|
|
||||||
pub fn with_medium_model(&mut self) -> &mut Self {
|
pub fn with_medium_model(&mut self) -> &mut Self {
|
||||||
if self.model_type.is_some() {
|
if let Some(model_type) = self.model_type {
|
||||||
eprintln!(
|
if model_type != ModelType::Marian {
|
||||||
"Model selection overwritten: was {:?}, replaced by {:?}",
|
eprintln!(
|
||||||
self.model_type.unwrap(),
|
"Model selection overwritten: was {:?}, replaced by {:?}",
|
||||||
ModelType::Marian
|
self.model_type.unwrap(),
|
||||||
);
|
ModelType::Marian
|
||||||
|
);
|
||||||
|
}
|
||||||
}
|
}
|
||||||
self.model_type = Some(ModelType::Marian);
|
self.model_type = Some(ModelType::Marian);
|
||||||
self.model_size = Some(ModelSize::Medium);
|
self.model_size = Some(ModelSize::Medium);
|
||||||
@ -889,12 +1016,14 @@ where
|
|||||||
}
|
}
|
||||||
|
|
||||||
pub fn with_large_model(&mut self) -> &mut Self {
|
pub fn with_large_model(&mut self) -> &mut Self {
|
||||||
if self.model_type.is_some() {
|
if let Some(model_type) = self.model_type {
|
||||||
eprintln!(
|
if model_type != ModelType::M2M100 {
|
||||||
"Model selection overwritten: was {:?}, replaced by {:?}",
|
eprintln!(
|
||||||
self.model_type.unwrap(),
|
"Model selection overwritten: was {:?}, replaced by {:?}",
|
||||||
ModelType::M2M100
|
self.model_type.unwrap(),
|
||||||
);
|
ModelType::M2M100
|
||||||
|
);
|
||||||
|
}
|
||||||
}
|
}
|
||||||
self.model_type = Some(ModelType::M2M100);
|
self.model_type = Some(ModelType::M2M100);
|
||||||
self.model_size = Some(ModelSize::Large);
|
self.model_size = Some(ModelSize::Large);
|
||||||
@ -902,12 +1031,14 @@ where
|
|||||||
}
|
}
|
||||||
|
|
||||||
pub fn with_xlarge_model(&mut self) -> &mut Self {
|
pub fn with_xlarge_model(&mut self) -> &mut Self {
|
||||||
if self.model_type.is_some() {
|
if let Some(model_type) = self.model_type {
|
||||||
eprintln!(
|
if model_type != ModelType::M2M100 {
|
||||||
"Model selection overwritten: was {:?}, replaced by {:?}",
|
eprintln!(
|
||||||
self.model_type.unwrap(),
|
"Model selection overwritten: was {:?}, replaced by {:?}",
|
||||||
ModelType::M2M100
|
self.model_type.unwrap(),
|
||||||
);
|
ModelType::M2M100
|
||||||
|
);
|
||||||
|
}
|
||||||
}
|
}
|
||||||
self.model_type = Some(ModelType::M2M100);
|
self.model_type = Some(ModelType::M2M100);
|
||||||
self.model_size = Some(ModelSize::XLarge);
|
self.model_size = Some(ModelSize::XLarge);
|
||||||
|
Loading…
Reference in New Issue
Block a user