# coding: UTF-8
import time
import torch
import numpy as np
from train_eval import train, init_network
from importlib import import_module
import argparse

parser = argparse.ArgumentParser(description='Chinese Text Classification')
parser.add_argument('--model', type=str, required=True, help='choose a model: TextCNN')
parser.add_argument('--embedding', default='pre_trained', type=str, help='random or pre_trained')
parser.add_argument('--word', default=False, type=bool, help='True for word, False for char')
args = parser.parse_args()


if __name__ == '__main__':
    dataset = 'THUCNews'  # 数据集

    # 搜狗新闻:embedding_SougouNews.npz, 腾讯:embedding_Tencent.npz, 随机初始化:random

    # embedding = 'random'
    model_name = args.model  # TextCNN

    from utils import build_dataset, build_iterator, get_time_dif

    x = import_module('models.' + model_name)
    from config import Config
    config = Config(dataset)
    np.random.seed(1)
    torch.manual_seed(1)
    torch.cuda.manual_seed_all(1)
    torch.backends.cudnn.deterministic = True  # 保证每次结果一样

    start_time = time.time()
    print("Loading data...")
    vocab, train_data, dev_data, test_data = build_dataset(config, args.word)
    train_iter = build_iterator(train_data, config)
    dev_iter = build_iterator(dev_data, config)
    test_iter = build_iterator(test_data, config)
    time_dif = get_time_dif(start_time)
    print("Time usage:", time_dif)

    # train
    config.n_vocab = len(vocab)
    model = x.Model().to(config.device)
    init_network(model)
    print(model.parameters)
    train(config, model, train_iter, dev_iter, test_iter)

[(9, 'argparse.ArgumentParser', 'argparse.ArgumentParser', False, 'import argparse\n'), (26, 'importlib.import_module', 'import_module', False, 'from importlib import import_module\n'), (28, 'config.Config', 'Config', False, 'from config import Config\n'), (29, 'numpy.random.seed', 'np.random.seed', True, 'import numpy as np\n'), (30, 'torch.manual_seed', 'torch.manual_seed', False, 'import torch\n'), (31, 'torch.cuda.manual_seed_all', 'torch.cuda.manual_seed_all', False, 'import torch\n'), (34, 'time.time', 'time.time', False, 'import time\n'), (36, 'utils.build_dataset', 'build_dataset', False, 'from utils import build_dataset, build_iterator, get_time_dif\n'), (37, 'utils.build_iterator', 'build_iterator', False, 'from utils import build_dataset, build_iterator, get_time_dif\n'), (38, 'utils.build_iterator', 'build_iterator', False, 'from utils import build_dataset, build_iterator, get_time_dif\n'), (39, 'utils.build_iterator', 'build_iterator', False, 'from utils import build_dataset, build_iterator, get_time_dif\n'), (40, 'utils.get_time_dif', 'get_time_dif', False, 'from utils import build_dataset, build_iterator, get_time_dif\n'), (46, 'train_eval.init_network', 'init_network', False, 'from train_eval import train, init_network\n'), (48, 'train_eval.train', 'train', False, 'from train_eval import train, init_network\n')]