Word2vec之skip-gram训练词向量

参考自哈工大车万翔等老师编写的《自然语言处理-基于预训练模型的方法》

# coding: utf-8
# Name:     tesst2
# Author:   dell
# Data:     2021/10/12

# 基于负采样的skip-garm模型
import torch
import torch.nn.functional as F
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
from collections import defaultdict
from tqdm import tqdm
import torch.optim as optim


BOS_TOKEN = "<bos>"
EOS_TOKEN = "<eos>"
PAD_TOKEN = "<pad>"


def load_reuters():
    # 从NLTK中导入Reuters数据处理模块
    from nltk.corpus import reuters
    # 获取Reuters数据中的所有句子（已完成标记解析）
    text = reuters.sents()
    # （可选）将预料中的词转换为小写
    text = [[word.lower() for word in sentence] for sentence in text]
    # 构建词表，并传入预留标记
    vocab = Vocab.build(text, reserved_tokens=[PAD_TOKEN, BOS_TOKEN, EOS_TOKEN])
    # 利用词表将文本数据转换为id表示
    corpus = [vocab.convert_tokens_to_ids(sentence) for sentence in text]
    return corpus, vocab

class Vocab:
    def __init__(self, tokens=None):
        self.idx_to_token = list()
        self.token_to_idx = dict()

        if tokens is not None:
            if "<unk>" not in tokens:
                tokens = tokens + ["<unk>"]
            for token in tokens:
                self.idx_to_token.append(token)
                self.token_to_idx[token] = len(self.idx_to_token) - 1

            self.unk = self.token_to_idx["<unk>"]

    @classmethod
    def build(cls, text, min_freq=1, reserved_tokens=None):
        token_freqs = defaultdict(int)
        for sentence in text:
            for token in sentence:
                token_freqs[token] += 1
        uniq_tokens = ["<unk>"] + (reserved_tokens if reserved_tokens else [])
        uniq_tokens += [token for token, freq in token_freqs.items() if freq >= min_freq and token != "<unk>"]

        return cls(uniq_tokens)

    def __len__(self):
        # 返回词表的大小，即词表中有多少个互不相同的标记
        return len(self.idx_to_token)

    def __getitem__(self, token):
        # 查找输入标记对应的索引值，如果该标记不存在，返回<unk>对应的索引值
        return self.token_to_idx.get(token, self.unk)

    def convert_tokens_to_ids(self, tokens):
        # 查找一系列输入标记对应的索引值，此处直接使用self即可，会调用__getitem__得到id
        return [self[token] for token in tokens]
        # 或者是
        # return [self.token_to_idx[token] for token in tokens]

    def convert_ids_to_tokens(self, indices):
        # 查找一系列索引值对应的标记
        return [self.idx_to_token[index] for index in indices]


class SGNSDataset(Dataset):
    def __init__(self, corpus, vocab, context_size=2, n_negatives=5, ns_dist=None):
        self.data = []
        self.bos = vocab[BOS_TOKEN]
        self.eos = vocab[EOS_TOKEN]
        self.pad = vocab[PAD_TOKEN]
        for sentence in tqdm(corpus, desc="Dataset Construction"):
            sentence = [self.bos] + sentence + [self.eos]
            for i in range(1, len(sentence)-1):
                # 模型输入: (w, context)
                w = sentence[i]
                left_context_index = max(0, i-context_size)
                right_context_index = min(len(sentence), i+context_size)
                context = sentence[left_context_index:i] + sentence[i+1:right_context_index+1]
                context += [self.pad] * (2*context_size - len(context_size))
                self.data.append((w, context))
        # 负采样数量
        self.n_negatives = n_negatives
        # 负采样分布：若参数ns_dist为None，则使用均匀分布（从词表中均匀采样）
        self.ns_dist = ns_dist if ns_dist else torch.ones(len(vocab))

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

    def __getitem__(self, i):
        return self.data[i]

    def collate_fn(self, examples):
        words = torch.tensor([ex[0] for ex in examples], dtype=torch.long)
        contexts = torch.tensor([ex[1] for ex in examples], dtype=torch.long)
        batch_size, context_size = contexts.shape
        neg_contexts = []
        # 对批次内的样本分别进行负采样
        for i in range(batch_size):
            # 保证负样本不包含当前样本中的context
            ns_dist = self.ns_dist.index_fill(0, contexts[i], .0)   # dim=0, 需要填充的tensor的索引-contexts[i]， vale=.0
            # torch.multinomial--对ns_dist中的值，有放回(replacement=True)地抽取self.n_negatives * context_size
            neg_contexts.append(torch.multinomial(ns_dist, self.n_negatives * context_size, replacement=True))

        neg_contexts = torch.stack(neg_contexts, dim=0)

        return words, contexts, neg_contexts


class SGNSModel(nn.Module):
    def __init__(self, vocab_size, embedding_dim):
        super(SGNSModel, self).__init__()
        # 词向量
        self.w_embeddings = nn.Embedding(vocab_size, embedding_dim)
        # 上下文向量
        self.c_embeddings = nn.Embedding(vocab_size, embedding_dim)

    def forward_w(self, words):
        w_embeds = self.w_embeddings(words)
        return w_embeds

    def forward_c(self, contexts):
        c_embeds = self.c_embeddings(contexts)
        return c_embeds

def get_unigram_distribution(corpus, vocab_size):
    # 从给定的语料中计算Unigram概率分布
    token_counts = torch.tensor([0]*vocab_size)
    total_count = 0
    for sentence in corpus:
        total_count += len(sentence)
        for token in sentence:
            token_counts[token] += 1
    unigram_dist = torch.div(token_counts.float(), total_count)
    return unigram_dist


def save_pretrained(vocab, embeds, save_path):
    with open(save_path, "w") as writer:
        # 记录词向量大小
        writer.write(f"{embeds.shape[0]} {embeds.shape[1]}\n")
        for idx, token in enumerate(vocab.idx_to_token):
            vec = " ".join([f"{x}" for x in embeds[idx]])
            # 每一行对应一个单词以及由空格分隔的词向量
            writer.write(f"{token} {vec}\n")

def main():
    # 设置超参数
    embedding_dim = 128
    context_size = 3
    batch_size = 1024
    n_negatives = 5         # 负样本数量
    num_epoch = 10

    # 读取文本数据
    corpus, vocab = load_reuters()
    # 计算Unigram概率分布
    unigram_dist = get_unigram_distribution(corpus, len(vocab))
    # 根据Unigram概率分布计算负采样分布: p(w)**0.75
    # 为了防止低频单词被忽略-->通过取 0.75 次方，低频单词的概率将稍微变大。
    negative_sampling_dist = unigram_dist ** 0.75
    negative_sampling_dist /= negative_sampling_dist.sum()

    # 构建SGNS训练数据集
    dataset = SGNSDataset(corpus, vocab, context_size=context_size, n_negatives=n_negatives, ns_dist=negative_sampling_dist)

    # data_loader = get_loader(dataset, batch_size)
    data_loader = DataLoader(dataset, batch_size)

    model = SGNSModel(len(vocab), embedding_dim)
    device = "cuda" if torch.cuda.is_available() else "cpu"
    model.to(device)
    optimizer = optim.Adam(model.parameters(), lr=0.001)

    model.train()
    for epoch in range(num_epoch):
        total_loss = 0
        for batch in tqdm(data_loader, desc=f"Training Epoch {epoch}"):
            words, contexts, neg_contexts = [x.to(device) for x in batch]
            optimizer.zero_grad()
            batch_size = words.shape[0]
            # 分贝提取batch内词、上下文和负样本的向量表示
            word_embeds = model.forward_w(words).unsqueeze(dim=2)       # [batch_size, word_embedding, 1]
            context_embeds = model.forward_c(contexts)                  # [batch_size, context_num, context_word_embedding]
            neg_context_embeds = model.forward_c(neg_contexts)          # [batch_size, neg_context_word_embedding]

            # 正样本的分类（对数）似然
            context_loss = F.logsigmoid(torch.bmm(context_embeds, word_embeds).seqeeze(dim=2))   # [batch_size, context_num]-->预测上下文的词
            context_loss = context_loss.mean(dim=1)

            # 负样本的分类（对数）似然
            # torch.neg()--->按元素取负-->output = -1 * input
            neg_context_loss = F.logsigmoid(torch.bmm(neg_context_embeds, word_embeds).squeeze(dim=2).neg())
            neg_context_loss = neg_context_loss.view(batch_size, -1, n_negatives).sum(2)        # [batch_size, context_size, n_negatives]
            neg_context_loss = neg_context_loss.mean(dim=1)
            
            # 总体损失
            loss = -(context_loss + neg_context_loss).mean()
            loss.backward()
            optimizer.step()
            total_loss += loss.item()
            print(f"Loss: {total_loss:.2f}")
        # 合并词向量矩阵与上下文向量矩阵，作为最终的预训练词向量
        combined_embeds = model.w_embeddings.weight + model.c_embeddings.weight
        # 将词向量保存至sgns,vec文件
        save_pretrained(vocab, combined_embeds.data, "sgns.vec")

if __name__ == "__main__":
    main()