from transformers import AutoTokenizer, pipeline, PretrainedConfig
from query2labels.infer import parser_args, Query2Label
# from Faster_VisualGenome import demo

#custokenizer = AutoTokenizer.from_pretrained("vinai/phobert-base", use_fast=False)
custokenizer = AutoTokenizer.from_pretrained("./GPT2/phoBert", use_fast=False)
# -------Load model saved-----------------#

parser = parser_args()
#parser.add_argument('--img_path', help='img path', default='./test_imgs/test.jpg')
parser.add_argument('--config', help='config file', default='./query2labels/output/config.json')
parser.add_argument('-f')
args = parser.parse_args()

vis_extractor = Query2Label(args)
# infer.main(args)


configuration = {'num_beams': 5, 'max_length': 256, "architectures": ["GPT2LMHeadModel"]}
config = PretrainedConfig()
config.from_dict(configuration)
poem = pipeline('text-generation', model="./GPT2/rkw_4sen",
                tokenizer=custokenizer,
                config=config)


def main(img):
    clses = vis_extractor.predict(img)
    
    keywords = clses
    print(keywords)
    keywords = ' '.join(keywords).replace('_', ' ')
    poem = generate_poem(keywords)
    return poem


def generate_poem(keywords):
    # Test
    input = '<s>' + keywords + ' [SEP]'
    a = poem(input)
    out = a[0]['generated_text']
    out = out.replace('<s>', '')
    out = out.replace('</s>', '')
    out = out.split('<unk>')
    return '\n'.join(out)
    
#if __name__ == '__main__':
#  print(main())