evaluate_wmt.py 3.58 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
import argparse
from pathlib import Path

import torch
from tqdm import tqdm

from sacrebleu import corpus_bleu
from transformers import T5ForConditionalGeneration, T5Tokenizer


def chunks(lst, n):
    """Yield successive n-sized chunks from lst."""
    for i in range(0, len(lst), n):
        yield lst[i : i + n]


17
18
def generate_translations(lns, output_file_path, model_size, batch_size, device):
    model = T5ForConditionalGeneration.from_pretrained(model_size)
19
20
    model.to(device)

21
    tokenizer = T5Tokenizer.from_pretrained(model_size)
22
23
24
25
26
27

    # update config with summarization specific params
    task_specific_params = model.config.task_specific_params
    if task_specific_params is not None:
        model.config.update(task_specific_params.get("translation_en_to_de", {}))

28
29
30
    with Path(output_file_path).open("w") as output_file:
        for batch in tqdm(list(chunks(lns, batch_size))):
            batch = [model.config.prefix + text for text in batch]
31

32
            dct = tokenizer.batch_encode_plus(batch, max_length=512, return_tensors="pt", pad_to_max_length=True)
33

34
35
            input_ids = dct["input_ids"].to(device)
            attention_mask = dct["attention_mask"].to(device)
36

37
38
39
40
            translations = model.generate(input_ids=input_ids, attention_mask=attention_mask)
            dec = [
                tokenizer.decode(g, skip_special_tokens=True, clean_up_tokenization_spaces=False) for g in translations
            ]
41

42
43
            for hypothesis in dec:
                output_file.write(hypothesis + "\n")
44
45
46
47
48


def calculate_bleu_score(output_lns, refs_lns, score_path):
    bleu = corpus_bleu(output_lns, [refs_lns])
    result = "BLEU score: {}".format(bleu.score)
49
50
    with Path(score_path).open("w") as score_file:
        score_file.write(result)
51
52
53
54


def run_generate():
    parser = argparse.ArgumentParser()
55
56
57
58
59
60
    parser.add_argument(
        "model_size",
        type=str,
        help="T5 model size, either 't5-small', 't5-base', 't5-large', 't5-3b', 't5-11b'. Defaults to 't5-base'.",
        default="t5-base",
    )
61
    parser.add_argument(
62
        "input_path", type=str, help="like wmt/newstest2014.en",
63
64
65
66
67
    )
    parser.add_argument(
        "output_path", type=str, help="where to save translation",
    )
    parser.add_argument(
68
        "reference_path", type=str, help="like wmt/newstest2014.de",
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
    )
    parser.add_argument(
        "score_path", type=str, help="where to save the bleu score",
    )
    parser.add_argument(
        "--batch_size", type=int, default=16, required=False, help="batch size: how many to summarize at a time",
    )
    parser.add_argument(
        "--no_cuda", default=False, type=bool, help="Whether to force the execution on CPU.",
    )

    args = parser.parse_args()
    args.device = torch.device("cuda" if torch.cuda.is_available() and not args.no_cuda else "cpu")

    dash_pattern = (" ##AT##-##AT## ", "-")

85
86
87
    # Read input lines into python
    with open(args.input_path, "r") as input_file:
        input_lns = [x.strip().replace(dash_pattern[0], dash_pattern[1]) for x in input_file.readlines()]
88

89
    generate_translations(input_lns, args.output_path, args.model_size, args.batch_size, args.device)
90

91
92
93
94
95
96
97
    # Read generated lines into python
    with open(args.output_path, "r") as output_file:
        output_lns = [x.strip() for x in output_file.readlines()]

    # Read reference lines into python
    with open(args.reference_path, "r") as reference_file:
        refs_lns = [x.strip().replace(dash_pattern[0], dash_pattern[1]) for x in reference_file.readlines()]
98
99
100
101
102
103

    calculate_bleu_score(output_lns, refs_lns, args.score_path)


if __name__ == "__main__":
    run_generate()