【pyspark系列】pyspark.ml Word2Vec使用总结

pyspark版word2vec

一、导入 word2vec

import pyspark
from pyspark.sql.functions import format_number as fmt
from pyspark.sql.types import FloatType,StringType,ArrayType
import pyspark.sql
from pyspark.sql import SparkSession
from scipy import spatial
from pyspark.sql.functions import regexp_replace ,collect_list,concat_ws,lit,udf,col
from pyspark.ml.feature import Word2Vec,Word2VecModel

二、训练词向量

#训练词向量
data = spark.sql(select * from xxxx)
mm = Word2Vec(vectorSize=128, minCount=1, numPartitions = 300,windowSize=5,inputCol="pkglist", outputCol="result")
model = mm.fit(data)

#保存向量
data1=model.getVectors()
data1.write.saveAsTable('xxxxxxx')

三、pyspark save model and load model

#保存模型
temp_path = "/user/xxxx"
modelPath = temp_path + "/word2vec-model"
model.save(modelPath)

# 读取模型
model = Word2VecModel.load(temp_path)

四、词向量映射
transform，对映射后的向量序列直接取了平均；

data2 = spark.sql('''select * from xxx''')
vec=model.transform(data2)