mirror of
https://github.com/guillaume-be/rust-bert.git
synced 2024-09-11 12:55:34 +03:00
* replace failure with anyhow * formatting update with `fmt` Co-authored-by: guillaume-be <guillaume.becquin@gmail.com>
25 lines
812 B
Rust
25 lines
812 B
Rust
use rust_bert::pipelines::translation::{Language, TranslationConfig, TranslationModel};
|
|
use tch::Device;
|
|
|
|
#[test]
|
|
#[cfg_attr(not(feature = "all-tests"), ignore)]
|
|
fn test_translation() -> anyhow::Result<()> {
|
|
// Set-up translation model
|
|
let translation_config = TranslationConfig::new(Language::EnglishToFrench, Device::Cpu);
|
|
let model = TranslationModel::new(translation_config)?;
|
|
|
|
let input_context_1 = "The quick brown fox jumps over the lazy dog";
|
|
let input_context_2 = "The dog did not wake up";
|
|
|
|
let output = model.translate(&[input_context_1, input_context_2]);
|
|
|
|
assert_eq!(output.len(), 2);
|
|
assert_eq!(
|
|
output[0],
|
|
" Le rapide renard brun saute sur le chien paresseux"
|
|
);
|
|
assert_eq!(output[1], " Le chien ne s'est pas réveillé.");
|
|
|
|
Ok(())
|
|
}
|