Word2Vec的Pytorch实现（二）

• 所需导入的包

import collections
import math
import random
import sys
import time
import os
import torch.utils.data as Data
import torch
import torch.nn as nn

• 1、读取并处理数据集

assert 'ptb.train.txt' in os.listdir("../data/ptb")

with open('../data/ptb/ptb.train.txt','r') as f:
    lines = f.readlines()
    raw_dataset = [sentence.split() for sentence in lines]

• 1、1建立词语索引

counter = collections.Counter([token for sentence in raw_dataset for token in sentence])
# print(counter.items())
counter = dict(filter(lambda x:x[1] >= 5,counter.items()))
# [token for sentence in raw_dataset for token in sentence] 代码作用同下
# a = []
# for sentence in raw_dataset:
#     for token in sentence:
#         a.append(token)
# 将词映射到整数索引
idx_to_token = [token for token,_ in counter.items()]
# print(idx_to_token) # ['pierre', '<unk>', 'N', 'years', 'old', 'will', 'join', 'the', 'board', 'as', 'a', 'nonexecutive', 'director', 'nov.',
token_to_idx = {token:idx for idx,token in enumerate(idx_to_token)}
# print(token_to_idx) # {'pierre': 0, '<unk>': 1, 'N': 2, 'years': 3, 'old': 4, 'will': 5, 'join': 6, 'the': 7, 'board': 8, 'as': 9,
dataset = [[token_to_idx[token] for token in sentence if token in token_to_idx] for sentence in raw_dataset]
num_tokens = sum([len(sentence) for sentence in dataset])
# print(num_tokens) # 二次采样前 887100

# [[token_to_idx[token] for token in sentence if token in token_to_idx] for sentence in raw_dataset]
# 代码作用同下
# print(len(dataset))
# print(dataset)
# b = []
# for sentence in raw_dataset:
#     c = []
#     for token in sentence:
#         if token in token_to_idx:
#             c.append(token_to_idx[token])
#     b.append(c)
# print(len(b))
# print(b)

• 1、2二次采样

# 1、2 二次采样
def discard(idx):
    return random.uniform(0,1) < 1 - math.sqrt(1e-4 / counter[idx_to_token[idx]] * num_tokens)

subsampled_dataset = [[token for token in sentence if not discard(token)] for sentence in dataset]
num_tokens_2 = sum([len(sentence) for sentence in subsampled_dataset])
print('二次采样后：',num_tokens_2) # 二次采样后 375930

# 比较一个词在二次采样前后出现在数据集中的次数
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')) # the: before=50770, after=2089
# print(compare_counts('join')) # join: before=45, after=45

• 1、3 提取中心词和背景词

# 1、3 提取中心词和背景词
def get_centers_and_contexts(dataset, max_window_size):
    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

# 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]
# # center 1 has contexts [0, 2]
# # center 2 has contexts [0, 1, 3, 4]
# # center 3 has contexts [1, 2, 4, 5]
# # center 4 has contexts [3, 5]
# # center 5 has contexts [3, 4, 6]
# # center 6 has contexts [4, 5]

# 设最大背景窗口大小为5，提取数据集中所有的中心词及其背景词
all_centers, all_contexts = get_centers_and_contexts(subsampled_dataset, 5)

• 2、负采样

def get_negatives(all_contexts, sampling_weights, K):
    all_negatives, neg_candidates, i = [], [], 0
    population = list(range(len(sampling_weights)))
    for contexts in all_contexts:
        negatives = []
        while len(negatives) < len(contexts) * K:
            if i == len(neg_candidates):
                # 根据每个词的权重（sampling_weights）随机生成k个词的索引作为噪声词。
                # 为了高效计算，可以将k设得稍大一点
                i, neg_candidates = 0, random.choices(
                    population, sampling_weights, k=int(1e5))
            neg, i = neg_candidates[i], i + 1
            # 噪声词不能是背景词
            if neg not in set(contexts):
                negatives.append(neg)
        all_negatives.append(negatives)
    return all_negatives

sampling_weights = [counter[w]**0.75 for w in idx_to_token]
all_negatives = get_negatives(all_contexts, sampling_weights, 5)

• 3、读取数据

class MyDataset(torch.utils.data.Dataset):
    def __init__(self, centers, contexts, negatives):
        assert len(centers) == len(contexts) == len(negatives)
        self.centers = centers
        self.contexts = contexts
        self.negatives = negatives

    def __getitem__(self, index):
        return (self.centers[index], self.contexts[index], self.negatives[index])

    def __len__(self):
        return len(self.centers)
        
# 小批量读取函数
def batchify(data):
    '''
      用作DataLoader的参数collate_fn: 输入是一个长为batchsize的list,list中的每个元素都是Dataset类调用__getitem__得到的结果
    '''
    max_len = max(len(c) + len(n) for _, c, n in data)
    centers, contexts_negatives, masks, labels = [], [], [], []
    for center, context, negative in data:
        cur_len = len(context) + len(negative)
        centers += [center]
        contexts_negatives += [context + negative + [0] * (max_len - cur_len)]
        masks += [[1] * cur_len + [0] * (max_len - cur_len)]
        labels += [[1] * len(context) + [0] * (max_len - len(context))]
    return (torch.tensor(centers).view(-1, 1), torch.tensor(contexts_negatives),
            torch.tensor(masks), torch.tensor(labels))

# 上面定义的batchify函数指定DataLoader实例中小批量的读取方式，然后打印读取的第一个批量中各个变量的形状
batch_size = 512
num_workers = 0 if sys.platform.startswith('win32') else 4

dataset = MyDataset(all_centers,
                    all_contexts,
                    all_negatives)
data_iter = Data.DataLoader(dataset, batch_size, shuffle=True,
                            collate_fn=batchify,
                            num_workers=num_workers)
for batch in data_iter:
    for name, data in zip(['centers', 'contexts_negatives', 'masks',
                           'labels'], batch):
        print(name, 'shape:', data.shape)
    break
# # centers shape: torch.Size([512, 1])
# # contexts_negatives shape: torch.Size([512, 60])
# # masks shape: torch.Size([512, 60])
# # labels shape: torch.Size([512, 60])

• 4、跳字模型

# 行数为词典大小（num_embeddings），列数为每个词向量的维度（embedding_dim）
embed = nn.Embedding(num_embeddings=20, embedding_dim=4) # [num_embeddings, embedding_dim]
# print(embed.weight)
# x = torch.tensor([[1,2,3],[4,5,6]], dtype=torch.long) # [2,3]
# print(embed(x)) # [2, 3, 4]

# 跳字模型前向计算
def skip_gram(center, contexts_and_negatives, embed_v, embed_u):
    v = embed_v(center)
    u = embed_u(contexts_and_negatives)
    pred = torch.bmm(v, u.permute(0, 2, 1)) #  u.permute(0, 2, 1) 将tensor的维度换位
    return pred

• 5、训练模型
• 5、1定义模型的损失函数（二元交叉熵损失函数）

class SigmoidBinaryCrossEntropyLoss(nn.Module):
    def __init__(self):
        super(SigmoidBinaryCrossEntropyLoss, self).__init__()
    def forward(self, inputs, targets, mask=None):
        '''
        :param inputs: Tensor shape: (batch_size, len)
        :param targets: Tensor of the same shape as input
        '''
        inputs, targets, mask = inputs.float(), targets.float(), mask.float()
        res = nn.functional.binary_cross_entropy_with_logits(
            input=inputs, target=targets, reduction='none', weight=mask
        )
        return res.mean(dim=1)

loss = SigmoidBinaryCrossEntropyLoss()

# pred = torch.tensor([[1.5,0.3,-1,2],[1.1,-0.6,2.2,0.4]])
# # 标签变量label中的1和0分别代表背景词和噪声词
# label = torch.tensor([[1,0,0,0],[1,1,0,0]])
# mask = torch.tensor([[1,1,1,1],[1,1,1,0]]) # 掩码变量
# res = loss(pred,label,mask) * mask.shape[1] / mask.float().sum(dim=1)
# print(res) # tensor([0.8740, 1.2100])

• 5、2初始化模型参数

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)
)

• 5、3定义训练函数

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, 0
        for 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])

            # 使用掩码变量mask来避免填充项对损失函数计算的影响
            l = (loss(pred.view(label.shape), label, mask) *
                 mask.shape[1] / mask.float().sum(dim=1)).mean()  # 一个batch的平均loss
            optimizer.zero_grad()
            l.backward()
            optimizer.step()
            l_sum += l.cpu().item()
            n += 1
            print('epoch %d, loss %.2f, time %.2fs'
                  % (epoch + 1, l_sum / n, time.time() - start))

• 训练

train(net, 0.01, 10)

• 6、应用词嵌入模型

def get_similar_tokens(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])