import sys

import argparse

from summarizer_test import summarizertest

def main(args):

    parser = argparse.ArgumentParser(description='To get arguments')

    parser.add_argument('--generator_max_length',

                        nargs='*',

                        type=int,

                        default=[80, 110],

                        help='Maximum length of generator text',

                        required=False)

    parser.add_argument('--generator_min_length',

                        nargs='*',

                        type=int,

                        default=[10, 20],

                        help='Minimum length of generator text',

                        required=False)

    parser.add_argument('--generator_top_k',

                        nargs='*',

                        type=int,

                        default=[50, 100],

                        help='top_k of generator model',

                        required=False)

    parser.add_argument('--generator_length_penalty',

                        nargs='*',

                        type=int,

                        default=[0.8, 1],

                        help='length_penalty of generator model',

                        required=False)

    parser.add_argument('--generator_no_repeat_ngram_size',

                        nargs='*',

                        type=int,

                        default=[2, 3],

                        help='no_repeat_ngram_size of generator model',

                        required=False)

    parser.add_argument('--generator_sequences',

                        nargs='*',

                        type=int,

                        default=[2],

                        help='count of sequences that model returned',

                        required=False)

    parser.add_argument('--summarizer_models',

                        nargs='*',

                        type=str,

                        default=["sshleifer/distilbart-cnn-12-6"],

                        help='Models that want test in data',

                        required=False)

    parser.add_argument('--report_out_dir',

                        type=str,

                        default='summarizer-report.json',

                        help='The report of run all models with diffrent parameters',

                        required=False)

    parser.add_argument('--data_dir',

                        type=str,

                        default='summarizer-context.json',

                        help='dataset that include long text for summarization',

                        required=False)

    args = parser.parse_args()

    report_generator = summarizertest.ReportGenerator(

        models_names=args.summarizer_models,

        val_contexts_path=args.data_dir,

        report_path=args.report_out_dir,

        max_lengths=args.generator_max_length,

        min_lengths=args.generator_min_length,

        top_k=args.generator_top_k,

        penalty_l=args.generator_length_penalty,

        no_repeat_ngram_size=args.generator_no_repeat_ngram_size,

        num_return_sequences=args.generator_sequences,

        )

    report_generator.get_report()

if __name__ == "__main__":

    main(sys.argv[1:])