将spark word2vec向量转储到文件中

时间:2016-12-06 09:32:12

标签: java apache-spark rdd apache-spark-mllib word2vec

我使用spark mllib生成单词向量。我希望能够拟合我的所有数据,然后获得经过训练的单词向量并将它们转储到文件中。

我这样做:

JavaRDD<List<String>> data = javaSparkContext.parallelize(streamingData, partitions);
Word2Vec word2vec = new Word2Vec();
Word2VecModel model = word2vec.fit(data);

所以,如果我的训练数据有像

这样的句子
I love Spark

我想将文件中的输出保存为:

I       0.03 0.53 0.12...
love    0.31 0.14 0.12...
Spark   0.41 0.18 0.84...

训练结束后,我从模型对象中获取矢量

Map<String, float[]> wordMap = JavaConverters.mapAsJavaMapConverter(model.getVectors()).asJava();
List<String> wordvectorlist = Lists.newArrayList();
for (String s : wordMap.keySet()) {
    StringBuilder wordvector = new StringBuilder(s);
    for (float f : wordMap.get(s)) {
        wordvector.append(" " + f);
    }
    wordvectorlist.add(wordvector.toString());
    if (wordvectorlist.size() > 1000000) {
        writeToFile(wordvectorlist);
        wordvectorlist.clear();
    }

}

我将为非常大的数据(~1.5 TB)生成这些单词向量,因此,我可能无法将返回的对象Word2VecModel保存在我的驱动程序的内存中。 如何将此wordvectors映射存储为rdd,以便我可以在不将完整映射存储在驱动程序内存中的情况下写入文件?

我查看了deeplearning4j的word2vec实现,但该实现还需要加载驱动程序内存中的所有向量。

1 个答案:

答案 0 :(得分:3)

Word2VecModel具有save功能,可以以自己的格式将其保存到磁盘 这将创建一个名为data的目录,其中包含数据的镶木地板文件和具有人类可读元数据的元数据文件。

您现在可以阅读镶木地板文件并自行转换,或者使用spark.read.parquet将其读取到数据框。每行都包含一些地图,您可以按照自己的意愿编写它。