#!python
from dialogpt2 import DialoGPT2

if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser()
    group = parser.add_mutually_exclusive_group(required=True)   
    group.add_argument('--question', '-q', type=str, default=None, help='A single question.')
    group.add_argument('--questions-file', '-qf', default=None, type=str, help='Questions separated by a line break character.')
    group.add_argument('--interactive', '-i', action='store_true', help='Talk with it!')
    parser.add_argument('--num_outputs', '-n', type=int, default=1)
    parser.add_argument('--cuda-device', '-c', type=int, default=None)
    parser.add_argument('--model_name_or_path', '-m', type=str, default="microsoft/DialoGPT-small")
    parser.add_argument('--use_context', '-uc', action='store_true')
    parser.add_argument('--reset_context_on_iter', '-rcoi', action='store_true', help='When num_outputs is bigger than 1 reset context on each output batch')

    args = parser.parse_args()
    dgpt = DialoGPT2(args.model_name_or_path, cuda_device=args.cuda_device, use_context=args.use_context)        
    if args.interactive:
        while True:
            print("Bot >> " + dgpt.gen(input("User >> ")))       
    elif args.question is not None:
        questions = [args.question]
    else:
        questions = [l.strip('\n') for l in open(args.questions_file, 'r').readlines()]

    for i in range(args.num_outputs):
        print("\n------- Output #" + str(i) + " ----------")
        for q in questions:
            print("Question: " + q)
            print("Answer: " + dgpt.gen(q))
            if args.reset_context_on_iter:
                dgpt.reset_context()

