eic_for_generative.py

import sys
from argparse import ArgumentParser


def main(args):
    if args.method == 'e2esr':
        from src.generative.e2esr.train import train as e2esr_main
        del sys.argv[sys.argv.index('--method'):sys.argv.index('--method') + 2]
        e2esr_main()
    elif args.method == 'snip':
        from src.generative.e2esr.train import train as snip_main
        try:
            assert sys.argv[sys.argv.index('--snip_loss')+1] > 0
        except:
            print("Set --snip_loss to train snip!")
        del sys.argv[sys.argv.index('--method'):sys.argv.index('--method') + 2]
        snip_main()
    elif args.method == 'sr4mdl':
        from src.generative.sr4mdl.train import train as sr4mdl_main
        del sys.argv[sys.argv.index('--method'):sys.argv.index('--method') + 2]
        sr4mdl_main()
    else:
        raise ValueError(f"Unknown method: {args.method}")

if __name__ == '__main__':
    parser = ArgumentParser()
    parser.add_argument('--method', type=str, choices=['e2esr', 'snip', 'sr4mdl'], default='e2esr')
    args, unknown = parser.parse_known_args()
    main(args)