Pyspark中的平均向量与查找表

Posted 2021-05-01

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

我有两个RDD：

• 一对RDD称为documents形式（K：[V]），其中K是文档ID，[V]是该文档中所有单词的列表，例如 ('testDoc1':'i am using spark') ('testDoc2':'testing spark')
• 一对称为words的RDD，表示以K：V形式嵌入的单词，其中K是单词，V是表示单词的向量，例如 ('i', [0.1, 0.1, 0.1]) ('spark': [0.2, 0.2, 0.2]) ('am', [0.3, 0.3, 0.3]) ('testing', [0.5, 0.5, 0.5]) ('using', [0.4, 0.4, 0.4])

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

答案

假设你有一个函数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

请注意，这是一个内部联接，因此您丢弃了没有嵌入的单词。因为你可能想要计算平均值中的那些单词，所以左连接可能更合适，然后你必须用零向量（或你选择的任何向量）替换任何得到的Nones。

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

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

（请原谅我受Scala影响的Python。已经有一段时间了，因为我使用了pyspark而且我没有检查它是否是flatMap或flat_map等等。）