如何使用 LSTM model 预测下一个词？

Question

我目前正在 Pytorch 中构建一个 LSTM model 来预测给定输入的下一个单词。

我的model：

class LSTM(nn.Module):
  def __init__(self, vocab_size, embedding_dim, hidden_dim, num_layers, dropout_rate, tie_weights):
    super().__init__()

    self.num_layers = num_layers
    self.hidden_dim = hidden_dim
    self.embedding_dim = embedding_dim

    self.embedding = nn.Embedding(vocab_size, embedding_dim)
    self.lstm = nn.LSTM(embedding_dim, hidden_dim, num_layers=num_layers, 
                    dropout=dropout_rate, batch_first=True)
    self.dropout = nn.Dropout(dropout_rate)
    self.linear = nn.Linear(hidden_dim, vocab_size)

    if tie_weights:
      #Embedding and hidden layer need to be same size for weight tieing
          assert embedding_dim == hidden_dim, 'cannot tie, check dims'
          self.linear.weight = self.embedding.weight
    self.init_weights()

  def forward(self, x):
    # x is a batch of input sequences
    x = self.embedding(x)
    x, _ = self.lstm(x)
    x = self.linear(x)
    return x

  def init_weights(self):
    init_range_emb = 0.1
    init_range_other = 1/math.sqrt(self.hidden_dim)
    self.embedding.weight.data.uniform_(-init_range_emb, init_range_emb)
    self.linear.weight.data.uniform_(-init_range_other, init_range_other)
    self.linear.bias.data.zero_()
    for i in range(self.num_layers):
        self.lstm.all_weights[i][0] = torch.FloatTensor(self.embedding_dim,
                self.hidden_dim).uniform_(-init_range_other, init_range_other) 
        self.lstm.all_weights[i][1] = torch.FloatTensor(self.hidden_dim, 
                self.hidden_dim).uniform_(-init_range_other, init_range_other)



vocab_size = len(vocab)
embedding_dim = 100
hidden_dim = 100
num_layers = 2
dropout_rate = 0.4
tie_weights = True
model = LSTM(vocab_size, embedding_dim, hidden_dim, num_layers, dropout_rate, tie_weights)
model.to(device)

培训和评估功能：

import copy
import time

criterion = nn.CrossEntropyLoss()
lr = 20.0  # learning rate
optimizer = torch.optim.SGD(model.parameters(), lr=lr)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, 1.0, gamma=0.95)

def train(model: nn.Module) -> None:
    model.train()  # turn on train mode
    total_loss = 0.
    log_interval = 200
    start_time = time.time()
    
    num_batches = len(train_data) // bptt
    for batch, i in enumerate(range(0, train_data.size(0) - 1, bptt)):
        data, targets = get_batch(train_data, i)
        seq_len = data.size(0)
        output = model(data)
        loss = criterion(output.view(-1, vocab_size), targets)

        optimizer.zero_grad()
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 0.5)
        optimizer.step()

        total_loss += loss.item()
        
        if batch % log_interval == 0 and batch > 0:
            lr = scheduler.get_last_lr()[0]
            ms_per_batch = (time.time() - start_time) * 1000 / log_interval
            cur_loss = total_loss / log_interval
            ppl = math.exp(cur_loss)
            print(f'| epoch {epoch:3d} | {batch:5d}/{num_batches:5d} batches | '
                  f'lr {lr:02.2f} | ms/batch {ms_per_batch:5.2f} | '
                  f'loss {cur_loss:5.2f} | ppl {ppl:8.2f}')
            total_loss = 0
            start_time = time.time()


def evaluate(model: nn.Module, eval_data: Tensor) -> float:
    model.eval()  # turn on evaluation mode
    total_loss = 0.
    with torch.no_grad():
        for i in range(0, eval_data.size(0) - 1, bptt):
            data, targets = get_batch(eval_data, i)
            seq_len = data.size(0)
            output = model(data)
            output_flat = output.view(-1, vocab_size)
            total_loss += seq_len * criterion(output_flat, targets).item()
    return total_loss / (len(eval_data) - 1)

训练循环

best_val_loss = float('inf')
epochs = 50
best_model = None

for epoch in range(1, epochs + 1):
    epoch_start_time = time.time()
    train(model)
    val_loss = evaluate(model, val_data)
    val_ppl = math.exp(val_loss)
    elapsed = time.time() - epoch_start_time
    print('-' * 89)
    print(f'| end of epoch {epoch:3d} | time: {elapsed:5.2f}s | '
          f'valid loss {val_loss:5.2f} | valid ppl {val_ppl:8.2f}')
    print('-' * 89)

    if val_loss < best_val_loss:
        best_val_loss = val_loss
        best_model = copy.deepcopy(model)

    scheduler.step()

我的问题是我不知道如何与 go 联系。 我看过一些基于字符的 LSTM 文本生成器的实现，但我正在寻找它是基于单词的。 例如，我想传递“你好吗”之类的输入，output 将包含下一个预测词，例如“你今天好吗”

任何帮助表示赞赏。

Answer 1

我建议尝试附加链接中的示例（ https://www.kaggle.com/code/ysthehurricane/next-word-prediction-bi-lstm-tutorial-easy-way ）。

您可以从下面的附加链接下载数据集。 ( https://www.kaggle.com/datasets/dorianlazar/medium-articles-dataset )

它尝试使用双向 LSTM 架构来预测下一个单词。 我认为这个例子最适合你的需要，这会给你一个进一步进行的想法。

您可以按照第一个链接中提供的说明进行操作。