test_modeling_mbart.py 7.8 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

Sam Shleifer's avatar
Sam Shleifer committed
7
from .test_modeling_bart import TOLERANCE, _long_tensor, assert_tensors_close
8
from .test_modeling_common import ModelTesterMixin
9
10
11
12


if is_torch_available():
    import torch
13

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


EN_CODE = 250004
RO_CODE = 250020


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
56
57
@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,
            return_dict=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)


58
@require_torch
59
60
@require_sentencepiece
@require_tokenizers
61
62
class AbstractSeq2SeqIntegrationTest(unittest.TestCase):
    maxDiff = 1000  # longer string compare tracebacks
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
    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
80
81
@require_sentencepiece
@require_tokenizers
82
class MBartEnroIntegrationTest(AbstractSeq2SeqIntegrationTest):
83
84
85
86
87
88
89
    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",
90
        '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.',
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
    ]
    expected_src_tokens = [8274, 127873, 25916, 7, 8622, 2071, 438, 67485, 53, 187895, 23, 51712, 2, EN_CODE]

    @slow
    @unittest.skip("This has been failing since June 20th at least.")
    def test_enro_forward(self):
        model = self.model
        net_input = {
            "input_ids": _long_tensor(
                [
                    [3493, 3060, 621, 104064, 1810, 100, 142, 566, 13158, 6889, 5, 2, 250004],
                    [64511, 7, 765, 2837, 45188, 297, 4049, 237, 10, 122122, 5, 2, 250004],
                ]
            ),
            "decoder_input_ids": _long_tensor(
                [
                    [250020, 31952, 144, 9019, 242307, 21980, 55749, 11, 5, 2, 1, 1],
                    [250020, 884, 9019, 96, 9, 916, 86792, 36, 18743, 15596, 5, 2],
                ]
            ),
        }
112
        net_input["attention_mask"] = net_input["input_ids"].ne(1)
113
114
115
116
117
        with torch.no_grad():
            logits, *other_stuff = model(**net_input)

        expected_slice = torch.tensor([9.0078, 10.1113, 14.4787], device=logits.device, dtype=logits.dtype)
        result_slice = logits[0, 0, :3]
Sam Shleifer's avatar
Sam Shleifer committed
118
119
120
121
122
123
124
125
126
127
128
        assert_tensors_close(expected_slice, result_slice, atol=TOLERANCE)

    @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])
129
130
131

    @slow
    def test_enro_generate(self):
132
        batch: BatchEncoding = self.tokenizer.prepare_seq2seq_batch(self.src_text).to(torch_device)
133
134
        translated_tokens = self.model.generate(**batch)
        decoded = self.tokenizer.batch_decode(translated_tokens, skip_special_tokens=True)
135
        assert self.tgt_text == decoded
136
137
138
139
140

    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:
141
            config = MBartConfig.from_pretrained(name)
142
143
144
145
146
147
148
149
150
            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):
151
        config = MBartConfig(
152
153
154
155
156
157
158
159
160
161
            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,
Sylvain Gugger's avatar
Sylvain Gugger committed
162
            return_dict=True,
163
        )
164
        lm_model = MBartForConditionalGeneration(config).to(torch_device)
165
166
        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
167
        result = lm_model(input_ids=context, decoder_input_ids=summary, labels=summary)
168
        expected_shape = (*summary.shape, config.vocab_size)
169
        self.assertEqual(result.logits.shape, expected_shape)
170
171


172
@require_torch
173
174
@require_sentencepiece
@require_tokenizers
175
class MBartCC25IntegrationTest(AbstractSeq2SeqIntegrationTest):
176
177
178
179
180
181
182
183
184
    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):
185
        inputs = self.tokenizer.prepare_seq2seq_batch([self.src_text[0]]).to(torch_device)
186
187
188
189
190
191
        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])
192
193
194

    @slow
    def test_fill_mask(self):
195
        inputs = self.tokenizer.prepare_seq2seq_batch(["One of the best <mask> I ever read!"]).to(torch_device)
196
197
198
199
200
201
202
        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!")