AI大模型之路 第三篇：從零實(shí)現(xiàn)詞嵌入模型，加深理解！

你好，我是郭震

今天我們研究「AI大模型第三篇」：詞維度預(yù)測(cè)，很多讀者聽過詞嵌入，這篇文章解答下面問題：

• 詞嵌入是什么意思？

• 怎么做到的？原理是什么？

• 從零實(shí)現(xiàn)一個(gè)專屬你數(shù)據(jù)集的詞嵌入

我們完整從零走一遍，根基的東西要理解透，這樣才能發(fā)明出更好的東西。

1 skip-gram模型

Skip-gram模型是一種廣泛使用的詞嵌入（Word Embedding）方法，由Mikolov等人在2013年提出。它是Word2Vec模型的一種形式，主要用于從大量文本中學(xué)習(xí)詞匯的高質(zhì)量向量表示。

Skip-gram模型的目標(biāo)是通過給定的目標(biāo)詞來預(yù)測(cè)其上下文中的詞匯，從而在這個(gè)過程中學(xué)習(xí)詞的嵌入表示。

因此，Skip-gram模型通過給定詞預(yù)測(cè)上下文，來最終學(xué)習(xí)到每個(gè)單詞的詞嵌入表示。

★

有些同學(xué)可能不理解，通過給定詞預(yù)測(cè)上下文，是什么意思？為什么要這么做？

因?yàn)?，某個(gè)單詞的上下文是有規(guī)律可尋的，比如 am單詞的上下文，一般就是 I，teacher，或tired，am后面一定不會(huì)出現(xiàn)：eat或walk，因?yàn)閮蓚€(gè)動(dòng)詞不可能出現(xiàn)在一起。

正是利用這個(gè)規(guī)律，也就是已知條件，我們學(xué)習(xí)到另一些好的特性，比如在這里，我們學(xué)習(xí)到每一個(gè)單詞的數(shù)學(xué)向量表示，計(jì)算機(jī)只認(rèn)得數(shù)字，它不認(rèn)識(shí)我們認(rèn)識(shí)的單詞。

2 求解問題

假設(shè)我們有一個(gè)簡(jiǎn)單的句子："the quick brown fox jumps over the lazy dog"，并且我們選擇Skip-gram模型進(jìn)行詞向量的訓(xùn)練。

我們可以挑選“fox”作為輸入詞，上下文窗口大小為2：

• 輸入："fox"
• 預(yù)測(cè)的上下文："quick"、"brown"、"jumps"、"over"

3 求解思路：

• 1 對(duì)“fox”進(jìn)行獨(dú)熱編碼。
• 2 使用Word2Vec模型預(yù)測(cè)“fox”的上下文詞。
• 3 通過調(diào)整模型權(quán)重來最小化預(yù)測(cè)誤差，使得模型可以更準(zhǔn)確地預(yù)測(cè)到“fox”的正確上下文。

4 訓(xùn)練模型

用到的包

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, Dataset

詞匯表和單詞索引

sentence = "the quick brown fox jumps over the lazy dog"
words = sentence.split()
word_to_ix = {word: i for i, word in enumerate(set(words))}

創(chuàng)建你的數(shù)據(jù)集

context_size = 2
data = []
for i in range(context_size, len(words) - context_size):
    target = words[i]
    context = [words[i - j - 1] for j in range(context_size)] + [words[i + j + 1] for j in range(context_size)]
    data.append((target, context))
    
class SkipGramDataset(Dataset):
    def __init__(self, data, word_to_ix):
        self.data = data
        self.word_to_ix = word_to_ix

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        target, context = self.data[idx]
        target_idx = self.word_to_ix[target]
        context_idx = torch.tensor([self.word_to_ix[w] for w in context], dtype=torch.long)
        return target_idx, context_idx

創(chuàng)建模型SkipGramModel

class SkipGramModel(nn.Module):
    def __init__(self, vocab_size, embedding_dim):
        super(SkipGramModel, self).__init__()
        self.embeddings = nn.Embedding(vocab_size, embedding_dim)
        self.predictions = nn.Linear(embedding_dim, vocab_size)

    def forward(self, input_words):
        embeds = self.embeddings(input_words)
        scores = self.predictions(embeds)
        log_probs = torch.log_softmax(scores, dim=1)
        return log_probs

訓(xùn)練模型SkipGramModel

# 初始化模型和優(yōu)化器
embedding_dim = 10
vocab_size = len(word_to_ix)
model = SkipGramModel(vocab_size, embedding_dim)
loss_function = nn.NLLLoss()
optimizer = optim.SGD(model.parameters(), lr=0.01)

# 數(shù)據(jù)加載器
dataset = SkipGramDataset(data, word_to_ix)
dataloader = DataLoader(dataset, batch_size=1, shuffle=True)

for epoch in range(50):
    total_loss = 0
    for target_idx, context_idx in dataloader:
        model.zero_grad()

        # 得到模型的預(yù)測(cè)對(duì)數(shù)概率輸出
        log_probs = model(target_idx)

        # 循環(huán)計(jì)算每個(gè)上下文詞的損失并累加
        loss = 0
        for context_word_idx in context_idx.view(-1):
            loss += loss_function(log_probs, context_word_idx.unsqueeze(0))

        # 反向傳播和優(yōu)化
        loss.backward()
        optimizer.step()

        total_loss += loss.item()
    print(f'Epoch {epoch + 1}, Loss: {total_loss}')

5 使用模型預(yù)測(cè)

使用模型SkipGramModel

def predict_context(model, input_word, word_to_ix, ix_to_word, top_n=3):
    # Check if the word is in the dictionary
    if input_word not in word_to_ix:
        return f"Word '{input_word}' not in vocabulary."

    # Prepare the model for evaluation
    model.eval()

    # Convert word to index and wrap in tensor
    word_idx = torch.tensor([word_to_ix[input_word]], dtype=torch.long)

    # Forward pass to get log probabilities
    with torch.no_grad():
        log_probs = model(word_idx)

    # Convert log probabilities to actual probabilities
    probs = torch.exp(log_probs).squeeze(0)  # Remove batch dimension

    # Get the indices of the top N probabilities
    top_indices = torch.topk(probs, top_n, dim=0)[1].tolist()

    # Convert indices back to words
    top_words = [ix_to_word[idx] for idx in top_indices]

    return top_words


# Create a reverse dictionary to map indices back to words
ix_to_word = {index: word for word, index in word_to_ix.items()}

# Example usage: predict context words for 'fox'
predicted_context_words = predict_context(model, 'fox', word_to_ix, ix_to_word, top_n=4)
print(f"Context words for 'fox': {predicted_context_words}")

6 結(jié)果分析

以上代碼完整可運(yùn)行，我們打印預(yù)測(cè)結(jié)果，看到預(yù)測(cè)fox的上下文是準(zhǔn)確的：

最后打印我們得到的fox單詞的嵌入詞向量：

# 確保'fox'在詞匯表中
if 'fox' in word_to_ix:
    # 獲取'fox'的索引
    fox_index = word_to_ix['fox']

    # 獲取嵌入層
    embeddings = model.embeddings

    # 提取'fox'的嵌入向量
    fox_vector = embeddings(torch.tensor([fox_index], dtype=torch.long))

    # 打印向量
    print("Embedding vector for 'fox':")
    print(fox_vector)
else:
    print("Word 'fox' not found in the vocabulary.")

我們這里嵌入詞向量長(zhǎng)度為10，見代碼，看到打印結(jié)果長(zhǎng)度也是10，這是正確的：