def test_pipeline(self):
pipeline = TranslationPipeline(self.model,
self.tokenizer,
framework="pt")
output = pipeline(self.src_text)
self.assertEqual(self.expected_text,
[x["translation_text"] for x in output])
def test_pipeline(self):
device = 0 if torch_device == "cuda" else -1
pipeline = TranslationPipeline(self.model,
self.tokenizer,
framework="pt",
device=device)
output = pipeline(self.src_text)
self.assertEqual(self.expected_text,
[x["translation_text"] for x in output])
def test_translation_pipeline(self, pair):
tokenizer, model, src_text, tgt_text = self.translation_setup(pair)
device = 0 if torch_device == "cuda" else -1
pipeline = TranslationPipeline(model,
tokenizer,
framework="pt",
device=device)
output = pipeline([src_text])
self.assertEqual([tgt_text], [x["translation_text"] for x in output])