run_eval.py 6.46 KB
Newer Older
1
import argparse
2
import datetime
3
import json
4
5
6
import time
import warnings
from logging import getLogger
7
from pathlib import Path
8
from typing import Dict, List
9
10
11
12
13

import torch
from tqdm import tqdm

from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
14
from utils import calculate_bleu, calculate_rouge, parse_numeric_n_bool_cl_kwargs, use_task_specific_params
15
16


17
18
logger = getLogger(__name__)

19
20
21
22
23
24
25
26
27
28

DEFAULT_DEVICE = "cuda" if torch.cuda.is_available() else "cpu"


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


29
def generate_summaries_or_translations(
30
    examples: List[str],
31
32
33
34
35
    out_file: str,
    model_name: str,
    batch_size: int = 8,
    device: str = DEFAULT_DEVICE,
    fp16=False,
36
    task="summarization",
37
    prefix=None,
38
39
40
    **generate_kwargs,
) -> Dict:
    """Save model.generate results to <out_file>, and return how long it took."""
41
    fout = Path(out_file).open("w", encoding="utf-8")
42
    model_name = str(model_name)
43
44
45
46
47
    model = AutoModelForSeq2SeqLM.from_pretrained(model_name).to(device)
    if fp16:
        model = model.half()

    tokenizer = AutoTokenizer.from_pretrained(model_name)
48
    logger.info(f"Inferred tokenizer type: {tokenizer.__class__}")  # if this is wrong, check config.model_type.
49

50
51
    start_time = time.time()
    # update config with task specific params
52
    use_task_specific_params(model, task)
53
54
    if prefix is None:
        prefix = prefix or getattr(model.config, "prefix", "") or ""
55
    for examples_chunk in tqdm(list(chunks(examples, batch_size))):
56
        examples_chunk = [prefix + text for text in examples_chunk]
57
        batch = tokenizer(examples_chunk, return_tensors="pt", truncation=True, padding="longest").to(device)
58
        summaries = model.generate(
59
60
61
            input_ids=batch.input_ids,
            attention_mask=batch.attention_mask,
            **generate_kwargs,
62
        )
63
64
65
66
        dec = tokenizer.batch_decode(summaries, skip_special_tokens=True, clean_up_tokenization_spaces=False)
        for hypothesis in dec:
            fout.write(hypothesis + "\n")
            fout.flush()
67
    fout.close()
68
    runtime = int(time.time() - start_time)  # seconds
69
70
    n_obs = len(examples)
    return dict(n_obs=n_obs, runtime=runtime, seconds_per_sample=round(runtime / n_obs, 4))
71
72


73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
def datetime_now():
    return datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")


def run_generate(verbose=True):
    """

    Takes input text, generates output, and then using reference calculates the BLEU scores.

    The results are saved to a file and returned to the caller, and printed out unless ``verbose=False`` is passed.

    Args:
        verbose (:obj:`bool`, `optional`, defaults to :obj:`True`): print results to stdout

    Returns:
        a tuple: ``(scores, params}``
        - ``scores``: a dict of scores data ``{'bleu': 39.6501, 'n_obs': 2000, 'runtime': 186, 'seconds_per_sample': 0.093}``
        - ``params``: a dict of custom params, e.g. ``{'num_beams': 5, 'length_penalty': 0.8}``
    """

93
94
    parser = argparse.ArgumentParser()
    parser.add_argument("model_name", type=str, help="like facebook/bart-large-cnn,t5-base, etc.")
95
96
    parser.add_argument("input_path", type=str, help="like cnn_dm/test.source")
    parser.add_argument("save_path", type=str, help="where to save summaries")
97
98
    parser.add_argument("--reference_path", type=str, required=False, help="like cnn_dm/test.target")
    parser.add_argument("--score_path", type=str, required=False, default="metrics.json", help="where to save metrics")
99
    parser.add_argument("--device", type=str, required=False, default=DEFAULT_DEVICE, help="cuda, cuda:1, cpu etc.")
100
101
102
    parser.add_argument(
        "--prefix", type=str, required=False, default=None, help="will be added to the begininng of src examples"
    )
103
    parser.add_argument("--task", type=str, default="summarization", help="used for task_specific_params + metrics")
104
    parser.add_argument("--bs", type=int, default=8, required=False, help="batch size")
105
106
107
    parser.add_argument(
        "--n_obs", type=int, default=-1, required=False, help="How many observations. Defaults to all."
    )
108
    parser.add_argument("--fp16", action="store_true")
109
110
111
112
113
114
115
116
    parser.add_argument("--dump-args", action="store_true", help="print the custom hparams with the results")
    parser.add_argument(
        "--info",
        nargs="?",
        type=str,
        const=datetime_now(),
        help="use in conjunction w/ --dump-args to print with the results whatever other info you'd like, e.g. lang=en-ru. If no value is passed, the current datetime string will be used.",
    )
117
118
    # Unspecified args like --num_beams=2 --decoder_start_token_id=4 are passed to model.generate
    args, rest = parser.parse_known_args()
119
120
121
    parsed_args = parse_numeric_n_bool_cl_kwargs(rest)
    if parsed_args and verbose:
        print(f"parsed the following generate kwargs: {parsed_args}")
122
    examples = [" " + x.rstrip() if "t5" in args.model_name else x.rstrip() for x in open(args.input_path).readlines()]
123
124
    if args.n_obs > 0:
        examples = examples[: args.n_obs]
125
    Path(args.save_path).parent.mkdir(exist_ok=True)
126
127
128
    if args.reference_path is None and Path(args.score_path).exists():
        warnings.warn(f"score_path {args.score_path} will be overwritten unless you type ctrl-c.")
    runtime_metrics = generate_summaries_or_translations(
129
130
131
132
133
134
135
        examples,
        args.save_path,
        args.model_name,
        batch_size=args.bs,
        device=args.device,
        fp16=args.fp16,
        task=args.task,
136
        prefix=args.prefix,
137
        **parsed_args,
138
    )
139

140
    if args.reference_path is None:
141
142
        return {}

143
    # Compute scores
144
    score_fn = calculate_bleu if "translation" in args.task else calculate_rouge
145
146
147
    output_lns = [x.rstrip() for x in open(args.save_path).readlines()]
    reference_lns = [x.rstrip() for x in open(args.reference_path).readlines()][: len(output_lns)]
    scores: dict = score_fn(output_lns, reference_lns)
148
    scores.update(runtime_metrics)
149
150
151
152
153
154
155

    if args.dump_args:
        scores.update(parsed_args)
    if args.info:
        scores["info"] = args.info

    if verbose:
156
        print(scores)
157

158
    if args.score_path is not None:
159
160
161
        path = args.score_path
        json.dump(scores, open(path, "w"))

162
    return scores
163
164
165


if __name__ == "__main__":
166
167
    # Usage for MT:
    # python run_eval.py MODEL_NAME $DATA_DIR/test.source $save_dir/test_translations.txt --reference_path $DATA_DIR/test.target --score_path $save_dir/test_bleu.json  --task translation $@
168
    run_generate(verbose=True)