逐句解析如何使用nnlm模型实现单词预测

实现了一个简单的神经网络语言模型（NNLM），用于预测给定前n-1个单词的下一个单词。通过训练，模型能够学习到单词之间的上下文关系，从而进行预测。

代码摘自A Neural Probabilistic Language Model 论文阅读及实战，感谢原作者

# -*- coding: utf-8 -*-
import torch
import torch.nn as nn
import torch.optim as optim
from torch.autograd import Variable


dtype = torch.FloatTensor
sentences = ["i like dog", "i love coffee", "i hate milk"]

word_list = " ".join(sentences).split()  # 制作词汇表-用空格连在一起，然后按照空格分割成单词列表
print(word_list)
word_list = list(set(word_list))  # 去除词汇表中的重复元素
print("去重后的word_list:", word_list)
word_dict = {w: i for i, w in enumerate(word_list)}  # 将每个单词对应于相应的索引
number_dict = {i: w for i, w in enumerate(word_list)}  # 将每个索引对应于相应的单词
n_class = len(word_dict)  # 唯一单词的总数

# NNLM parameters
n_step = 2   # 根据前两个单词预测第3个单词
n_hidden = 2  # 隐藏层神经元的个数
m = 2  # 词向量的维度


# 由于pytorch中输入的数据是以batch小批量进行输入的，下面的函数就是将原始数据以一个batch为基本单位喂给模型
def make_batch(sentences):
    # 将句子转换为输入和目标数据批次
    input_batch = []  # 每个句子的前n-1个单词
    target_batch = []  # 每个句子的最后一个单词
    for sentence in sentences:
        word = sentence.split()
        input = [word_dict[w] for w in word[:-1]]
        target = word_dict[word[-1]]
        input_batch.append(input)
        target_batch.append(target)
    return input_batch, target_batch


# 模型
class NNLM(nn.Module):
    # 神经网络语言模型类
    def __init__(self):
        super(NNLM, self).__init__()
        self.C = nn.Embedding(n_class, embedding_dim=m)  # 嵌入层，讲单词索引转化为固定大小的向量
        self.H = nn.Parameter(torch.randn(n_step * m, n_hidden).type(dtype))  # 输入到隐藏层的权重
        self.W = nn.Parameter(torch.randn(n_step * m, n_class).type(dtype))  # 从输入层到输出层的权重
        self.d = nn.Parameter(torch.randn(n_hidden).type(dtype))  # 输入到隐藏层的偏置
        self.U = nn.Parameter(torch.randn(n_hidden, n_class).type(dtype))  # 从隐藏层到输出层的权重
        self.b = nn.Parameter(torch.randn(n_class).type(dtype))  # 输出层的偏置

    def forward(self, x):
        #  前向传播方法，定义数据如何通过模型
        x = self.C(x)  # 将输入的单词索引转化为嵌入向量
        x = x.view(-1, n_step * m)  # 将嵌入向量重新排列为合适的形状
        # x: [batch_size, n_step*n_class]
        tanh = torch.tanh(self.d + torch.mm(x, self.H))  # 使用tanh激活函数处理隐藏层
        # tanh: [batch_size, n_hidden]
        output = self.b + torch.mm(x, self.W) + torch.mm(tanh, self.U)  # 计算输出层的logits
        # output: [batch_size, n_class]
        return output


model = NNLM()  # 实例化NNLM模型

criterion = nn.CrossEntropyLoss()  # 定义损失函数，使用交叉熵损失
optimizer = optim.Adam(model.parameters(), lr=0.001)  # 定义优化器，使用adam优化器

# 制作输入
input_batch, target_batch = make_batch(sentences)  # 调用make_batch函数生成批次数据
input_batch = Variable(torch.LongTensor(input_batch))  # 将输入转化为pytorch变量
target_batch = Variable(torch.LongTensor(target_batch))  # 将目标数据转化为pytorch变量


# 开始训练
for epoch in range(5000):
    # 训练模型5000个epoch
    optimizer.zero_grad()  # 清空梯度
    output = model(input_batch)  # 向前传播
    loss = criterion(output, target_batch)  # 计算损失
    if (epoch + 1) % 1000 == 0:
        # 每1000个epoch打印一次损失
        print("Epoch:{}".format(epoch + 1), "Loss:{:.3f}".format(loss))
    loss.backward()  # 反向传播
    optimizer.step()  # 更新模型参数

# 预测
predict = model(input_batch).data.max(
    1, keepdim=True)[1]  # [batch_size, n_class]
# 使用训练好的模型进行预测 model(input_batch)：前向传播，返回每个单词的预测概率，data.max(1, keepdim=True)[1]找到概率最大的单词的索引
print("predict: \n", predict)
# 测试
print([sentence.split()[:2] for sentence in sentences], "---->",
      [number_dict[n.item()] for n in predict.squeeze()])
# predict.squeeze(): 去除维度为1的维度
# number_dict[n.item()]: 将预测的索引转换为对应的单词