test_modeling_mbart.py 6.61 KB
Newer Older
1
2
3
4
import unittest

from transformers import is_torch_available
from transformers.file_utils import cached_property
5
from transformers.testing_utils import require_sentencepiece, require_tokenizers, require_torch, slow, torch_device
6

7
from .test_modeling_common import ModelTesterMixin
8
9
10
11


if is_torch_available():
    import torch
12

13
14
    from transformers import (
        AutoModelForSeq2SeqLM,
15
16
        AutoTokenizer,
        BatchEncoding,
17
18
        MBartConfig,
        MBartForConditionalGeneration,
19
20
21
22
23
24
25
    )


EN_CODE = 250004
RO_CODE = 250020


26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
@require_torch
class ModelTester:
    def __init__(self, parent):
        self.config = MBartConfig(
            vocab_size=99,
            d_model=24,
            encoder_layers=2,
            decoder_layers=2,
            encoder_attention_heads=2,
            decoder_attention_heads=2,
            encoder_ffn_dim=32,
            decoder_ffn_dim=32,
            max_position_embeddings=48,
            add_final_layer_norm=True,
        )

    def prepare_config_and_inputs_for_common(self):
        return self.config, {}


@require_torch
class SelectiveCommonTest(unittest.TestCase):
    all_model_classes = (MBartForConditionalGeneration,) if is_torch_available() else ()

    test_save_load_keys_to_never_save = ModelTesterMixin.test_save_load_keys_to_never_save

    def setUp(self):
        self.model_tester = ModelTester(self)


56
@require_torch
57
58
@require_sentencepiece
@require_tokenizers
59
60
class AbstractSeq2SeqIntegrationTest(unittest.TestCase):
    maxDiff = 1000  # longer string compare tracebacks
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
    checkpoint_name = None

    @classmethod
    def setUpClass(cls):
        cls.tokenizer = AutoTokenizer.from_pretrained(cls.checkpoint_name)
        return cls

    @cached_property
    def model(self):
        """Only load the model if needed."""
        model = AutoModelForSeq2SeqLM.from_pretrained(self.checkpoint_name).to(torch_device)
        if "cuda" in torch_device:
            model = model.half()
        return model


@require_torch
78
79
@require_sentencepiece
@require_tokenizers
80
class MBartEnroIntegrationTest(AbstractSeq2SeqIntegrationTest):
81
82
83
84
85
86
87
    checkpoint_name = "facebook/mbart-large-en-ro"
    src_text = [
        " UN Chief Says There Is No Military Solution in Syria",
        """ Secretary-General Ban Ki-moon says his response to Russia's stepped up military support for Syria is that "there is no military solution" to the nearly five-year conflict and more weapons will only worsen the violence and misery for millions of people.""",
    ]
    tgt_text = [
        "艦eful ONU declar膬 c膬 nu exist膬 o solu牛ie militar膬 卯n Siria",
88
        'Secretarul General Ban Ki-moon declar膬 c膬 r膬spunsul s膬u la intensificarea sprijinului militar al Rusiei pentru Siria este c膬 "nu exist膬 o solu牛ie militar膬" la conflictul de aproape cinci ani 艧i c膬 noi arme nu vor face dec芒t s膬 卯nr膬ut膬牛easc膬 violen牛a 艧i mizeria pentru milioane de oameni.',
89
90
91
    ]
    expected_src_tokens = [8274, 127873, 25916, 7, 8622, 2071, 438, 67485, 53, 187895, 23, 51712, 2, EN_CODE]

Sam Shleifer's avatar
Sam Shleifer committed
92
93
94
95
96
97
98
99
100
    @slow
    def test_enro_generate_one(self):
        batch: BatchEncoding = self.tokenizer.prepare_seq2seq_batch(
            ["UN Chief Says There Is No Military Solution in Syria"]
        ).to(torch_device)
        translated_tokens = self.model.generate(**batch)
        decoded = self.tokenizer.batch_decode(translated_tokens, skip_special_tokens=True)
        self.assertEqual(self.tgt_text[0], decoded[0])
        # self.assertEqual(self.tgt_text[1], decoded[1])
101
102

    @slow
Sam Shleifer's avatar
Sam Shleifer committed
103
    def test_enro_generate_batch(self):
104
        batch: BatchEncoding = self.tokenizer.prepare_seq2seq_batch(self.src_text).to(torch_device)
105
106
        translated_tokens = self.model.generate(**batch)
        decoded = self.tokenizer.batch_decode(translated_tokens, skip_special_tokens=True)
107
        assert self.tgt_text == decoded
108
109
110
111
112

    def test_mbart_enro_config(self):
        mbart_models = ["facebook/mbart-large-en-ro"]
        expected = {"scale_embedding": True, "output_past": True}
        for name in mbart_models:
113
            config = MBartConfig.from_pretrained(name)
114
115
116
117
118
119
120
121
122
            self.assertTrue(config.is_valid_mbart())
            for k, v in expected.items():
                try:
                    self.assertEqual(v, getattr(config, k))
                except AssertionError as e:
                    e.args += (name, k)
                    raise

    def test_mbart_fast_forward(self):
123
        config = MBartConfig(
124
125
126
127
128
129
130
131
132
133
134
            vocab_size=99,
            d_model=24,
            encoder_layers=2,
            decoder_layers=2,
            encoder_attention_heads=2,
            decoder_attention_heads=2,
            encoder_ffn_dim=32,
            decoder_ffn_dim=32,
            max_position_embeddings=48,
            add_final_layer_norm=True,
        )
135
        lm_model = MBartForConditionalGeneration(config).to(torch_device)
136
137
        context = torch.Tensor([[71, 82, 18, 33, 46, 91, 2], [68, 34, 26, 58, 30, 2, 1]]).long().to(torch_device)
        summary = torch.Tensor([[82, 71, 82, 18, 2], [58, 68, 2, 1, 1]]).long().to(torch_device)
Sylvain Gugger's avatar
Sylvain Gugger committed
138
        result = lm_model(input_ids=context, decoder_input_ids=summary, labels=summary)
139
        expected_shape = (*summary.shape, config.vocab_size)
140
        self.assertEqual(result.logits.shape, expected_shape)
141
142


143
@require_torch
144
145
@require_sentencepiece
@require_tokenizers
146
class MBartCC25IntegrationTest(AbstractSeq2SeqIntegrationTest):
147
148
149
150
151
152
153
154
155
    checkpoint_name = "facebook/mbart-large-cc25"
    src_text = [
        " UN Chief Says There Is No Military Solution in Syria",
        " I ate lunch twice yesterday",
    ]
    tgt_text = ["艦eful ONU declar膬 c膬 nu exist膬 o solu牛ie militar膬 卯n Siria", "to be padded"]

    @unittest.skip("This test is broken, still generates english")
    def test_cc25_generate(self):
156
        inputs = self.tokenizer.prepare_seq2seq_batch([self.src_text[0]]).to(torch_device)
157
158
159
160
161
162
        translated_tokens = self.model.generate(
            input_ids=inputs["input_ids"].to(torch_device),
            decoder_start_token_id=self.tokenizer.lang_code_to_id["ro_RO"],
        )
        decoded = self.tokenizer.batch_decode(translated_tokens, skip_special_tokens=True)
        self.assertEqual(self.tgt_text[0], decoded[0])
163
164
165

    @slow
    def test_fill_mask(self):
166
        inputs = self.tokenizer.prepare_seq2seq_batch(["One of the best <mask> I ever read!"]).to(torch_device)
167
168
169
170
171
172
173
        outputs = self.model.generate(
            inputs["input_ids"], decoder_start_token_id=self.tokenizer.lang_code_to_id["en_XX"], num_beams=1
        )
        prediction: str = self.tokenizer.batch_decode(
            outputs, clean_up_tokenization_spaces=True, skip_special_tokens=True
        )[0]
        self.assertEqual(prediction, "of the best books I ever read!")