#!/usr/bin/env python
# -*- coding:utf-8 -*-
# author:roger
# datetime:19-11-12 下午4:59
# software: PyCharm

import pickle

import tensorflow as tf

from loader import input_from_line
from model import Model
from train import FLAGS, load_config
from utils import create_model, get_logger


def main(_):
    config = load_config(FLAGS.config_file)
    logger = get_logger(FLAGS.log_file)
    # limit GPU memory
    tf_config = tf.ConfigProto()
    tf_config.gpu_options.allow_growth = True
    with open(FLAGS.map_file, "rb") as f:
        tag_to_id, id_to_tag = pickle.load(f)
    with tf.Session(config=tf_config) as sess:
        model = create_model(sess, Model, FLAGS.ckpt_path, config, logger)
        while True:
            line = input("input sentence, please:")
            result = model.evaluate_line(sess, input_from_line(line, FLAGS.max_seq_len, tag_to_id), id_to_tag)
            print(result['entities'])


if __name__ == '__main__':
    # os.environ['CUDA_VISIBLE_DEVICES'] = '0'
    tf.app.run(main)