| 58 | |
| 59 | |
| 60 | def run_generate(): |
| 61 | parser = argparse.ArgumentParser() |
| 62 | parser.add_argument("model_name", type=str, help="like facebook/bart-large-cnn,t5-base, etc.") |
| 63 | parser.add_argument("input_path", type=str, help="like cnn_dm/test.source") |
| 64 | parser.add_argument("save_path", type=str, help="where to save summaries") |
| 65 | |
| 66 | parser.add_argument("--reference_path", type=str, required=False, help="like cnn_dm/test_reference_summaries.txt") |
| 67 | parser.add_argument("--score_path", type=str, required=False, help="where to save the rouge score in json format") |
| 68 | parser.add_argument("--device", type=str, required=False, default=DEFAULT_DEVICE, help="cuda, cuda:1, cpu etc.") |
| 69 | parser.add_argument("--task", type=str, default="summarization", help="typically translation or summarization") |
| 70 | parser.add_argument("--bs", type=int, default=8, required=False, help="batch size") |
| 71 | parser.add_argument( |
| 72 | "--n_obs", type=int, default=-1, required=False, help="How many observations. Defaults to all." |
| 73 | ) |
| 74 | parser.add_argument("--fp16", action="store_true") |
| 75 | args = parser.parse_args() |
| 76 | examples = [" " + x.rstrip() if "t5" in args.model_name else x.rstrip() for x in open(args.input_path).readlines()] |
| 77 | if args.n_obs > 0: |
| 78 | examples = examples[: args.n_obs] |
| 79 | |
| 80 | generate_summaries_or_translations( |
| 81 | examples, |
| 82 | args.save_path, |
| 83 | args.model_name, |
| 84 | batch_size=args.bs, |
| 85 | device=args.device, |
| 86 | fp16=args.fp16, |
| 87 | task=args.task, |
| 88 | ) |
| 89 | if args.reference_path is None: |
| 90 | return |
| 91 | # Compute scores |
| 92 | score_fn = calculate_bleu_score if "translation" in args.task else calculate_rouge |
| 93 | output_lns = [x.rstrip() for x in open(args.save_path).readlines()] |
| 94 | reference_lns = [x.rstrip() for x in open(args.reference_path).readlines()][: len(output_lns)] |
| 95 | scores: dict = score_fn(output_lns, reference_lns) |
| 96 | if args.score_path is not None: |
| 97 | json.dump(scores, open(args.score_path, "w+")) |
| 98 | return scores |
| 99 | |
| 100 | |
| 101 | if __name__ == "__main__": |