Struct TranslationModel

pub struct TranslationModel { /* private fields */ }

TranslationModel to perform translation

Implementations

impl TranslationModel

pub fn new( translation_config: TranslationConfig, ) -> Result<TranslationModel, RustBertError>

Build a new TranslationModel

Arguments

• translation_config - TranslationConfig object containing the resource references (model, vocabulary, configuration), translation options and device placement (CPU/GPU)

Example

use rust_bert::marian::{
    MarianConfigResources, MarianModelResources, MarianSourceLanguages, MarianSpmResources,
    MarianTargetLanguages, MarianVocabResources,
};
use rust_bert::pipelines::common::{ModelResource, ModelType};
use rust_bert::pipelines::translation::{TranslationConfig, TranslationModel};
use rust_bert::resources::RemoteResource;
use tch::Device;

let model_resource = ModelResource::Torch(Box::new(RemoteResource::from_pretrained(
    MarianModelResources::ROMANCE2ENGLISH,
)));
let config_resource = RemoteResource::from_pretrained(MarianConfigResources::ROMANCE2ENGLISH);
let vocab_resource = RemoteResource::from_pretrained(MarianVocabResources::ROMANCE2ENGLISH);
let spm_resource = RemoteResource::from_pretrained(MarianSpmResources::ROMANCE2ENGLISH);

let source_languages = MarianSourceLanguages::ROMANCE2ENGLISH;
let target_languages = MarianTargetLanguages::ROMANCE2ENGLISH;

let translation_config = TranslationConfig::new(
    ModelType::Marian,
    model_resource,
    config_resource,
    vocab_resource,
    Some(spm_resource),
    source_languages,
    target_languages,
    Device::cuda_if_available(),
);
let mut summarization_model = TranslationModel::new(translation_config)?;

pub fn new_with_tokenizer( translation_config: TranslationConfig, tokenizer: TokenizerOption, ) -> Result<TranslationModel, RustBertError>

Build a new TranslationModel with a provided tokenizer.

Arguments

• translation_config - TranslationConfig object containing the resource references (model, vocabulary, configuration), translation options and device placement (CPU/GPU)
• tokenizer - TokenizerOption tokenizer to use for translation

Example

use rust_bert::marian::{
    MarianConfigResources, MarianModelResources, MarianSourceLanguages, MarianSpmResources,
    MarianTargetLanguages, MarianVocabResources,
};
use rust_bert::pipelines::common::{ModelResource, ModelType, TokenizerOption};
use rust_bert::pipelines::translation::{TranslationConfig, TranslationModel};
use rust_bert::resources::{RemoteResource, ResourceProvider};
use tch::Device;

let model_resource = ModelResource::Torch(Box::new(RemoteResource::from_pretrained(
    MarianModelResources::ROMANCE2ENGLISH,
)));
let config_resource = RemoteResource::from_pretrained(MarianConfigResources::ROMANCE2ENGLISH);
let vocab_resource = RemoteResource::from_pretrained(MarianVocabResources::ROMANCE2ENGLISH);
let spm_resource = RemoteResource::from_pretrained(MarianSpmResources::ROMANCE2ENGLISH);

let tokenizer = TokenizerOption::from_file(
    ModelType::Marian,
    vocab_resource.get_local_path()?.to_str().unwrap(),
    Some(spm_resource.get_local_path()?.to_str().unwrap()),
    false,
    None,
    None,
)?;

let source_languages = MarianSourceLanguages::ROMANCE2ENGLISH;
let target_languages = MarianTargetLanguages::ROMANCE2ENGLISH;

let translation_config = TranslationConfig::new(
    ModelType::Marian,
    model_resource,
    config_resource,
    vocab_resource,
    Some(spm_resource),
    source_languages,
    target_languages,
    Device::cuda_if_available(),
);
let mut summarization_model =
    TranslationModel::new_with_tokenizer(translation_config, tokenizer)?;

pub fn get_tokenizer(&self) -> &TokenizerOption

Get a reference to the model tokenizer.

pub fn get_tokenizer_mut(&mut self) -> &mut TokenizerOption

Get a mutable reference to the model tokenizer.

pub fn translate<S>( &self, texts: &[S], source_language: impl Into<Option<Language>>, target_language: impl Into<Option<Language>>, ) -> Result<Vec<String>, RustBertError>

where S: AsRef<str> + Send + Sync,

Translates texts provided

Arguments

• input - &[&str] Array of texts to summarize.

Returns

• Vec<String> Translated texts

Example

use rust_bert::marian::{
    MarianConfigResources, MarianModelResources, MarianSourceLanguages, MarianSpmResources,
    MarianTargetLanguages, MarianVocabResources,
};
use rust_bert::pipelines::common::{ModelResource, ModelType};
use rust_bert::pipelines::translation::{Language, TranslationConfig, TranslationModel};
use rust_bert::resources::RemoteResource;
use tch::Device;

let model_resource = ModelResource::Torch(Box::new(RemoteResource::from_pretrained(
    MarianModelResources::ENGLISH2ROMANCE,
)));
let config_resource = RemoteResource::from_pretrained(MarianConfigResources::ENGLISH2ROMANCE);
let vocab_resource = RemoteResource::from_pretrained(MarianVocabResources::ENGLISH2ROMANCE);
let merges_resource = RemoteResource::from_pretrained(MarianSpmResources::ENGLISH2ROMANCE);
let source_languages = MarianSourceLanguages::ENGLISH2ROMANCE;
let target_languages = MarianTargetLanguages::ENGLISH2ROMANCE;

let translation_config = TranslationConfig::new(
    ModelType::Marian,
    model_resource,
    config_resource,
    vocab_resource,
    Some(merges_resource),
    source_languages,
    target_languages,
    Device::cuda_if_available(),
);
let model = TranslationModel::new(translation_config)?;

let input = ["This is a sentence to be translated"];
let source_language = None;
let target_language = Language::French;

let output = model.translate(&input, source_language, target_language);