Embedding之Word2Vec的代码实现

目录

一、Embedding与Word2Vec的联系

二、代码示例及注释

三、运行结果

---

一、Embedding与Word2Vec的联系

嵌入（Embedding）是一种通用的概念，用于将高维度的离散数据映射到低维度的连续向量空间。这种映射允许我们更好地表示数据的特征，并在深度学习中广泛应用于各种领域，不仅仅是NLP。嵌入的目标是通过学习数据的分布式表示来捕捉数据之间的关系。

在NLP中，嵌入通常是指将单词或字符映射为连续向量，这就是Word2Vec的一个例子。嵌入可以通过深度学习模型自动学习，也可以使用预训练的模型（如Word2Vec、GloVe、FastText）来获得。因此，Word2Vec是嵌入的一种具体形式，专门用于将单词嵌入为连续向量。可以把Word2Vec看作是嵌入的一个子集，特别适用于NLP领域。嵌入是一个更广泛的概念，适用于多种数据类型和领域。

详情可以查看我的另一篇博客; 词向量Embedding的原理

二、代码示例及注释

import torch  # PyTorch深度学习框架
import torch.nn as nn  # 用于定义神经网络结构的模块
import torch.nn.functional as F  # 包含神经网络操作的函数
import torch.optim as optim  # 包含各种优化器
from tqdm import tqdm, trange  # 用于在命令行中显示进度条
import numpy as np  # 用于科学计算和数据处理的核心库


# 设置词左边和右边选择的个数(即上下文词汇个数)
CONTEXT_SIZE = 2

# 原始文本数据
raw_text = """We are about to study the idea of a computational process.
Computational processes are abstract beings that inhabit computers.
As they evolve, processes manipulate other abstract things called data.
The evolution of a process is directed by a pattern of rules
called a program. People create programs to direct processes. In effect,
we conjure the spirits of the computer with our spells.""".split()

# 创建词汇表，去除重复的单词
vocab = set(raw_text)
vocab_size = len(vocab)

# 创建单词到索引和索引到单词的映射
word_to_idx = {word: i for i, word in enumerate(vocab)}
idx_to_word = {i: word for i, word in enumerate(vocab)}

# 创建训练数据，使用上下文窗口和目标单词的方式
data = []
for i in range(CONTEXT_SIZE, len(raw_text) - CONTEXT_SIZE):
    context = (
        [raw_text[i - (2-j)] for j in range(CONTEXT_SIZE)]  # 上下文窗口左侧的单词
        + [raw_text[i + j + 1] for j in range(CONTEXT_SIZE)]  # 上下文窗口右侧的单词
    )
    target = raw_text[i]  # 目标单词
    data.append((context, target))

# 定义函数，将上下文转换为索引张量
def make_context_vector(context, word_to_ix):
    idxs = [word_to_ix[w] for w in context]
    return torch.tensor(idxs, dtype=torch.long)

# 创建一个上下文向量示例
make_context_vector(data[0][0], word_to_idx)

# 定义一个CBOW模型类
class CBOW(nn.Module):
    def __init__(self, vocab_size, embedding_dim):
        super(CBOW, self).__init__()

        # 创建嵌入层
        self.embeddings = nn.Embedding(vocab_size, embedding_dim)

        # 创建线性投影层
        self.proj = nn.Linear(embedding_dim, 128)

        # 创建输出层
        self.output = nn.Linear(128, vocab_size)

    def forward(self, inputs):
        # 计算输入上下文单词的嵌入向量的和
        embeds = sum(self.embeddings(inputs)).view(1, -1)

        # 通过线性投影层并应用ReLU激活函数
        out = F.relu(self.proj(embeds))

        # 通过输出层并应用log_softmax得到概率分布
        out = self.output(out)
        nll_prob = F.log_softmax(out, dim=-1)
        return nll_prob

# 检查是否有GPU可用，如果没有则使用CPU
device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
# 创建一个具有词汇大小 'vocab_size' 和嵌入维度为 10 的CBOW模型，并将其移动到指定的 'device'
model = CBOW(vocab_size, 10).to(device)

# 使用Adam优化器，学习率为0.001，用于更新模型参数
optimizer = optim.Adam(model.parameters(), lr=0.001)

# 存储训练期间的损失的列表
losses = []

# 定义负对数似然损失函数，适用于多类分类任务
loss_function = nn.NLLLoss()

# 开始进行 200 个周期的训练
for epoch in trange(200):
    total_loss = 0

    # 遍历训练数据
    for context, target in tqdm(data):
        # 将训练数据移到GPU
        context_vector = make_context_vector(context, word_to_idx).to(device)
        target = torch.tensor([word_to_idx[target]]).to(device)

        # 将梯度清零，以防止梯度累积
        model.zero_grad()

        # 正向传播以获得预测
        train_predict = model(context_vector)
        
        # 计算预测和目标之间的负对数似然损失
        loss = loss_function(train_predict, target)

        # 将优化器的梯度清零
        optimizer.zero_grad()

        # 反向传播以计算梯度
        loss.backward()

        # 使用优化器更新模型参数
        optimizer.step()

        # 累积每个周期的总损失
        total_loss += loss.item()

    # 将每个周期的总损失添加到 'losses' 列表中
    losses.append(total_loss)

# 使用特定上下文测试训练过的模型
context = ['People', 'create', 'to', 'direct']
context_vector = make_context_vector(context, word_to_idx).to(device)

# 将模型设置为评估模式
model.eval()

# 对给定上下文进行预测
predict = model(context_vector)

# 找到预测中最大值的索引
max_idx = predict.argmax(1)

# 获取嵌入层的权重矩阵
W = model.embeddings.weight.cpu().detach().numpy()

# 创建一个将单词映射到其对应单词向量的字典
word_2_vec = {}
for word in word_to_idx.keys():
    word_2_vec[word] = W[word_to_idx[word], :]

# 打印生成的单词向量字典
print(word_2_vec)

三、运行结果