让我们从零开始使用PyTorch构建一个轻量级的词嵌入模型

图片来源：Phil Hearing（Unsplash）

在我之前写的一篇文章中，我们学习了如何使用PyTorch的nn.Embedding层将单词转换为稠密向量。但由于该嵌入层是未经训练的，这些向量并没有任何语义含义。默认情况下，PyTorch会使用随机权重初始化嵌入层，这使得生成的向量基本上是毫无意义的。今天，我们要改变这一点，并希望赋予我们的嵌入一些真正的含义。

我们将创建一个微型Word2Vec模型。换句话说，我们将使用Word2Vec方法来训练我们的嵌入层。

Word2Vec是一种非常流行的方法，用于从大型文本语料库中学习单词的向量表示。这些向量携带语义信息，使得具有相似含义的单词被映射到向量空间中的相邻点。例如，“happy”、“joyful”和“cheerful”这些词的嵌入会彼此接近。同样，“car”、“automobile”和“vehicle”也会聚集在一起。

这些向量实际上反映了训练好的Word2Vec模型如何内化各种概念，比如性别差异、地理关系以及类别分组（比如动物、水果、衣物等），这些概念对于自然语言理解至关重要。

测试并使用自己训练的单词嵌入是一种奇妙的体验。想象一下，你可以用向量数学进行推理，比如这样：

Technology + Nature - Industry = Sustainability.

当然，构建一个高性能、全球规模的Word2Vec模型通常是那些拥有大量资源的公司在做的事情。这需要处理TB级的文本，并利用数百个GPU或TPU进行训练。因此，我们将创建一个微型的玩具模型，仅作为一个概念验证（POC）。

下载数据集

我的训练数据来自一本书。这本书可以在Project Gutenberg免费下载，它的名字是 “The Oxford Book of American Essays by Brander Matthews et al”. 你也可以选择其他书籍，甚至可以合并多本书的内容。

我使用wget下载了这本书，并统计了其中的行数、单词数和字符数：

$ wget -O dataset.txt 'https://www.gutenberg.org/ebooks/40196.txt.utf-8'

$ wc -lwc dataset.txt

16020 166444 974113 dataset.txt

处理数据集

现在，我们需要确定数据集中包含的词汇。我们将使用nltk来完成这个任务。思路很简单，就是找到数据集中所有唯一的标记（tokens）。

import nltk

from nltk.tokenize import word_tokenize

from pathlib import Path

nltk.download('punkt_tab')

dataset = Path("./dataset.txt").read_text()

tokens = word_tokenize(dataset.lower())

vocab = sorted(list(set(tokens)))

vocab_size = len(vocab)

print(f"Number of tokens: {len(tokens)}")

print(f"Vocabulary size: {vocab_size}")

p