test_summarization_examples.py 9.54 KB
Newer Older
1
import argparse
2
import logging
3
import os
4
5
6
7
8
9
import sys
import tempfile
import unittest
from pathlib import Path
from unittest.mock import patch

10
import torch
11
12
13
from torch.utils.data import DataLoader

from transformers import BartTokenizer
14

15
from .distillation import distill_main, evaluate_checkpoint
16
from .finetune import main
17
18
from .run_eval import generate_summaries, run_generate
from .utils import SummarizationDataset, lmap, pickle_load
19
20
21
22
23


logging.basicConfig(level=logging.DEBUG)

logger = logging.getLogger()
24
25
26
FP16_EVER = False
CHEAP_ARGS = {
    "logger": "default",
27
    "num_workers": 2,
28
29
30
31
32
33
34
35
    "alpha_hid": 0,
    "freeze_embeds": True,
    "enc_only": False,
    "tgt_suffix": "",
    "resume_from_checkpoint": None,
    "sortish_sampler": True,
    "student_decoder_layers": 1,
    "val_check_interval": 1.0,
36
37
    "output_dir": "",
    "fp16": False,
38
    "no_teacher": False,
39
    "fp16_opt_level": "O1",
40
    "gpus": 1 if torch.cuda.is_available() else 0,
41
42
43
    "n_tpu_cores": 0,
    "max_grad_norm": 1.0,
    "do_train": True,
44
    "do_predict": True,
45
46
47
48
49
50
51
    "gradient_accumulation_steps": 1,
    "server_ip": "",
    "server_port": "",
    "seed": 42,
    "model_type": "bart",
    "model_name_or_path": "sshleifer/bart-tiny-random",
    "config_name": "",
52
    "tokenizer_name": "facebook/bart-large",
53
54
55
56
57
58
59
60
61
62
63
    "cache_dir": "",
    "do_lower_case": False,
    "learning_rate": 3e-05,
    "weight_decay": 0.0,
    "adam_epsilon": 1e-08,
    "warmup_steps": 0,
    "num_train_epochs": 1,
    "train_batch_size": 2,
    "eval_batch_size": 2,
    "max_source_length": 12,
    "max_target_length": 12,
64
65
66
67
68
69
70
71
72
73
74
    "val_max_target_length": 12,
    "test_max_target_length": 12,
    "fast_dev_run": False,
    "no_cache": False,
    "n_train": -1,
    "n_val": -1,
    "n_test": -1,
    "student_encoder_layers": 1,
    "alpha_loss_encoder": 0.0,
    "freeze_encoder": False,
    "auto_scale_batch_size": False,
75
76
}

77

78
79
80
81
82
def _dump_articles(path: Path, articles: list):
    with path.open("w") as f:
        f.write("\n".join(articles))


83
84
MSG = "T5 is broken at the moment"
T5_TINY = "patrickvonplaten/t5-tiny-random"
85
86


87
88
89
90
91
92
93
94
95
96
def make_test_data_dir():
    tmp_dir = Path(tempfile.gettempdir())
    articles = [" Sam ate lunch today", "Sams lunch ingredients"]
    summaries = ["A very interesting story about what I ate for lunch.", "Avocado, celery, turkey, coffee"]
    for split in ["train", "val", "test"]:
        _dump_articles((tmp_dir / f"{split}.source"), articles)
        _dump_articles((tmp_dir / f"{split}.target"), summaries)
    return tmp_dir


97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
class TestSummarizationDistiller(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        logging.disable(logging.CRITICAL)  # remove noisy download output from tracebacks
        return cls

    @unittest.skipUnless(torch.cuda.device_count() > 1, "skipping multiGPU test")
    def test_bdc_multigpu(self):
        updates = dict(
            student_encoder_layers=2,
            student_decoder_layers=1,
            no_teacher=True,
            freeze_encoder=True,
            gpus=2,
            sortish_sampler=False,
            fp16_opt_level="O1",
            fp16=FP16_EVER,
        )
115
        self._bart_distiller_cli(updates)
116

117
    def test_bdc_t5_train(self):
118
119
        updates = dict(
            fp16=FP16_EVER,
120
            gpus=1 if torch.cuda.is_available() else 0,
121
            model_type="t5",
122
            model_name_or_path=T5_TINY,
123
124
            do_train=True,
            do_predict=True,
125
            tokenizer_name=T5_TINY,
126
            no_teacher=True,
127
            alpha_hid=2.0,
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
        )
        self._bart_distiller_cli(updates)

    def test_bdc_no_teacher(self):
        updates = dict(student_encoder_layers=2, student_decoder_layers=1, no_teacher=True,)
        self._bart_distiller_cli(updates)

    def test_bdc_yes_teacher(self):
        updates = dict(student_encoder_layers=2, student_decoder_layers=1,)
        self._bart_distiller_cli(updates)

    def test_bdc_checkpointing(self):
        updates = dict(
            student_encoder_layers=2,
            student_decoder_layers=1,
            num_train_epochs=4,
            val_check_interval=0.25,
            alpha_hid=2.0,
        )
        model = self._bart_distiller_cli(updates, check_contents=False)

        ckpts = list(Path(model.output_dir).glob("*.ckpt"))
        self.assertEqual(1, len(ckpts))
        transformer_ckpts = list(Path(model.output_dir).glob("**/*.bin"))
        self.assertEqual(len(transformer_ckpts), len(ckpts))
        new_transformer_ckpts = list(Path(model.output_dir).glob("**/*.bin"))
        self.assertEqual(len(new_transformer_ckpts), 1)
        examples = lmap(str.strip, model.hparams.data_dir.joinpath("test.source").open().readlines())
        out_path = tempfile.mktemp()
        generate_summaries(examples, out_path, new_transformer_ckpts[0].parent)
        self.assertTrue(Path(out_path).exists())

        evaluate_checkpoint(ckpts[0], dest_dir=Path(tempfile.mkdtemp()))

    def _bart_distiller_cli(self, updates, check_contents=True):
        default_updates = dict(
            train_batch_size=1,
            eval_batch_size=2,
            num_train_epochs=2,
            alpha_mlm=0.2,
            alpha_ce=0.8,
            do_predict=True,
            gpus=1 if torch.cuda.is_available() else 0,
            model_name_or_path="sshleifer/tinier_bart",
            teacher=CHEAP_ARGS["model_name_or_path"],
            val_check_interval=0.5,
            alpha_encoder_loss=0.4,
        )
        default_updates.update(updates)
        args_d: dict = CHEAP_ARGS.copy()
        tmp_dir = make_test_data_dir()
        output_dir = tempfile.mkdtemp(prefix="output_")

        args_d.update(data_dir=tmp_dir, output_dir=output_dir, **default_updates)
        model = distill_main(argparse.Namespace(**args_d))
        if not check_contents:
            return model
        contents = os.listdir(output_dir)
        ckpt_name = "val_avg_rouge2=0.0000-step_count=2.ckpt"  # "val_avg_rouge2=0.0000-epoch=1.ckpt"  # "epoch=1-val_avg_rouge2=0.0000.ckpt"
        contents = {os.path.basename(p) for p in contents}
        self.assertIn(ckpt_name, contents)
        self.assertIn("metrics.pkl", contents)
        self.assertIn("test_generations.txt", contents)
191
192
        self.assertIn("val_generations_00001.txt", contents)
        self.assertIn("val_results_00001.txt", contents)
193
194
195
        self.assertIn("test_results.txt", contents)

        metrics = pickle_load(Path(output_dir) / "metrics.pkl")
196
197
198
        desired_n_evals = int(args_d["num_train_epochs"] * (1 / args_d["val_check_interval"]) + 1)
        self.assertEqual(len(metrics["val"]), desired_n_evals)
        self.assertEqual(len(metrics["train"]), 0)  # doesn't get logged here
199
200
201
        return model


202
class TestBartExamples(unittest.TestCase):
203
204
    @classmethod
    def setUpClass(cls):
205
206
        stream_handler = logging.StreamHandler(sys.stdout)
        logger.addHandler(stream_handler)
207
208
209
210
        logging.disable(logging.CRITICAL)  # remove noisy download output from tracebacks
        return cls

    def test_bart_cnn_cli(self):
211
        tmp = Path(tempfile.gettempdir()) / "utest_generations_bart_sum.hypo"
212
        output_file_name = Path(tempfile.gettempdir()) / "utest_output_bart_sum.hypo"
213
214
        articles = [" New York (CNN)When Liana Barrientos was 23 years old, she got married in Westchester County."]
        _dump_articles(tmp, articles)
215
        testargs = ["run_eval.py", str(tmp), str(output_file_name), "sshleifer/bart-tiny-random"]
216
        with patch.object(sys, "argv", testargs):
217
            run_generate()
218
219
220
            self.assertTrue(Path(output_file_name).exists())
            os.remove(Path(output_file_name))

221
    def test_t5_run_sum_cli(self):
222
223
        args_d: dict = CHEAP_ARGS.copy()

224
225
226
227
        tmp_dir = make_test_data_dir()
        output_dir = tempfile.mkdtemp(prefix="output_")
        args_d.update(
            data_dir=tmp_dir,
228
229
            model_name_or_path=T5_TINY,
            tokenizer_name=None,  # T5_TINY,
230
231
            train_batch_size=2,
            eval_batch_size=2,
232
            gpus=0,
233
234
235
            output_dir=output_dir,
            do_predict=True,
        )
236
237
238
        assert "n_train" in args_d
        args = argparse.Namespace(**args_d)
        main(args)
239
240
241
242
243
244
245

    def test_bart_summarization_dataset(self):
        tmp_dir = Path(tempfile.gettempdir())
        articles = [" Sam ate lunch today", "Sams lunch ingredients"]
        summaries = ["A very interesting story about what I ate for lunch.", "Avocado, celery, turkey, coffee"]
        _dump_articles((tmp_dir / "train.source"), articles)
        _dump_articles((tmp_dir / "train.target"), summaries)
246
        tokenizer = BartTokenizer.from_pretrained("facebook/bart-large")
247
248
249
250
251
252
253
254
        max_len_source = max(len(tokenizer.encode(a)) for a in articles)
        max_len_target = max(len(tokenizer.encode(a)) for a in summaries)
        trunc_target = 4
        train_dataset = SummarizationDataset(
            tokenizer, data_dir=tmp_dir, type_path="train", max_source_length=20, max_target_length=trunc_target,
        )
        dataloader = DataLoader(train_dataset, batch_size=2, collate_fn=train_dataset.collate_fn)
        for batch in dataloader:
255
            self.assertEqual(batch["attention_mask"].shape, batch["input_ids"].shape)
256
            # show that articles were trimmed.
257
258
            self.assertEqual(batch["input_ids"].shape[1], max_len_source)
            self.assertGreater(20, batch["input_ids"].shape[1])  # trimmed significantly
259
260

            # show that targets were truncated
261
            self.assertEqual(batch["decoder_input_ids"].shape[1], trunc_target)  # Truncated
262
            self.assertGreater(max_len_target, trunc_target)  # Truncated
263
264


265
266
267
def list_to_text_file(lst, path):
    dest = Path(path)
    dest.open("w+").writelines(lst)