import collections
import math
import random
import sys
import time
import os
import numpy as np
import torch
from torch import nn
import torch.utils.data asData
1|1PTB 数据集
简单来说,Word2Vec 能从语料中学到如何将离散的词映射为连续空间中的向量,并保留其语义上的相似关系。那么为了训练 Word2Vec 模型,我们就需要一个自然语言语料库,模型将从中学习各个单词间的关系,这里我们使用经典的 PTB 语料库进行训练。PTB (Penn Tree Bank) 是一个常用的小型语料库,它采样自《华尔街日报》的文章,包括训练集、验证集和测试集。我们将在PTB训练集上训练词嵌入模型。
载入数据集
数据集训练文件 ptb.train.txt 示例:
aer banknote berlitz calloway centrust cluett fromstein gitano guterman ...
pierre N years old will join the board as a nonexecutive director nov. N
mr.is chairman of n.v. the dutch publishing group...
with open('/home/kesci/input/ptb_train1020/ptb.train.txt','r')as f:
lines = f.readlines()# 该数据集中句子以换行符为分割
raw_dataset =[st.split()for st in lines]# st是sentence的缩写,单词以空格为分割print('# sentences: %d'% len(raw_dataset))# 对于数据集的前3个句子,打印每个句子的词数和前5个词# 句尾符为 '' ,生僻词全用 '' 表示,数字则被替换成了 'N'for st in raw_dataset[:3]:print('# tokens:', len(st), st[:5])
counter = collections.Counter([tk for st in raw_dataset for tk in st])# tk是token的缩写
counter = dict(filter(lambda x: x[1]>=5, counter.items()))# 只保留在数据集中至少出现5次的词
idx_to_token =[tk for tk, _ in counter.items()]
token_to_idx ={tk: idx for idx, tk in enumerate(idx_to_token)}
dataset =[[token_to_idx[tk]for tk in st if tk in token_to_idx]for st in raw_dataset]# raw_dataset中的单词在这一步被转换为对应的idx
num_tokens = sum([len(st)for st in dataset])'# tokens: %d'% num_tokens
'# tokens: 887100'
二次采样
文本数据中一般会出现一些高频词,如英文中的“the”“a”和“in”。通常来说,在一个背景窗口中,一个词(如“chip”)和较低频词(如“microprocessor”)同时出现比和较高频词(如“the”)同时出现对训练词嵌入模型更有益。因此,训练词嵌入模型时可以对词进行二次采样。 具体来说,数据集中每个被索引词 wi 将有一定概率被丢弃,该丢弃概率为
P(wi)=max(1−√tf(wi),0)
其中 f(wi) 是数据集中词 wi 的个数与总词数之比,常数 t 是一个超参数(实验中设为 10−4)。可见,只有当 f(wi)>t 时,我们才有可能在二次采样中丢弃词 wi,并且越高频的词被丢弃的概率越大。具体的代码如下:
def discard(idx):'''
@params:
idx: 单词的下标
@return: True/False 表示是否丢弃该单词
'''return random.uniform(0,1)<1- math.sqrt(1e-4/ counter[idx_to_token[idx]]* num_tokens)
subsampled_dataset =[[tk for tk in st ifnot discard(tk)]for st in dataset]print('# tokens: %d'% sum([len(st)for st in subsampled_dataset]))def compare_counts(token):return'# %s: before=%d, after=%d'%(token, sum([st.count(token_to_idx[token])for st in dataset]), sum([st.count(token_to_idx[token])for st in subsampled_dataset]))print(compare_counts('the'))print(compare_counts('join'))
def get_centers_and_contexts(dataset, max_window_size):'''
@params:
dataset: 数据集为句子的集合,每个句子则为单词的集合,此时单词已经被转换为相应数字下标
max_window_size: 背景词的词窗大小的最大值
@return:
centers: 中心词的集合
contexts: 背景词窗的集合,与中心词对应,每个背景词窗则为背景词的集合
'''
centers, contexts =[],[]for st in dataset:if len(st)<2:# 每个句子至少要有2个词才可能组成一对“中心词-背景词”continue
centers += st
for center_i in range(len(st)):
window_size = random.randint(1, max_window_size)# 随机选取背景词窗大小
indices = list(range(max(0, center_i - window_size),
min(len(st), center_i +1+ window_size)))
indices.remove(center_i)# 将中心词排除在背景词之外
contexts.append([st[idx]for idx in indices])return centers, contexts
all_centers, all_contexts = get_centers_and_contexts(subsampled_dataset,5)
tiny_dataset =[list(range(7)), list(range(7,10))]print('dataset', tiny_dataset)for center, context in zip(*get_centers_and_contexts(tiny_dataset,2)):print('center', center,'has contexts', context)
dataset [[0,1,2,3,4,5,6],[7,8,9]]
center 0 has contexts [1,2]
center 1 has contexts [0,2,3]
center 2 has contexts [0,1,3,4]
center 3 has contexts [2,4]
center 4 has contexts [3,5]
center 5 has contexts [4,6]
center 6 has contexts [5]
center 7 has contexts [8]
center 8 has contexts [7,9]
center 9 has contexts [7,8]
注:数据批量读取的实现需要依赖负采样近似的实现,故放于负采样近似部分进行讲解。
1|2Skip-Gram 跳字模型
在跳字模型中,每个词被表示成两个 d 维向量,用来计算条件概率。假设这个词在词典中索引为 i ,当它为中心词时向量表示为 vi∈Rd,而为背景词时向量表示为 ui∈Rd 。设中心词 wc 在词典中索引为 c,背景词 wo 在词典中索引为 o,我们假设给定中心词生成背景词的条件概率满足下式:
P(wo∣wc)=exp(u⊤ovc)∑i∈Vexp(u⊤ivc)
PyTorch 预置的 Embedding 层
embed = nn.Embedding(num_embeddings=10, embedding_dim=4)print(embed.weight)
x = torch.tensor([[1,2,3],[4,5,6]], dtype=torch.long)print(embed(x))
embed_size =100
net = nn.Sequential(nn.Embedding(num_embeddings=len(idx_to_token), embedding_dim=embed_size),
nn.Embedding(num_embeddings=len(idx_to_token), embedding_dim=embed_size))
训练模型
def train(net, lr, num_epochs):
device = torch.device('cuda'if torch.cuda.is_available()else'cpu')print("train on", device)
net = net.to(device)
optimizer = torch.optim.Adam(net.parameters(), lr=lr)for epoch in range(num_epochs):
start, l_sum, n = time.time(),0.0,0for batch in data_iter:
center, context_negative, mask, label =[d.to(device)for d in batch]
pred = skip_gram(center, context_negative, net[0], net[1])
l = loss(pred.view(label.shape), label, mask).mean()# 一个batch的平均loss
optimizer.zero_grad()
l.backward()
optimizer.step()
l_sum += l.cpu().item()
n +=1print('epoch %d, loss %.2f, time %.2fs'%(epoch +1, l_sum / n, time.time()- start))
train(net,0.01,5)
train on cpu
epoch 1, loss 0.61, time 221.30s
epoch 2, loss 0.42, time 227.70s
epoch 3, loss 0.38, time 240.50s
epoch 4, loss 0.36, time 253.79s
epoch 5, loss 0.34, time 238.51s
注:由于本地CPU上训练时间过长,故只截取了运行的结果,后同。大家可以自行在网站上训练。
测试模型
def get_similar_tokens(query_token, k, embed):'''
@params:
query_token: 给定的词语
k: 近义词的个数
embed: 预训练词向量
'''
W = embed.weight.data
x = W[token_to_idx[query_token]]# 添加的1e-9是为了数值稳定性
cos = torch.matmul(W, x)/(torch.sum(W * W, dim=1)* torch.sum(x * x)+1e-9).sqrt()
_, topk = torch.topk(cos, k=k+1)
topk = topk.cpu().numpy()for i in topk[1:]:# 除去输入词print('cosine sim=%.3f: %s'%(cos[i],(idx_to_token[i])))
get_similar_tokens('chip',3, net[0])
【推荐】国内首个AI IDE,深度理解中文开发场景,立即下载体验Trae
【推荐】编程新体验,更懂你的AI,立即体验豆包MarsCode编程助手
【推荐】抖音旗下AI助手豆包,你的智能百科全书,全免费不限次数
【推荐】轻量又高性能的 SSH 工具 IShell:AI 加持,快人一步
· AI与.NET技术实操系列:向量存储与相似性搜索在 .NET 中的实现
· 基于Microsoft.Extensions.AI核心库实现RAG应用
· Linux系列:如何用heaptrack跟踪.NET程序的非托管内存泄露
· 开发者必知的日志记录最佳实践
· SQL Server 2025 AI相关能力初探
· 震惊!C++程序真的从main开始吗?99%的程序员都答错了
· winform 绘制太阳,地球,月球 运作规律
· 【硬核科普】Trae如何「偷看」你的代码?零基础破解AI编程运行原理
· 上周热点回顾(3.3-3.9)
· 超详细:普通电脑也行Windows部署deepseek R1训练数据并当服务器共享给他人