
import torch
import numpy

class Config(object):
    """模型参数配置"""

    def __init__(self, base_path, embedding_name):

        self.vocab_path = base_path + '/vocab.pkl'
        self.train_path = base_path + '/train.txt'  # 训练集
        self.dev_path = base_path + '/dev.txt'  # 验证集
        self.test_path = base_path + '/test.txt'
        self.save_path = base_path + '/TextCNN.ckpt'  # 模型训练结果
        class_path = base_path + '/class.txt'
        self.class_list = [x.strip() for x in open(class_path, encoding='utf-8').readlines()]  # 类别名单

        # 预训练词向量
        if embedding_name == 'random':
            self.embedding_pretrained = None
        else:
            # 预训练词向量
            embed_path = base_path + '/' + embedding_name
            self.embedding_pretrained = torch.tensor(numpy.load(embed_path)["embeddings"].astype('float32'))

        self.embedding_size = 300
        if self.embedding_pretrained is not None:
            self.embedding_size = self.embedding_pretrained.size(1)
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

        self.dropout = 0.5  # 随机失活
        self.require_improvement = 1000  # 若超过1000batch效果还没提升，则提前结束训练
        self.num_classes = len(self.class_list)  # 类别数
        self.n_vocab = 0  # 词表大小，在运行时赋值
        self.num_epochs = 20  # epoch数
        self.batch_size = 128  # mini-batch大小
        self.pad_size = 32  # 每句话处理成的长度(短填长切)
        self.learning_rate = 1e-3  # 学习率
        self.filter_sizes = (2, 3, 4)  # 卷积核尺寸

        self.hidden_size = 256
        self.n_gram_vocab = 250499


class Model(torch.nn.Module):

    def __init__(self, config : Config):
        super(Model, self).__init__()
        if config.embedding_pretrained is None:
            self.embedding = torch.nn.Embedding(config.n_vocab, config.embedding_size, padding_idx=config.n_vocab-1)
        else:
            self.embedding = torch.nn.Embedding.from_pretrained(config.embedding_pretrained, freeze=False)
        self.embedding_ngram2 = torch.nn.Embedding(config.n_gram_vocab, config.embedding_size)
        self.embedding_ngram3 = torch.nn.Embedding(config.n_gram_vocab, config.embedding_size)
        self.dropout = torch.nn.Dropout(config.dropout)
        self.full_connect1 = torch.nn.Linear(config.embedding_size * 3, config.hidden_size)
        self.full_connect2 = torch.nn.Linear(config.hidden_size, config.num_classes)

    def forward(self, x):
        outword = self.embedding(x[0])
        out_bigram = self.embedding_ngram2(x[2])
        out_trigram = self.embedding_ngram3(x[3])
        out = torch.cat((outword, out_bigram, out_trigram), -1)

        out = out.mean(dim = 1)
        out = self.dropout(out)
        out = self.full_connect1(out)
        out = torch.nn.functional.relu(out)
        out = self.full_connect2(out)
        return out