# 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, 9, 9, 75), 'argparse.ArgumentParser', 'argparse.ArgumentParser', (), '', False, 'import argparse\n'), ((26, 8, 26, 45), 'importlib.import_module', 'import_module', ({(26, 22, 26, 44): "'models.' + model_name"}, {}), "('models.' + model_name)", False, 'from importlib import import_module\n'), ((28, 13, 28, 28), 'config.Config', 'Config', ({(28, 20, 28, 27): 'dataset'}, {}), '(dataset)', False, 'from config import Config\n'), ((29, 4, 29, 21), 'numpy.random.seed', 'np.random.seed', ({(29, 19, 29, 20): '(1)'}, {}), '(1)', True, 'import numpy as np\n'), ((30, 4, 30, 24), 'torch.manual_seed', 'torch.manual_seed', ({(30, 22, 30, 23): '(1)'}, {}), '(1)', False, 'import torch\n'), ((31, 4, 31, 33), 'torch.cuda.manual_seed_all', 'torch.cuda.manual_seed_all', ({(31, 31, 31, 32): '(1)'}, {}), '(1)', False, 'import torch\n'), ((34, 17, 34, 28), 'time.time', 'time.time', ({}, {}), '()', False, 'import time\n'), ((36, 45, 36, 77), 'utils.build_dataset', 'build_dataset', ({(36, 59, 36, 65): 'config', (36, 67, 36, 76): 'args.word'}, {}), '(config, args.word)', False, 'from utils import build_dataset, build_iterator, get_time_dif\n'), ((37, 17, 37, 51), 'utils.build_iterator', 'build_iterator', ({(37, 32, 37, 42): 'train_data', (37, 44, 37, 50): 'config'}, {}), '(train_data, config)', False, 'from utils import build_dataset, build_iterator, get_time_dif\n'), ((38, 15, 38, 47), 'utils.build_iterator', 'build_iterator', ({(38, 30, 38, 38): 'dev_data', (38, 40, 38, 46): 'config'}, {}), '(dev_data, config)', False, 'from utils import build_dataset, build_iterator, get_time_dif\n'), ((39, 16, 39, 49), 'utils.build_iterator', 'build_iterator', ({(39, 31, 39, 40): 'test_data', (39, 42, 39, 48): 'config'}, {}), '(test_data, config)', False, 'from utils import build_dataset, build_iterator, get_time_dif\n'), ((40, 15, 40, 39), 'utils.get_time_dif', 'get_time_dif', ({(40, 28, 40, 38): 'start_time'}, {}), '(start_time)', False, 'from utils import build_dataset, build_iterator, get_time_dif\n'), ((46, 4, 46, 23), 'train_eval.init_network', 'init_network', ({(46, 17, 46, 22): 'model'}, {}), '(model)', False, 'from train_eval import train, init_network\n'), ((48, 4, 48, 57), 'train_eval.train', 'train', ({(48, 10, 48, 16): 'config', (48, 18, 48, 23): 'model', (48, 25, 48, 35): 'train_iter', (48, 37, 48, 45): 'dev_iter', (48, 47, 48, 56): 'test_iter'}, {}), '(config, model, train_iter, dev_iter, test_iter)', False, 'from train_eval import train, init_network\n')]