基于查找表在PySpark中实现向量平均

嘿，这事儿其实不难，咱们就从最基础的词向量平均法入手实现简单Doc2Vec——毕竟很多场景下这种朴素的方法效果已经够用了。结合你现有的两个RDD，咱们一步步来：

核心思路 #

简单Doc2Vec的核心逻辑就是：把文档里所有单词对应的GloVe向量做平均，得到的结果就作为这个文档的向量。当然你也可以后续扩展成加权平均（比如TF-IDF加权），但先从基础版开始。

具体实现步骤 #

1. 预处理documents RDD #

首先得把documents RDD（格式：(doc_id, word_list)）拆成每个单词关联文档ID的形式，这样才能和words RDD（格式：(word, vector)）做关联。

# 把每个文档拆成(word, doc_id)的键值对，方便后续join
word_doc_pairs = documents.flatMap(lambda x: [(word, x[0]) for word in x[1]])

2. 关联词向量与文档ID #

接下来把上面的RDD和words RDD做join，这样就能得到每个单词对应的文档ID和它的GloVe向量：

# 关联词向量，得到(word, (doc_id, vector))
doc_word_vectors = word_doc_pairs.join(words)

# 转换格式为(doc_id, vector)，方便按文档分组
doc_vector_pairs = doc_word_vectors.map(lambda x: (x[1][0], x[1][1]))

3. 计算文档向量（词向量平均） #

现在按文档ID分组，对每个组里的所有词向量求平均值，就是咱们要的文档向量了：

import numpy as np

def average_vectors(vectors):
    # 将向量列表转换成numpy数组，求平均后再转回列表（如果需要的话）
    return list(np.mean(np.array(vectors), axis=0))

# 按文档ID分组，然后对每组的向量求平均
doc2vec_rdd = doc_vector_pairs.groupByKey().mapValues(average_vectors)

4. 处理OOV词（可选） #

上面的逻辑会自动过滤掉不在GloVe里的单词（因为join的时候匹配不到）。如果想给OOV词分配随机向量，可以先给documents里的每个单词都生成向量，再做平均：

import numpy as np

# 假设GloVe向量的维度是100，你可以改成自己的维度
VECTOR_DIM = 100

def get_word_vector(word, word_vectors_dict):
    # 如果单词在GloVe里就返回对应向量，否则返回随机向量
    return word_vectors_dict.get(word, np.random.rand(VECTOR_DIM).tolist())

# 先把words RDD转成广播变量，提高效率
word_vectors_broadcast = sc.broadcast(words.collectAsMap())

# 直接对每个文档的单词列表计算平均向量
doc2vec_with_oov_rdd = documents.map(lambda x: (
    x[0],
    average_vectors([get_word_vector(word, word_vectors_broadcast.value) for word in x[1]])
))

验证结果 #

你可以用take()方法看看结果：

print(doc2vec_rdd.take(2))

输出应该是类似[('testDoc1', [0.123, 0.456, ...]), ('testDoc2', [0.789, 0.012, ...])]的格式。

内容的提问来源于stack exchange，提问作者Bill Chen