import argparse
from termcolor import colored
from idiomifier import Idiomifier
from idiomify.fetchers import fetch_config, fetch_alpha
from transformers import BartTokenizer


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--model", type=str,
                        default="alpha")
    parser.add_argument("--ver", type=str,
                        default="overfit")
    parser.add_argument("--src", type=str,
                        default="If there's any benefits to losing my job, it's that I'll now be able to go to school full-time and finish my degree earlier.")
    args = parser.parse_args()
    config = fetch_config()[args.model][args.ver]
    config.update(vars(args))
    model = fetch_alpha(config['ver'])
    tokenizer = BartTokenizer.from_pretrained(config['bart'])
    idiomifier = Idiomifier(model, tokenizer)
    src = config['src']
    tgt = idiomifier(src=config['src'])
    print(src, "\n->", colored(tgt, "blue"))


if __name__ == '__main__':
    main()