Struct rust_bert::pipelines::translation::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::ModelType;
use rust_bert::pipelines::translation::{TranslationConfig, TranslationModel};
use rust_bert::resources::RemoteResource;
use tch::Device;

let model_resource = 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 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> + 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::ModelType;
use rust_bert::pipelines::translation::{Language, TranslationConfig, TranslationModel};
use rust_bert::resources::RemoteResource;
use tch::Device;

let model_resource = 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);