python - 使用查找表在 Pyspark 中平均向量

Question

我正在尝试使用来自 https://nlp.stanford.edu/projects/glove/ 的预训练 GloVe 模型在 PySpark 中实现一个简单的 Doc2Vec 算法.

我有两个 RDD:

• 一个名为 documents 的 RDD 对，形式为 (K:[V])，其中 K 是文档 ID，[V] 是该文档中所有单词的列表，例如 ('testDoc1':'我正在使用 spark') ('testDoc2':'测试 Spark ')

• 一对名为 words 的 RDD 表示 K:V 形式的词嵌入，其中 K 是一个词，V 是表示该词的向量，例如 (‘我’, [0.1, 0.1, 0.1]) (' Spark ': [0.2, 0.2, 0.2]) ('我', [0.3, 0.3, 0.3]) ('测试', [0.5, 0.5, 0.5]) ('使用', [0.4, 0.4, 0.4])

遍历 documents 中的单词以获得所有单词的平均向量和的正确方法是什么？在上面的示例中，最终结果如下所示: ('testDoc1':[0.25, 0.25, 0.25]) ('testDoc2':[0.35, 0.35, 0.35])

Answer 1

假设你有一个函数 tokenize将字符串转换为单词列表。然后你可以flatMap documents得到一个RDD元组 (word, document id) :

flattened_docs = documents.flatMap(lambda x: [(word, x[0]) for word in tokenize(x[1])])

然后加入words会给你(word, (document id, vector))元组，此时您可以删除单词:

doc_vectors = flattened_docs.join(words).values

请注意，这是一个内部连接，因此您将丢弃一个没有嵌入的单词。由于您可能想在平均数中计算这些单词，因此左连接可能更合适，然后您必须替换所有结果 None s 与零向量(或您选择的任何向量)。

我们可以按文档 id 分组以获得 (document id, [list of vectors]) 的 rdd然后求平均值(我假设您有一个名为 average 的函数)。

final_vectors = doc_vectors.groupByKey().mapValues(average)

(请原谅我受 Scala 影响的 Python。自从我使用 pyspark 以来已经有一段时间了，我还没有检查它是 flatMap 还是 flat_map 等等。)