import copyimport mathimport matplotlib.pyplot as pltimport numpy as npimport osimport seaborn as snsimport timeimport torchimport torch.nn as nnimport torch.nn.functional as F
from collections import Counterfrom langconv import Converterfrom nltk import word_tokenizefrom torch.autograd import Variable
def cht_to_chs(sent):    sent = Converter("zh-hans").convert(sent)    sent.encode("utf-8")    return sent

class PrepareData:    def __init__(self, train_file, dev_file):        # 读取数据、分词        self.train_en, self.train_cn = self.load_data(train_file)        self.dev_en, self.dev_cn = self.load_data(dev_file)        # 构建词表        self.en_word_dict, self.en_total_words, self.en_index_dict = \            self.build_dict(self.train_en)        self.cn_word_dict, self.cn_total_words, self.cn_index_dict = \            self.build_dict(self.train_cn)        # 单词映射为索引        self.train_en, self.train_cn = self.word2id(self.train_en, self.train_cn, self.en_word_dict, self.cn_word_dict)        self.dev_en, self.dev_cn = self.word2id(self.dev_en, self.dev_cn, self.en_word_dict, self.cn_word_dict)        # 划分批次、填充、掩码        self.train_data = self.split_batch(self.train_en, self.train_cn, BATCH_SIZE)        self.dev_data = self.split_batch(self.dev_en, self.dev_cn, BATCH_SIZE)
    def load_data(self, path):        """        读取英文、中文数据        对每条样本分词并构建包含起始符和终止符的单词列表        形式如：en = [['BOS', 'i', 'love', 'you', 'EOS'], ['BOS', 'me', 'too', 'EOS'], ...]                cn = [['BOS', '我', '爱', '你', 'EOS'], ['BOS', '我', '也', '是', 'EOS'], ...]        """        en = []        cn = []        with open(path, mode="r", encoding="utf-8") as f:            for line in f.readlines():                sent_en, sent_cn = line.strip().split("\t")                sent_en = sent_en.lower()                sent_cn = cht_to_chs(sent_cn)                sent_en = ["BOS"] + word_tokenize(sent_en) + ["EOS"]                # 中文按字符切分                sent_cn = ["BOS"] + [char for char in sent_cn] + ["EOS"]                en.append(sent_en)                cn.append(sent_cn)        return en, cn
    def build_dict(self, sentences, max_words=5e4):        """        构造分词后的列表数据        构建单词-索引映射（key为单词，value为id值）        """        # 统计数据集中单词词频        word_count = Counter([word for sent in sentences for word in sent])        # 按词频保留前max_words个单词构建词典        # 添加UNK和PAD两个单词        ls = word_count.most_common(int(max_words))        total_words = len(ls) + 2        word_dict = {w[0]: index + 2 for index, w in enumerate(ls)}        word_dict['UNK'] = UNK        word_dict['PAD'] = PAD        # 构建id2word映射        index_dict = {v: k for k, v in word_dict.items()}        return word_dict, total_words, index_dict
    def word2id(self, en, cn, en_dict, cn_dict, sort=True):        """        将英文、中文单词列表转为单词索引列表        `sort=True`表示以英文语句长度排序，以便按批次填充时，同批次语句填充尽量少        """        length = len(en)        # 单词映射为索引        out_en_ids = [[en_dict.get(word, UNK) for word in sent] for sent in en]        out_cn_ids = [[cn_dict.get(word, UNK) for word in sent] for sent in cn]
        # 按照语句长度排序        def len_argsort(seq):            """            传入一系列语句数据(分好词的列表形式)，            按照语句长度排序后，返回排序后原来各语句在数据中的索引下标            """            return sorted(range(len(seq)), key=lambda x: len(seq[x]))
        # 按相同顺序对中文、英文样本排序        if sort:            # 以英文语句长度排序            sorted_index = len_argsort(out_en_ids)            out_en_ids = [out_en_ids[idx] for idx in sorted_index]            out_cn_ids = [out_cn_ids[idx] for idx in sorted_index]        return out_en_ids, out_cn_ids
    def split_batch(self, en, cn, batch_size, shuffle=True):        """        划分批次        `shuffle=True`表示对各批次顺序随机打乱        """        # 每隔batch_size取一个索引作为后续batch的起始索引        idx_list = np.arange(0, len(en), batch_size)        # 起始索引随机打乱        if shuffle:            np.random.shuffle(idx_list)        # 存放所有批次的语句索引        batch_indexs = []        for idx in idx_list:            """            形如[array([4, 5, 6, 7]),                  array([0, 1, 2, 3]),                  array([8, 9, 10, 11]),                 ...]            """            # 起始索引最大的批次可能发生越界，要限定其索引            batch_indexs.append(np.arange(idx, min(idx + batch_size, len(en))))        # 构建批次列表        batches = []        for batch_index in batch_indexs:            # 按当前批次的样本索引采样            batch_en = [en[index] for index in batch_index]            batch_cn = [cn[index] for index in batch_index]            # 对当前批次中所有语句填充、对齐长度            # 维度为：batch_size * 当前批次中语句的最大长度            batch_cn = seq_padding(batch_cn)            batch_en = seq_padding(batch_en)            # 将当前批次添加到批次列表            # Batch类用于实现注意力掩码            batches.append(Batch(batch_en, batch_cn))        return batches

class Embeddings(nn.Module):    def __init__(self, d_model, vocab):        super(Embeddings, self).__init__()        # Embedding层        self.lut = nn.Embedding(vocab, d_model)        # Embedding维数        self.d_model = d_model
    def forward(self, x):        # 返回x的词向量（需要乘以math.sqrt(d_model)）        return self.lut(x) * math.sqrt(self.d_model)

class PositionalEncoding(nn.Module):    def __init__(self, d_model, dropout, max_len=5000):        super(PositionalEncoding, self).__init__()        self.dropout = nn.Dropout(p=dropout)        # 位置编码矩阵，维度[max_len, embedding_dim]        pe = torch.zeros(max_len, d_model, device=DEVICE)        # 单词位置        position = torch.arange(0.0, max_len, device=DEVICE)        position.unsqueeze_(1)        # 使用exp和log实现幂运算        div_term = torch.exp(torch.arange(0.0, d_model, 2, device=DEVICE) * (- math.log(1e4) / d_model))        div_term.unsqueeze_(0)        # 计算单词位置沿词向量维度的纹理值        pe[:, 0 : : 2] = torch.sin(torch.mm(position, div_term))        pe[:, 1 : : 2] = torch.cos(torch.mm(position, div_term))        # 增加批次维度，[1, max_len, embedding_dim]        pe.unsqueeze_(0)        # 将位置编码矩阵注册为buffer(不参加训练)        self.register_buffer('pe', pe)
    def forward(self, x):        # 将一个批次中语句所有词向量与位置编码相加        # 注意，位置编码不参与训练，因此设置requires_grad=False        x += Variable(self.pe[:, : x.size(1), :], requires_grad=False)        return self.dropout(x)

def attention(query, key, value, mask=None, dropout=None):    """    Scaled Dot-Product Attention    """    # q、k、v向量长度为d_k    d_k = query.size(-1)    # 矩阵乘法实现q、k点积注意力，sqrt(d_k)归一化    scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k)    # 注意力掩码机制    if mask is not None:        scores = scores.masked_fill(mask==0, -1e9)    # 注意力矩阵softmax归一化    p_attn = F.softmax(scores, dim=-1)    # dropout    if dropout is not None:        p_attn = dropout(p_attn)    # 注意力对v加权    return torch.matmul(p_attn, value), p_attn

def clones(module, N):    """    克隆基本单元，克隆的单元之间参数不共享    """    return nn.ModuleList([        copy.deepcopy(module) for _ in range(N)    ])

class MultiHeadedAttention(nn.Module):    """    Multi-Head Attention    """    def __init__(self, h, d_model, dropout=0.1):        super(MultiHeadedAttention, self).__init__()        """        `h`：注意力头的数量        `d_model`：词向量维数        """        # 确保整除        assert d_model % h == 0        # q、k、v向量维数        self.d_k = d_model // h        # 头的数量        self.h = h        # WQ、WK、WV矩阵及多头注意力拼接变换矩阵WO        self.linears = clones(nn.Linear(d_model, d_model), 4)        self.attn = None        self.dropout = nn.Dropout(p=dropout)
    def forward(self, query, key, value, mask=None):        if mask is not None:            mask = mask.unsqueeze(1)        # 批次大小        nbatches = query.size(0)        # WQ、WK、WV分别对词向量线性变换，并将结果拆成h块        query, key, value = [            l(x).view(nbatches, -1, self.h, self.d_k).transpose(1, 2)            for l, x in zip(self.linears, (query, key, value))        ]        # 注意力加权        x, self.attn = attention(query, key, value, mask=mask, dropout=self.dropout)        # 多头注意力加权拼接        x = x.transpose(1, 2).contiguous().view(nbatches, -1, self.h * self.d_k)        # 对多头注意力加权拼接结果线性变换        return self.linears[-1](x)

class SublayerConnection(nn.Module):    """    通过层归一化和残差连接，连接Multi-Head Attention和Feed Forward    """    def __init__(self, size, dropout):        super(SublayerConnection, self).__init__()        self.norm = LayerNorm(size)        self.dropout = nn.Dropout(dropout)
    def forward(self, x, sublayer):        # 层归一化        x_ = self.norm(x)        x_ = sublayer(x_)        x_ = self.dropout(x_)        # 残差连接        return x + x_

class PositionwiseFeedForward(nn.Module):    def __init__(self, d_model, d_ff, dropout=0.1):        super(PositionwiseFeedForward, self).__init__()        self.w_1 = nn.Linear(d_model, d_ff)        self.w_2 = nn.Linear(d_ff, d_model)        self.dropout = nn.Dropout(dropout)
    def forward(self, x):        x = self.w_1(x)        x = F.relu(x)        x = self.dropout(x)        x = self.w_2(x)        return x

class EncoderLayer(nn.Module):    def __init__(self, size, self_attn, feed_forward, dropout):        super(EncoderLayer, self).__init__()        self.self_attn = self_attn        self.feed_forward = feed_forward        # SublayerConnection作用连接multi和ffn        self.sublayer = clones(SublayerConnection(size, dropout), 2)        # d_model        self.size = size
    def forward(self, x, mask):        # 将embedding层进行Multi head Attention        x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, mask))        # attn的结果直接作为下一层输入        return self.sublayer[1](x, self.feed_forward)

class Encoder(nn.Module):    def __init__(self, layer, N):        """        layer = EncoderLayer        """        super(Encoder, self).__init__()        # 复制N个编码器基本单元        self.layers = clones(layer, N)        # 层归一化        self.norm = LayerNorm(layer.size)
    def forward(self, x, mask):        """        循环编码器基本单元N次        """        for layer in self.layers:            x = layer(x, mask)        return self.norm(x)

def subsequent_mask(size):    "Mask out subsequent positions."    # 设定subsequent_mask矩阵的shape    attn_shape = (1, size, size)    # 生成一个右上角(不含主对角线)为全1，左下角(含主对角线)为全0的subsequent_mask矩阵    subsequent_mask = np.triu(np.ones(attn_shape), k=1).astype('uint8')    # 返回一个右上角(不含主对角线)为全False，左下角(含主对角线)为全True的subsequent_mask矩阵    return torch.from_numpy(subsequent_mask) == 0

class DecoderLayer(nn.Module):    def __init__(self, size, self_attn, src_attn, feed_forward, dropout):        super(DecoderLayer, self).__init__()        self.size = size        # 自注意力机制        self.self_attn = self_attn        # 上下文注意力机制        self.src_attn = src_attn        self.feed_forward = feed_forward        self.sublayer = clones(SublayerConnection(size, dropout), 3)
    def forward(self, x, memory, src_mask, tgt_mask):        # memory为编码器输出隐表示        m = memory        # 自注意力机制，q、k、v均来自解码器隐表示        x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, tgt_mask))        # 上下文注意力机制：q为来自解码器隐表示，而k、v为编码器隐表示        x = self.sublayer[1](x, lambda x: self.self_attn(x, m, m, src_mask))        return self.sublayer[2](x, self.feed_forward)

class Decoder(nn.Module):    def __init__(self, layer, N):        super(Decoder, self).__init__()        self.layers = clones(layer, N)        self.norm = LayerNorm(layer.size)
    def forward(self, x, memory, src_mask, tgt_mask):        """        循环解码器基本单元N次        """        for layer in self.layers:            x = layer(x, memory, src_mask, tgt_mask)        return self.norm(x)

class Generator(nn.Module):    """    解码器输出经线性变换和softmax函数映射为下一时刻预测单词的概率分布    """    def __init__(self, d_model, vocab):        super(Generator, self).__init__()        # decode后的结果，先进入一个全连接层变为词典大小的向量        self.proj = nn.Linear(d_model, vocab)
    def forward(self, x):        # 然后再进行log_softmax操作(在softmax结果上再做多一次log运算)        return F.log_softmax(self.proj(x), dim=-1)

class Transformer(nn.Module):    def __init__(self, encoder, decoder, src_embed, tgt_embed, generator):        super(Transformer, self).__init__()        self.encoder = encoder        self.decoder = decoder        self.src_embed = src_embed        self.tgt_embed = tgt_embed        self.generator = generator
    def encode(self, src, src_mask):        return self.encoder(self.src_embed(src), src_mask)
    def decode(self, memory, src_mask, tgt, tgt_mask):        return self.decoder(self.tgt_embed(tgt), memory, src_mask, tgt_mask)
    def forward(self, src, tgt, src_mask, tgt_mask):        # encoder的结果作为decoder的memory参数传入，进行decode        return self.decode(self.encode(src, src_mask), src_mask, tgt, tgt_mask)

def make_model(src_vocab, tgt_vocab, N=6, d_model=512, d_ff=2048, h=8, dropout=0.1):    c = copy.deepcopy    # 实例化Attention对象    attn = MultiHeadedAttention(h, d_model).to(DEVICE)    # 实例化FeedForward对象    ff = PositionwiseFeedForward(d_model, d_ff, dropout).to(DEVICE)    # 实例化PositionalEncoding对象    position = PositionalEncoding(d_model, dropout).to(DEVICE)    # 实例化Transformer模型对象    model = Transformer(        Encoder(EncoderLayer(d_model, c(attn), c(ff), dropout).to(DEVICE), N).to(DEVICE),        Decoder(DecoderLayer(d_model, c(attn), c(attn), c(ff), dropout).to(DEVICE), N).to(DEVICE),        nn.Sequential(Embeddings(d_model, src_vocab).to(DEVICE), c(position)),        nn.Sequential(Embeddings(d_model, tgt_vocab).to(DEVICE), c(position)),        Generator(d_model, tgt_vocab)).to(DEVICE)
    # This was important from their code.    # Initialize parameters with Glorot / fan_avg.    for p in model.parameters():        if p.dim() > 1:            # 这里初始化采用的是nn.init.xavier_uniform            nn.init.xavier_uniform_(p)    return model.to(DEVICE)

class LabelSmoothing(nn.Module):    """    标签平滑    """
    def __init__(self, size, padding_idx, smoothing=0.0):        super(LabelSmoothing, self).__init__()        self.criterion = nn.KLDivLoss(reduction='sum')        self.padding_idx = padding_idx        self.confidence = 1.0 - smoothing        self.smoothing = smoothing        self.size = size        self.true_dist = None
    def forward(self, x, target):        assert x.size(1) == self.size        true_dist = x.data.clone()        true_dist.fill_(self.smoothing / (self.size - 2))        true_dist.scatter_(1, target.data.unsqueeze(1), self.confidence)        true_dist[:, self.padding_idx] = 0        mask = torch.nonzero(target.data == self.padding_idx)        if mask.dim() > 0:            true_dist.index_fill_(0, mask.squeeze(), 0.0)        self.true_dist = true_dist        return self.criterion(x, Variable(true_dist, requires_grad=False))