word2vec

在 TensorFlow.org 上查看 在 Google Colab 中运行 在 GitHub 上查看 下载笔记本

word2vec 并不是一种单一的算法,而是一系列模型架构和优化方法的统称,可用于从大型数据集中学习词嵌入(word embeddings)。通过 word2vec 学习到的嵌入已被证明在各种下游自然语言处理任务中非常成功。

这些论文提出了两种学习词表示的方法:

  • 连续词袋模型 (Continuous bag-of-words model):根据周围的上下文词预测中间词。上下文由当前(中间)词之前和之后的几个词组成。这种架构被称为词袋模型,因为上下文中词的顺序并不重要。
  • 连续 Skip-gram 模型 (Continuous skip-gram model):预测同一句子中当前词之前和之后一定范围内的词。下面给出了一个具体的工作示例。

在本教程中,你将使用 Skip-gram 方法。首先,你将使用一个单句来探索 Skip-gram 和其他概念。接下来,你将在一个小数据集上训练自己的 word2vec 模型。本教程还包含用于导出训练好的嵌入并在 TensorFlow Embedding Projector 中进行可视化的代码。

Skip-gram 与负采样

虽然词袋模型是在给定邻近上下文的情况下预测一个词,但 Skip-gram 模型是在给定词本身的情况下预测其上下文(或邻居)。该模型基于 Skip-gram 进行训练,Skip-gram 是一种允许跳过词元(token)的 n-gram(如下图示例所示)。一个词的上下文可以用一组 (target_word, context_word) 的 Skip-gram 对来表示,其中 context_word 出现在 target_word 的邻近上下文中。

考虑以下包含八个词的句子:

The wide road shimmered in the hot sun.

句子中 8 个词的每一个词的上下文词都是由窗口大小定义的。窗口大小决定了 target_word 两侧可以被视为 context word 的词的范围。下表列出了基于不同窗口大小的目标词的 Skip-gram。

word2vec_skipgrams

Skip-gram 模型的训练目标是最大化在给定目标词的情况下预测上下文词的概率。对于词序列 w1, w2, ... wT,目标可以写成平均对数概率:

word2vec_skipgram_objective

其中 c 是训练上下文的大小。基本的 Skip-gram 公式使用 softmax 函数定义此概率。

word2vec_full_softmax

其中 vv' 分别是词的目标向量和上下文向量表示,W 是词汇表的大小。

计算该公式的分母涉及对整个词汇表执行完整的 softmax,而词汇表通常非常大(105-107 个词项)。

噪声对比估计 (NCE) 损失函数是对完整 softmax 的有效近似。由于目标是学习词嵌入而不是对词分布进行建模,因此 NCE 损失可以简化为使用负采样。

目标词的简化负采样目标是从词的噪声分布 Pn(w) 中抽取 num_ns 个负样本,从而将上下文词与这些负样本区分开来。更准确地说,对词汇表完整 softmax 的一种有效近似是:对于一个 Skip-gram 对,将目标词的损失设定为一个分类问题,即在上下文词和 num_ns 个负样本之间进行分类。

负样本定义为一个 (target_word, context_word) 对,使得 context_word 不出现在 target_wordwindow_size 邻域内。对于示例句子,当 window_size2 时,以下是一些潜在的负样本。

(hot, shimmered)
(wide, hot)
(wide, sun)

在下一节中,你将为一个句子生成 Skip-gram 和负样本。你还将了解子采样技术,并在本教程的稍后部分训练一个用于正样本和负样本训练的分类模型。

设置

import io
import re
import string
import tqdm

import numpy as np

import tensorflow as tf
from tensorflow.keras import layers
2024-07-19 11:19:50.266209: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:485] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered
2024-07-19 11:19:50.287751: E external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:8454] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered
2024-07-19 11:19:50.294095: E external/local_xla/xla/stream_executor/cuda/cuda_blas.cc:1452] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered
# Load the TensorBoard notebook extension
%load_ext tensorboard
SEED = 42
AUTOTUNE = tf.data.AUTOTUNE

向量化一个示例句子

考虑以下句子:

The wide road shimmered in the hot sun.

对句子进行词元化(Tokenize)

sentence = "The wide road shimmered in the hot sun"
tokens = list(sentence.lower().split())
print(len(tokens))
8

创建一个词汇表,用于保存从词元到整数索引的映射

vocab, index = {}, 1  # start indexing from 1
vocab['<pad>'] = 0  # add a padding token
for token in tokens:
  if token not in vocab:
    vocab[token] = index
    index += 1
vocab_size = len(vocab)
print(vocab)
{'<pad>': 0, 'the': 1, 'wide': 2, 'road': 3, 'shimmered': 4, 'in': 5, 'hot': 6, 'sun': 7}

创建一个反向词汇表,用于保存从整数索引到词元的映射

inverse_vocab = {index: token for token, index in vocab.items()}
print(inverse_vocab)
{0: '<pad>', 1: 'the', 2: 'wide', 3: 'road', 4: 'shimmered', 5: 'in', 6: 'hot', 7: 'sun'}

向量化你的句子

example_sequence = [vocab[word] for word in tokens]
print(example_sequence)
[1, 2, 3, 4, 5, 1, 6, 7]

从一个句子生成 Skip-gram

tf.keras.preprocessing.sequence 模块提供了一些有用的函数,可以简化 word2vec 的数据准备工作。你可以使用 tf.keras.preprocessing.sequence.skipgramsexample_sequence 中生成给定 window_size 的 Skip-gram 对,词元范围为 [0, vocab_size)

window_size = 2
positive_skip_grams, _ = tf.keras.preprocessing.sequence.skipgrams(
      example_sequence,
      vocabulary_size=vocab_size,
      window_size=window_size,
      negative_samples=0)
print(len(positive_skip_grams))
26

打印一些正 Skip-gram

for target, context in positive_skip_grams[:5]:
  print(f"({target}, {context}): ({inverse_vocab[target]}, {inverse_vocab[context]})")
(2, 1): (wide, the)
(2, 3): (wide, road)
(1, 6): (the, hot)
(3, 5): (road, in)
(1, 7): (the, sun)

单个 Skip-gram 的负采样

skipgrams 函数通过在给定窗口跨度上滑动,返回所有的正 Skip-gram 对。为了生成额外的 Skip-gram 对作为训练用的负样本,你需要从词汇表中随机采样词。使用 tf.random.log_uniform_candidate_sampler 函数为窗口中给定的目标词采样 num_ns 个负样本。你可以对一个 Skip-gram 的目标词调用此函数,并将上下文词作为“真类”(true class)传入,以将其排除在采样之外。

# Get target and context words for one positive skip-gram.
target_word, context_word = positive_skip_grams[0]

# Set the number of negative samples per positive context.
num_ns = 4

context_class = tf.reshape(tf.constant(context_word, dtype="int64"), (1, 1))
negative_sampling_candidates, _, _ = tf.random.log_uniform_candidate_sampler(
    true_classes=context_class,  # class that should be sampled as 'positive'
    num_true=1,  # each positive skip-gram has 1 positive context class
    num_sampled=num_ns,  # number of negative context words to sample
    unique=True,  # all the negative samples should be unique
    range_max=vocab_size,  # pick index of the samples from [0, vocab_size]
    seed=SEED,  # seed for reproducibility
    name="negative_sampling"  # name of this operation
)
print(negative_sampling_candidates)
print([inverse_vocab[index.numpy()] for index in negative_sampling_candidates])
WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
I0000 00:00:1721387992.808839    9940 cuda_executor.cc:1015] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355
I0000 00:00:1721387992.812783    9940 cuda_executor.cc:1015] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355
I0000 00:00:1721387992.816433    9940 cuda_executor.cc:1015] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355
I0000 00:00:1721387992.819973    9940 cuda_executor.cc:1015] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355
I0000 00:00:1721387992.831306    9940 cuda_executor.cc:1015] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355
I0000 00:00:1721387992.834967    9940 cuda_executor.cc:1015] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355
I0000 00:00:1721387992.838234    9940 cuda_executor.cc:1015] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355
I0000 00:00:1721387992.841501    9940 cuda_executor.cc:1015] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355
I0000 00:00:1721387992.844863    9940 cuda_executor.cc:1015] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355
I0000 00:00:1721387992.848342    9940 cuda_executor.cc:1015] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355
I0000 00:00:1721387992.851679    9940 cuda_executor.cc:1015] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355
I0000 00:00:1721387992.855085    9940 cuda_executor.cc:1015] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355
I0000 00:00:1721387994.077319    9940 cuda_executor.cc:1015] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355
I0000 00:00:1721387994.079463    9940 cuda_executor.cc:1015] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355
I0000 00:00:1721387994.081515    9940 cuda_executor.cc:1015] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355
I0000 00:00:1721387994.083524    9940 cuda_executor.cc:1015] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355
I0000 00:00:1721387994.085555    9940 cuda_executor.cc:1015] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355
I0000 00:00:1721387994.087555    9940 cuda_executor.cc:1015] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355
I0000 00:00:1721387994.089481    9940 cuda_executor.cc:1015] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355
I0000 00:00:1721387994.091388    9940 cuda_executor.cc:1015] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355
I0000 00:00:1721387994.093328    9940 cuda_executor.cc:1015] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355
I0000 00:00:1721387994.095306    9940 cuda_executor.cc:1015] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355
I0000 00:00:1721387994.097243    9940 cuda_executor.cc:1015] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355
I0000 00:00:1721387994.099177    9940 cuda_executor.cc:1015] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355
I0000 00:00:1721387994.136509    9940 cuda_executor.cc:1015] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355
I0000 00:00:1721387994.138618    9940 cuda_executor.cc:1015] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355
I0000 00:00:1721387994.140595    9940 cuda_executor.cc:1015] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355
I0000 00:00:1721387994.142565    9940 cuda_executor.cc:1015] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355
I0000 00:00:1721387994.144526    9940 cuda_executor.cc:1015] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355
I0000 00:00:1721387994.146537    9940 cuda_executor.cc:1015] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355
I0000 00:00:1721387994.148480    9940 cuda_executor.cc:1015] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355
tf.Tensor([2 1 4 3], shape=(4,), dtype=int64)
['wide', 'the', 'shimmered', 'road']
I0000 00:00:1721387994.150422    9940 cuda_executor.cc:1015] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355
I0000 00:00:1721387994.152419    9940 cuda_executor.cc:1015] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355
I0000 00:00:1721387994.154834    9940 cuda_executor.cc:1015] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355
I0000 00:00:1721387994.157208    9940 cuda_executor.cc:1015] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355
I0000 00:00:1721387994.159570    9940 cuda_executor.cc:1015] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355

构建一个训练示例

对于给定的正 (target_word, context_word) Skip-gram,你现在还有 num_ns 个未出现在 target_word 窗口邻域内的负采样上下文词。将 1 个正 context_wordnum_ns 个负上下文词批处理到一个张量中。这将为每个目标词生成一组正 Skip-gram(标记为 1)和负样本(标记为 0)。

# Reduce a dimension so you can use concatenation (in the next step).
squeezed_context_class = tf.squeeze(context_class, 1)

# Concatenate a positive context word with negative sampled words.
context = tf.concat([squeezed_context_class, negative_sampling_candidates], 0)

# Label the first context word as `1` (positive) followed by `num_ns` `0`s (negative).
label = tf.constant([1] + [0]*num_ns, dtype="int64")
target = target_word

查看上述 Skip-gram 示例中目标词的上下文及其对应的标签

print(f"target_index    : {target}")
print(f"target_word     : {inverse_vocab[target_word]}")
print(f"context_indices : {context}")
print(f"context_words   : {[inverse_vocab[c.numpy()] for c in context]}")
print(f"label           : {label}")
target_index    : 2
target_word     : wide
context_indices : [1 2 1 4 3]
context_words   : ['the', 'wide', 'the', 'shimmered', 'road']
label           : [1 0 0 0 0]

(target, context, label) 张量元组构成了训练你的 Skip-gram 负采样 word2vec 模型的一个训练示例。请注意,target 的形状为 (1,),而 context 和 label 的形状为 (1+num_ns,)

print("target  :", target)
print("context :", context)
print("label   :", label)
target  : 2
context : tf.Tensor([1 2 1 4 3], shape=(5,), dtype=int64)
label   : tf.Tensor([1 0 0 0 0], shape=(5,), dtype=int64)

总结

此图总结了从句子生成训练示例的过程:

word2vec_negative_sampling

请注意,单词 temperaturecode 不在输入句子中。它们与上图中使用的某些其他索引一样,属于词汇表。

将所有步骤整合到一个函数中

Skip-gram 采样表

大型数据集意味着更大的词汇表,其中包含更多的常用词(如停用词)。从频繁出现的词(如 theison)采样得到的训练示例并不能为模型学习增加太多有用的信息。Mikolov 等人建议将常用词的子采样作为提高嵌入质量的一种有效实践。

tf.keras.preprocessing.sequence.skipgrams 函数接受一个采样表参数来编码采样任何词元的概率。你可以使用 tf.keras.preprocessing.sequence.make_sampling_table 生成一个基于词频排序的概率采样表,并将其传递给 skipgrams 函数。查看 vocab_size 为 10 时的采样概率。

sampling_table = tf.keras.preprocessing.sequence.make_sampling_table(size=10)
print(sampling_table)
[0.00315225 0.00315225 0.00547597 0.00741556 0.00912817 0.01068435
 0.01212381 0.01347162 0.01474487 0.0159558 ]

sampling_table[i] 表示采样数据集中第 i 个最常用词的概率。该函数假设词频遵循齐夫定律 (Zipf's distribution)

生成训练数据

将上述所有步骤整合到一个函数中,该函数可以对从任何文本数据集中获得的向量化句子列表进行调用。请注意,采样表是在对 Skip-gram 词对进行采样之前构建的。你将在后面的章节中使用此函数。

# Generates skip-gram pairs with negative sampling for a list of sequences
# (int-encoded sentences) based on window size, number of negative samples
# and vocabulary size.
def generate_training_data(sequences, window_size, num_ns, vocab_size, seed):
  # Elements of each training example are appended to these lists.
  targets, contexts, labels = [], [], []

  # Build the sampling table for `vocab_size` tokens.
  sampling_table = tf.keras.preprocessing.sequence.make_sampling_table(vocab_size)

  # Iterate over all sequences (sentences) in the dataset.
  for sequence in tqdm.tqdm(sequences):

    # Generate positive skip-gram pairs for a sequence (sentence).
    positive_skip_grams, _ = tf.keras.preprocessing.sequence.skipgrams(
          sequence,
          vocabulary_size=vocab_size,
          sampling_table=sampling_table,
          window_size=window_size,
          negative_samples=0)

    # Iterate over each positive skip-gram pair to produce training examples
    # with a positive context word and negative samples.
    for target_word, context_word in positive_skip_grams:
      context_class = tf.expand_dims(
          tf.constant([context_word], dtype="int64"), 1)
      negative_sampling_candidates, _, _ = tf.random.log_uniform_candidate_sampler(
          true_classes=context_class,
          num_true=1,
          num_sampled=num_ns,
          unique=True,
          range_max=vocab_size,
          seed=seed,
          name="negative_sampling")

      # Build context and label vectors (for one target word)
      context = tf.concat([tf.squeeze(context_class,1), negative_sampling_candidates], 0)
      label = tf.constant([1] + [0]*num_ns, dtype="int64")

      # Append each element from the training example to global lists.
      targets.append(target_word)
      contexts.append(context)
      labels.append(label)

  return targets, contexts, labels

为 word2vec 准备训练数据

在了解了如何处理单个句子以用于 Skip-gram 负采样 word2vec 模型后,你可以继续从更大的句子列表中生成训练示例!

下载文本语料库

在本教程中,你将使用一份莎士比亚作品的文本文件。若要在你自己的数据上运行此代码,请修改以下行。

path_to_file = tf.keras.utils.get_file('shakespeare.txt', 'https://storage.googleapis.com/download.tensorflow.org/data/shakespeare.txt')

读取文件中的文本并打印前几行

with open(path_to_file) as f:
  lines = f.read().splitlines()
for line in lines[:20]:
  print(line)
First Citizen:
Before we proceed any further, hear me speak.

All:
Speak, speak.

First Citizen:
You are all resolved rather to die than to famish?

All:
Resolved. resolved.

First Citizen:
First, you know Caius Marcius is chief enemy to the people.

All:
We know't, we know't.

First Citizen:
Let us kill him, and we'll have corn at our own price.

使用非空行构建一个 tf.data.TextLineDataset 对象,以进行后续步骤

text_ds = tf.data.TextLineDataset(path_to_file).filter(lambda x: tf.cast(tf.strings.length(x), bool))

向量化语料库中的句子

你可以使用 TextVectorization 层来向量化语料库中的句子。在此 文本分类 教程中详细了解如何使用此层。从上面的前几个句子中可以看出,文本需要统一大小写并移除标点符号。为此,请定义一个可以在 TextVectorization 层中使用的 custom_standardization 函数

# Now, create a custom standardization function to lowercase the text and
# remove punctuation.
def custom_standardization(input_data):
  lowercase = tf.strings.lower(input_data)
  return tf.strings.regex_replace(lowercase,
                                  '[%s]' % re.escape(string.punctuation), '')


# Define the vocabulary size and the number of words in a sequence.
vocab_size = 4096
sequence_length = 10

# Use the `TextVectorization` layer to normalize, split, and map strings to
# integers. Set the `output_sequence_length` length to pad all samples to the
# same length.
vectorize_layer = layers.TextVectorization(
    standardize=custom_standardization,
    max_tokens=vocab_size,
    output_mode='int',
    output_sequence_length=sequence_length)

在文本数据集上调用 TextVectorization.adapt 以创建词汇表。

vectorize_layer.adapt(text_ds.batch(1024))

一旦该层的状态已适应于表示文本语料库,就可以使用 TextVectorization.get_vocabulary 访问词汇表。此函数返回按频率(降序)排序的所有词汇表词元列表。

# Save the created vocabulary for reference.
inverse_vocab = vectorize_layer.get_vocabulary()
print(inverse_vocab[:20])
['', '[UNK]', 'the', 'and', 'to', 'i', 'of', 'you', 'my', 'a', 'that', 'in', 'is', 'not', 'for', 'with', 'me', 'it', 'be', 'your']

vectorize_layer 现在可以用于为 text_ds(一个 tf.data.Dataset)中的每个元素生成向量。应用 Dataset.batchDataset.prefetchDataset.mapDataset.unbatch

# Vectorize the data in text_ds.
text_vector_ds = text_ds.batch(1024).prefetch(AUTOTUNE).map(vectorize_layer).unbatch()

从数据集中获取序列

你现在拥有一个整数编码句子的 tf.data.Dataset。为准备训练 word2vec 模型的数据集,将数据集展平为句子向量序列列表。由于你需要遍历数据集中的每个句子以生成正样本和负样本,因此这一步是必需的。

sequences = list(text_vector_ds.as_numpy_iterator())
print(len(sequences))
32777

检查 sequences 中的一些示例

for seq in sequences[:5]:
  print(f"{seq} => {[inverse_vocab[i] for i in seq]}")
[ 89 270   0   0   0   0   0   0   0   0] => ['first', 'citizen', '', '', '', '', '', '', '', '']
[138  36 982 144 673 125  16 106   0   0] => ['before', 'we', 'proceed', 'any', 'further', 'hear', 'me', 'speak', '', '']
[34  0  0  0  0  0  0  0  0  0] => ['all', '', '', '', '', '', '', '', '', '']
[106 106   0   0   0   0   0   0   0   0] => ['speak', 'speak', '', '', '', '', '', '', '', '']
[ 89 270   0   0   0   0   0   0   0   0] => ['first', 'citizen', '', '', '', '', '', '', '', '']

从序列生成训练示例

sequences 现在是一个整数编码句子的列表。只需调用前面定义的 generate_training_data 函数即可为 word2vec 模型生成训练示例。概括一下,该函数会遍历每个序列中的每个词,以收集正向和负向的上下文词。target、contexts 和 labels 的长度应该相同,代表训练示例的总数。

targets, contexts, labels = generate_training_data(
    sequences=sequences,
    window_size=2,
    num_ns=4,
    vocab_size=vocab_size,
    seed=SEED)

targets = np.array(targets)
contexts = np.array(contexts)
labels = np.array(labels)

print('\n')
print(f"targets.shape: {targets.shape}")
print(f"contexts.shape: {contexts.shape}")
print(f"labels.shape: {labels.shape}")
100%|██████████| 32777/32777 [00:38<00:00, 850.63it/s]
targets.shape: (65204,)
contexts.shape: (65204, 5)
labels.shape: (65204, 5)

配置数据集以获得高性能

为了对可能数量巨大的训练示例进行高效的批处理,请使用 tf.data.Dataset API。完成此步骤后,你将拥有一个 (target_word, context_word), (label) 元素的 tf.data.Dataset 对象,用于训练你的 word2vec 模型!

BATCH_SIZE = 1024
BUFFER_SIZE = 10000
dataset = tf.data.Dataset.from_tensor_slices(((targets, contexts), labels))
dataset = dataset.shuffle(BUFFER_SIZE).batch(BATCH_SIZE, drop_remainder=True)
print(dataset)
<_BatchDataset element_spec=((TensorSpec(shape=(1024,), dtype=tf.int64, name=None), TensorSpec(shape=(1024, 5), dtype=tf.int64, name=None)), TensorSpec(shape=(1024, 5), dtype=tf.int64, name=None))>

应用 Dataset.cacheDataset.prefetch 以提高性能

dataset = dataset.cache().prefetch(buffer_size=AUTOTUNE)
print(dataset)
<_PrefetchDataset element_spec=((TensorSpec(shape=(1024,), dtype=tf.int64, name=None), TensorSpec(shape=(1024, 5), dtype=tf.int64, name=None)), TensorSpec(shape=(1024, 5), dtype=tf.int64, name=None))>

模型与训练

word2vec 模型可以实现为一个分类器,用于区分 Skip-gram 中的真实上下文词和通过负采样获得的虚假上下文词。你可以对目标词和上下文词的嵌入进行点积乘法,以获得标签的预测值,并根据数据集中的真实标签计算损失函数。

子类化 word2vec 模型

使用 Keras 子类化 API 定义你的 word2vec 模型,包含以下层:

  • target_embedding:一个 tf.keras.layers.Embedding 层,当单词作为目标词出现时,它会查找该词的嵌入。该层的参数数量为 (vocab_size * embedding_dim)
  • context_embedding:另一个 tf.keras.layers.Embedding 层,当单词作为上下文词出现时,它会查找该词的嵌入。该层的参数数量与 target_embedding 相同,即 (vocab_size * embedding_dim)
  • dots:一个 tf.keras.layers.Dot 层,用于计算训练对中目标嵌入和上下文嵌入的点积。
  • flatten:一个 tf.keras.layers.Flatten 层,用于将 dots 层的输出结果展平为 logits。

使用子类化模型,你可以定义接受 (target, context) 对的 call() 函数,这些对随后可以传入相应的嵌入层。重塑 context_embedding 以便与 target_embedding 执行点积,并返回展平后的结果。

class Word2Vec(tf.keras.Model):
  def __init__(self, vocab_size, embedding_dim):
    super(Word2Vec, self).__init__()
    self.target_embedding = layers.Embedding(vocab_size,
                                      embedding_dim,
                                      name="w2v_embedding")
    self.context_embedding = layers.Embedding(vocab_size,
                                       embedding_dim)

  def call(self, pair):
    target, context = pair
    # target: (batch, dummy?)  # The dummy axis doesn't exist in TF2.7+
    # context: (batch, context)
    if len(target.shape) == 2:
      target = tf.squeeze(target, axis=1)
    # target: (batch,)
    word_emb = self.target_embedding(target)
    # word_emb: (batch, embed)
    context_emb = self.context_embedding(context)
    # context_emb: (batch, context, embed)
    dots = tf.einsum('be,bce->bc', word_emb, context_emb)
    # dots: (batch, context)
    return dots

定义损失函数并编译模型

为简单起见,你可以使用 tf.keras.losses.CategoricalCrossEntropy 作为负采样损失的替代方案。如果你想编写自己的自定义损失函数,也可以按如下方式进行:

def custom_loss(x_logit, y_true):
      return tf.nn.sigmoid_cross_entropy_with_logits(logits=x_logit, labels=y_true)

现在是构建模型的时候了!实例化你的 word2vec 类,并将嵌入维度设置为 128(你可以尝试不同的值)。使用 tf.keras.optimizers.Adam 优化器编译模型。

embedding_dim = 128
word2vec = Word2Vec(vocab_size, embedding_dim)
word2vec.compile(optimizer='adam',
                 loss=tf.keras.losses.CategoricalCrossentropy(from_logits=True),
                 metrics=['accuracy'])

同时定义一个回调函数,以记录 TensorBoard 的训练统计信息

tensorboard_callback = tf.keras.callbacks.TensorBoard(log_dir="logs")

dataset 上训练模型若干个 epoch

word2vec.fit(dataset, epochs=20, callbacks=[tensorboard_callback])
Epoch 1/20
WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
I0000 00:00:1721388042.283243   10107 service.cc:146] XLA service 0x7f343003f620 initialized for platform CUDA (this does not guarantee that XLA will be used). Devices:
I0000 00:00:1721388042.283295   10107 service.cc:154]   StreamExecutor device (0): Tesla T4, Compute Capability 7.5
I0000 00:00:1721388042.283299   10107 service.cc:154]   StreamExecutor device (1): Tesla T4, Compute Capability 7.5
I0000 00:00:1721388042.283302   10107 service.cc:154]   StreamExecutor device (2): Tesla T4, Compute Capability 7.5
I0000 00:00:1721388042.283305   10107 service.cc:154]   StreamExecutor device (3): Tesla T4, Compute Capability 7.5
33/63 ━━━━━━━━━━━━━━━━━━━━ 0s 3ms/step - accuracy: 0.2126 - loss: 1.6091
I0000 00:00:1721388042.958698   10107 device_compiler.h:188] Compiled cluster using XLA!  This line is logged at most once for the lifetime of the process.
63/63 ━━━━━━━━━━━━━━━━━━━━ 2s 8ms/step - accuracy: 0.2203 - loss: 1.6088
Epoch 2/20
63/63 ━━━━━━━━━━━━━━━━━━━━ 0s 2ms/step - accuracy: 0.5959 - loss: 1.5896
Epoch 3/20
63/63 ━━━━━━━━━━━━━━━━━━━━ 0s 2ms/step - accuracy: 0.6070 - loss: 1.5326
Epoch 4/20
63/63 ━━━━━━━━━━━━━━━━━━━━ 0s 2ms/step - accuracy: 0.5575 - loss: 1.4445
Epoch 5/20
63/63 ━━━━━━━━━━━━━━━━━━━━ 0s 2ms/step - accuracy: 0.5715 - loss: 1.3493
Epoch 6/20
63/63 ━━━━━━━━━━━━━━━━━━━━ 0s 2ms/step - accuracy: 0.6062 - loss: 1.2539
Epoch 7/20
63/63 ━━━━━━━━━━━━━━━━━━━━ 0s 2ms/step - accuracy: 0.6419 - loss: 1.1641
Epoch 8/20
63/63 ━━━━━━━━━━━━━━━━━━━━ 0s 2ms/step - accuracy: 0.6795 - loss: 1.0810
Epoch 9/20
63/63 ━━━━━━━━━━━━━━━━━━━━ 0s 2ms/step - accuracy: 0.7109 - loss: 1.0043
Epoch 10/20
63/63 ━━━━━━━━━━━━━━━━━━━━ 0s 2ms/step - accuracy: 0.7388 - loss: 0.9336
Epoch 11/20
63/63 ━━━━━━━━━━━━━━━━━━━━ 0s 2ms/step - accuracy: 0.7625 - loss: 0.8683
Epoch 12/20
63/63 ━━━━━━━━━━━━━━━━━━━━ 0s 2ms/step - accuracy: 0.7834 - loss: 0.8081
Epoch 13/20
63/63 ━━━━━━━━━━━━━━━━━━━━ 0s 2ms/step - accuracy: 0.8017 - loss: 0.7527
Epoch 14/20
63/63 ━━━━━━━━━━━━━━━━━━━━ 0s 2ms/step - accuracy: 0.8163 - loss: 0.7018
Epoch 15/20
63/63 ━━━━━━━━━━━━━━━━━━━━ 0s 2ms/step - accuracy: 0.8322 - loss: 0.6552
Epoch 16/20
63/63 ━━━━━━━━━━━━━━━━━━━━ 0s 2ms/step - accuracy: 0.8457 - loss: 0.6125
Epoch 17/20
63/63 ━━━━━━━━━━━━━━━━━━━━ 0s 2ms/step - accuracy: 0.8576 - loss: 0.5734
Epoch 18/20
63/63 ━━━━━━━━━━━━━━━━━━━━ 0s 2ms/step - accuracy: 0.8688 - loss: 0.5377
Epoch 19/20
63/63 ━━━━━━━━━━━━━━━━━━━━ 0s 2ms/step - accuracy: 0.8783 - loss: 0.5052
Epoch 20/20
63/63 ━━━━━━━━━━━━━━━━━━━━ 0s 2ms/step - accuracy: 0.8878 - loss: 0.4754
<keras.src.callbacks.history.History at 0x7f35ec96eee0>

TensorBoard 现在显示了 word2vec 模型的准确率和损失

#docs_infra: no_execute
%tensorboard --logdir logs

嵌入查找与分析

使用 Model.get_layerLayer.get_weights 从模型中获取权重。TextVectorization.get_vocabulary 函数提供了用于构建元数据文件的词汇表,每行包含一个词元。

weights = word2vec.get_layer('w2v_embedding').get_weights()[0]
vocab = vectorize_layer.get_vocabulary()

创建并保存向量和元数据文件

out_v = io.open('vectors.tsv', 'w', encoding='utf-8')
out_m = io.open('metadata.tsv', 'w', encoding='utf-8')

for index, word in enumerate(vocab):
  if index == 0:
    continue  # skip 0, it's padding.
  vec = weights[index]
  out_v.write('\t'.join([str(x) for x in vec]) + "\n")
  out_m.write(word + "\n")
out_v.close()
out_m.close()

下载 vectors.tsvmetadata.tsv,以便在 Embedding Projector 中分析获得的嵌入。

try:
  from google.colab import files
  files.download('vectors.tsv')
  files.download('metadata.tsv')
except Exception:
  pass

后续步骤

本教程向你展示了如何从零开始实现带负采样的 Skip-gram word2vec 模型,并可视化所获得的词嵌入。