语言模型

import numpy as np
import tensorflow as tf
from sklearn.model_selection import train_test_split

# 读取数据，需要将数据导入当前目录
data_path = 'nnlm_comments_data.txt'
with open(data_path, 'r', encoding='utf-8') as f:
    lines = f.read().split('评论')

for i in range(len(lines)):
    lines[i] = lines[i].strip().split(' ')[:-1]

# 构建数据词典
word_index = {}
word_idx = 0
for i in range(len(lines)):
    for j in range(len(lines[i])):
        if lines[i][j] not in word_index.keys():
            word_index[lines[i][j]] = word_idx
            word_idx += 1
word_index['[END]'] = len(word_index.keys())

# 构建反向词典
index_word = {word_index[word]: word for word in word_index.keys()}

# 构建训练数据，每生成一个词时，考虑前四个词
lines_index = []
for i in range(len(lines)):
    lines[i].append('[END]')
    line_index = [word_index[word] for word in lines[i]]
    lines_index.append(line_index)

X_data = []
y_data = []
for i in range(len(lines_index)):
    if len(lines_index[i]) > 4:
        for j in range(0, len(lines_index[i]) - 4):
            X_data.append(lines_index[i][j:j + 4])
            y_data.append(lines_index[i][j + 4])

X_data = np.array(X_data)
y_data = np.array(y_data)
print(X_data.shape, y_data.shape)

X_train, X_test, y_train, y_test = train_test_split(X_data, y_data)
print(X_train.shape, y_train.shape, X_test.shape, y_test.shape)
X_train = X_train[:100000]
X_test = X_test[:10000]
y_train = y_train[:100000]
y_test = y_test[:10000]

# 将数据转化为tf.data.Dataset格式
train_ds = tf.data.Dataset.from_tensor_slices((X_train, y_train)).shuffle(10000).batch(64)
test_ds = tf.data.Dataset.from_tensor_slices((X_test, y_test)).batch(64)


# 定义NNLM模型
class NNLM(tf.keras.models.Model):
    def __init__(self, vocab_size, embedding_dim, hidden_size):
        super(NNLM, self).__init__()
        self.embed_layer = tf.keras.layers.Embedding(vocab_size, embedding_dim, mask_zero=True)
        self.reshape_layer = tf.keras.layers.Reshape((-1,))
        self.fc_layer = tf.keras.layers.Dense(hidden_size, activation='relu')
        self.out_layer = tf.keras.layers.Dense(vocab_size)

    def call(self, x):
        # x shape: (batch_size, seq_len:4)
        embed = self.embed_layer(x)  # (batch_size, seq_len:4, embedding_dim)
        reshape = self.reshape_layer(embed)  # (batch_size, seq_len * embedding_dim)
        fc = self.fc_layer(reshape)  # (batch_size, hidden_size)
        out = self.out_layer(fc)  # (batch_size, vocab_size)
        return out


# 网络参数
vocab_size = len(index_word.keys())
embedding_dim = 256
hidden_size = 256
batch_size = 128

# 定义NNLM对象
nnlm = NNLM(vocab_size, embedding_dim, hidden_size)

# 定义优化器与损失函数
optimizer = tf.keras.optimizers.Adam()
loss_function = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
# 记录每个epoch的平均loss
train_loss = tf.keras.metrics.Mean(name='train_loss')
test_loss = tf.keras.metrics.Mean(name='test_loss')


# 定义训练步
# @tf.function
def train_step(batch_X, batch_y):
    with tf.GradientTape() as tape:
        predictions = nnlm(batch_X)
        loss = loss_function(batch_y, predictions)
    gradients = tape.gradient(loss, nnlm.trainable_variables)
    optimizer.apply_gradients(zip(gradients, nnlm.trainable_variables))

    train_loss(loss)


# 定义测试步
# @tf.function
def test_step(batch_X, batch_y):
    predictions = nnlm(batch_X)
    loss = loss_function(batch_y, predictions)

    test_loss(loss)


EPOCHS = 5  # 训练轮次

for epoch in range(EPOCHS):
    # 每个epoch开始时将指标重置
    train_loss.reset_states()
    test_loss.reset_states()

    for batch_X, batch_y in train_ds:
        train_step(batch_X, batch_y)

    for batch_X, batch_y in test_ds:
        test_step(batch_X, batch_y)

    template = 'Epoch {}, Loss: {}, Test Loss: {}'
    print(template.format(epoch + 1, train_loss.result(), test_loss.result()))

# 使用NNLM生成句子：从前四个单词出发，生成第五个单词，然后后取第2-5个单词，生成第六个单词，以此类推