Struct rust_bert::pipelines::translation::TranslationModel

pub struct TranslationModel { /* fields omitted */ }

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

let model_resource = Resource::Remote(RemoteResource::from_pretrained(
    MarianModelResources::ROMANCE2ENGLISH,
));
let config_resource = Resource::Remote(RemoteResource::from_pretrained(
    MarianConfigResources::ROMANCE2ENGLISH,
));
let vocab_resource = Resource::Remote(RemoteResource::from_pretrained(
    MarianVocabResources::ROMANCE2ENGLISH,
));

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

let translation_config = TranslationConfig::new(
    ModelType::Marian,
    model_resource,
    config_resource,
    vocab_resource.clone(),
    vocab_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, Resource};
use tch::Device;

let model_resource = Resource::Remote(RemoteResource::from_pretrained(
    MarianModelResources::ENGLISH2ROMANCE,
));
let config_resource = Resource::Remote(RemoteResource::from_pretrained(
    MarianConfigResources::ENGLISH2ROMANCE,
));
let vocab_resource = Resource::Remote(RemoteResource::from_pretrained(
    MarianVocabResources::ENGLISH2ROMANCE,
));
let merges_resource = Resource::Remote(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,
    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);