[cleanup/marian] pipelines test and new kwarg (#4812)
This commit is contained in:
@@ -38,6 +38,7 @@ if is_torch_available():
|
||||
convert_opus_name_to_hf_name,
|
||||
ORG_NAME,
|
||||
)
|
||||
from transformers.pipelines import TranslationPipeline
|
||||
|
||||
|
||||
class ModelManagementTests(unittest.TestCase):
|
||||
@@ -189,6 +190,7 @@ class TestMarian_RU_FR(MarianIntegrationTest):
|
||||
src_text = ["Он показал мне рукопись своей новой пьесы."]
|
||||
expected_text = ["Il m'a montré le manuscrit de sa nouvelle pièce."]
|
||||
|
||||
@slow
|
||||
def test_batch_generation_ru_fr(self):
|
||||
self._assert_generated_batch_equal_expected()
|
||||
|
||||
@@ -199,6 +201,7 @@ class TestMarian_MT_EN(MarianIntegrationTest):
|
||||
src_text = ["Billi messu b'mod ġentili, Ġesù fejjaq raġel li kien milqut bil - marda kerha tal - ġdiem."]
|
||||
expected_text = ["Touching gently, Jesus healed a man who was affected by the sad disease of leprosy."]
|
||||
|
||||
@slow
|
||||
def test_batch_generation_mt_en(self):
|
||||
self._assert_generated_batch_equal_expected()
|
||||
|
||||
@@ -229,6 +232,11 @@ class TestMarian_en_ROMANCE(MarianIntegrationTest):
|
||||
with self.assertRaises(ValueError):
|
||||
self.tokenizer.prepare_translation_batch([""])
|
||||
|
||||
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])
|
||||
|
||||
|
||||
@require_torch
|
||||
class TestConversionUtils(unittest.TestCase):
|
||||
|
||||
Reference in New Issue
Block a user