用于预训练词嵌入的数据集

用于预训练词嵌入的数据集

在了解word2vec模型的技术细节和大致的训练方法，让我们来看看它们的实现，具体地说，我们将以跳元模型和负采样为例

在本节中，我们将从用于预训练词嵌入模型的数据集开始：数据的原始格式将被转换为可以在训练期间迭代的小批量

import math
import os
import random
import torch
from d2l import torch as d2l

1 - 读取数据集

这里使用的数据集是Penn Tree Bank（PTB）。该语料库取自“华尔街日报”的文章，分为训练集、验证集和测试集。在原始格式中，文本文件的每一行表示由空格分隔的一句话，在这里，我们将每个单词视为一个词元

#@save
d2l.DATA_HUB['ptb'] = (d2l.DATA_URL + 'ptb.zip','319d85e578af0cdc590547f26231e4e31cdf1e42')

#@save
def read_ptb():
    """将PTB数据集加载到⽂本⾏的列表中"""
    data_dir = d2l.download_extract('ptb')
    # Readthetrainingset.
    with open(os.path.join(data_dir, 'ptb.train.txt')) as f:
        raw_text = f.read()
    return [line.split() for line in raw_text.split('\n')]

sentences = read_ptb()
f'# sentences数: {len(sentences)}'

'# sentences数: 42069'

在读取训练集之后，我们为语料库构建了一个词表，其中出现次数少于10次的任何单词都将由“”词元替换，请注意，原始数据集还包含表示稀有（未知）单词的“”词元

vocab = d2l.Vocab(sentences, min_freq=10)
f'vocab size: {len(vocab)}'

'vocab size: 6719'

2 - 下采样

#@save
def subsample(sentences, vocab):
    """下采样⾼频词"""
    # 排除未知词元''
    sentences = [[token for token in line if vocab[token] != vocab.unk] for line in sentences]
    counter = d2l.count_corpus(sentences)
    num_tokens = sum(counter.values())
    # 如果在下采样期间保留词元，则返回True
    def keep(token):
        return(random.uniform(0, 1) < math.sqrt(1e-4 / counter[token] * num_tokens))

    return ([[token for token in line if keep(token)] for line in sentences],counter)

subsampled, counter = subsample(sentences, vocab)

下面的代码片段绘制了下采样前后每句话的词元数量的直方图，正如预期那样，下采样通过删除高频词来显著缩短句子，这将使训练加速

d2l.show_list_len_pair_hist(['origin', 'subsampled'], '# tokens per sentence','count', sentences, subsampled);

​
[外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-bxKYpQKT-1662542915766)(https://yingziimage.oss-cn-beijing.aliyuncs.com/img/202209071728691.svg)]
​

对于单个词元，高频词“the”的采样率不到1/20

def compare_counts(token):
    return (f'"{token}"的数量：'
        f'之前={sum([l.count(token) for l in sentences])}, '
        f'之后={sum([l.count(token) for l in subsampled])}')

compare_counts('the')

'"the"的数量：之前=50770, 之后=2168'

相比之下，低频词“join”则被完全保留

compare_counts('join')

'"join"的数量：之前=45, 之后=45'

下采样之后，我们将词元映射到它们的语料库中的索引

corpus = [vocab[line] for line in subsampled]
corpus[:3]

[[], [392, 2115, 5, 406], [140, 5277, 3054, 1580]]

3 - 中心词和上下文词的提取

下面的get_centers_and_contexts函数从corpus中提取所有中心词及其上下文词。它随机采样1到max_window_size之间的整数作为上下文窗口。对于任一中心词，与其距离不超过采样上下文窗口大小词为其上下文词

#@save
def get_centers_and_contexts(corpus, max_window_size):
    """返回跳元模型中的中⼼词和上下⽂词"""
    centers, contexts = [], []
    for line in corpus:
        # 要形成“中⼼词-上下⽂词”对，每个句⼦⾄少需要有2个词
        if len(line) < 2:
            continue
        centers += line
        for i in range(len(line)): # 上下⽂窗⼝中间i
            window_size = random.randint(1, max_window_size)
            indices = list(range(max(0, i - window_size),
                                 min(len(line), i + 1 + window_size)))
            # 从上下⽂词中排除中⼼词
            indices.remove(i)
            contexts.append([line[idx] for idx in indices])
    return centers, contexts

接下来，我们创建一个人工数据集，分别包含7个和3个单词的两个句子。设置最大上下文窗口大小为2，并打印所有中心词及其上下文词

tiny_dataset = [list(range(7)), list(range(7, 10))]
print('数据集', tiny_dataset)
for center, context in zip(*get_centers_and_contexts(tiny_dataset, 2)):
    print('中⼼词', center, '的上下⽂词是', context)

数据集 [[0, 1, 2, 3, 4, 5, 6], [7, 8, 9]]
中⼼词 0 的上下⽂词是 [1, 2]
中⼼词 1 的上下⽂词是 [0, 2, 3]
中⼼词 2 的上下⽂词是 [1, 3]
中⼼词 3 的上下⽂词是 [1, 2, 4, 5]
中⼼词 4 的上下⽂词是 [2, 3, 5, 6]
中⼼词 5 的上下⽂词是 [4, 6]
中⼼词 6 的上下⽂词是 [5]
中⼼词 7 的上下⽂词是 [8]
中⼼词 8 的上下⽂词是 [7, 9]
中⼼词 9 的上下⽂词是 [8]

在PTB数据集上进行训练时，我们将最大上下文窗口大小设置为5。下面提取数据集中的所有中心词及其上下文词

all_centers, all_contexts = get_centers_and_contexts(corpus, 5)
f'# “中⼼词-上下⽂词对”的数量: {sum([len(contexts) for contexts in all_contexts])}'

'# “中⼼词-上下⽂词对”的数量: 1502843'

4 - 负采样

我们使用负采样进行近似训练，为了根据预定义的分布对噪声词进行采样，我们定义一下RandomGenerator类，其中采样分布通过变量sampling_weights传递

#@save
class RandomGenerator:
    """根据n个采样权重在{1,...,n}中随机抽取"""
    def __init__(self, sampling_weights):
        # Exclude
        self.population = list(range(1, len(sampling_weights) + 1))
        self.sampling_weights = sampling_weights
        self.candidates = []
        self.i = 0
    
    def draw(self):
        if self.i == len(self.candidates):
            # 缓存k个随机采样结果
            self.candidates = random.choices(self.population, self.sampling_weights, k=10000)
            self.i = 0
        self.i += 1
        return self.candidates[self.i - 1]

例如，我们可以在索引1、2和3中绘制10个随机变量X，采样概率为P(X = 1) = 2/9, P(X = 2) = 3/9和P(X = 3) = 4/9，如下所⽰

#@save
generator = RandomGenerator([2, 3, 4])
[generator.draw() for _ in range(10)]

[2, 3, 2, 2, 1, 3, 3, 3, 2, 1]

对于中心词和上下文词，我们随机抽取了K个（实验中为5个）噪声词，根据word2vec论文中的建议，将噪声词w的采样概率P(w)设置为其在字典中的相对频率，其幂为0.75

#@save
def get_negatives(all_contexts, vocab, counter, K):
    """返回负采样中的噪声词"""
    # 索引为1、2、...（索引0是词表中排除的未知标记）
    sampling_weights = [counter[vocab.to_tokens(i)]**0.75 for i in range(1, len(vocab))]
    
    all_negatives, generator = [], RandomGenerator(sampling_weights)
    for contexts in all_contexts:
        negatives = []
        while len(negatives) < len(contexts) * K:
            neg = generator.draw()
            # 噪声词不能是上下⽂词
            if neg not in contexts:
                negatives.append(neg)
        all_negatives.append(negatives)
    return all_negatives

all_negatives = get_negatives(all_contexts, vocab, counter, 5)

5 - 小批量加载训练实例

词center、其上下文词context和其噪声词negative组成的样本。此函数返回一个可以在训练期间加载用于计算的小批量，例如包括掩码变量

#@save
def batchify(data):
    """返回带有负采样的跳元模型的⼩批量样本"""
    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).reshape((-1, 1)), torch.tensor(contexts_negatives), torch.tensor(masks), torch.tensor(labels))

让我们使用一个小批量的两个样本来测试此函数

x_1 = (1, [2, 2], [3, 3, 3, 3])
x_2 = (1, [2, 2, 2], [3, 3])
batch = batchify((x_1, x_2))

names = ['centers', 'contexts_negatives', 'masks', 'labels']
for name, data in zip(names, batch):
    print(name, '=', data)

centers = tensor([[1],
        [1]])
contexts_negatives = tensor([[2, 2, 3, 3, 3, 3],
        [2, 2, 2, 3, 3, 0]])
masks = tensor([[1, 1, 1, 1, 1, 1],
        [1, 1, 1, 1, 1, 0]])
labels = tensor([[1, 1, 0, 0, 0, 0],
        [1, 1, 1, 0, 0, 0]])

6 - 整合代码

最后，我们定义了读取PTB数据集并返回数据迭代器和词表的load_data_ptb函数

#@save
def load_data_ptb(batch_size, max_window_size, num_noise_words):
    """下载PTB数据集，然后将其加载到内存中"""
    ##num_workers = d2l.get_dataloader_workers()
    num_workers = 0
    sentences = read_ptb()
    vocab = d2l.Vocab(sentences, min_freq=10)
    subsampled, counter = subsample(sentences, vocab)
    corpus = [vocab[line] for line in subsampled]
    all_centers, all_contexts = get_centers_and_contexts(corpus, max_window_size)
    all_negatives = get_negatives(all_contexts, vocab, counter, num_noise_words)
    
    class PTBDataset(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)
    
    dataset = PTBDataset(all_centers, all_contexts, all_negatives)
    
    data_iter = torch.utils.data.DataLoader(
        dataset, batch_size, shuffle=True,
        collate_fn=batchify, num_workers=num_workers)
    return data_iter, vocab

让我们打印数据迭代器的第一个小批量

data_iter, vocab = load_data_ptb(512, 5, 5)
for batch in data_iter:
    for name, data in zip(names, 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])