Spark Word2Vec算法代码实现

 1 import com.hankcs.hanlp.tokenizer.NLPTokenizer
 import org.apache.hadoop.io.{LongWritable, Text}
 import org.apache.hadoop.mapred.TextInputFormat
 import org.apache.log4j.{Level, Logger}
 import org.apache.spark.ml.feature.Word2Vec
 import org.apache.spark.sql.SparkSession

 /**
   * Created by zhen on 2018/11/20.
   */
 object Word2Vec {
   Logger.getLogger("org").setLevel(Level.WARN) // 设置日志级别
   def main(args: Array[String]) {
     val spark = SparkSession.builder()
       .appName("Word2Vec")
       .master("local[2]")
       .getOrCreate()
     val sc = spark.sparkContext

     val trainDataPath = "E://BDS/newsparkml/src/news_tensite_xml.smarty.dat"
     // 数据预处理
     val rdd = sc.hadoopFile(trainDataPath, classOf[TextInputFormat], classOf[LongWritable], classOf[Text])
       .map(pair => new String(pair._2.getBytes, 0, pair._2.getLength, "GBK"))
       .filter(row => row.contains("content"))
       .map(row =>
         if(row.contains("content")){
           row.substring(row.indexOf(">")+1, row.lastIndexOf("<")).trim()
         }else{
           null
         }
       )
       .filter(row => !row.equals(null))
     // 分词
     val segmentResult = rdd.mapPartitions( row =>{
       row.map(word => {
         val nlpList = NLPTokenizer.segment(word)
         import scala.collection.JavaConverters._
         nlpList.asScala.map(term => {
           term.word.trim()
         })
         .filter(word => word.length>1) //过滤掉长度小于2的词
         .mkString(" ")
       })
     })
     val regex = """^\d+$""".r
     //val size = 5
     segmentResult.saveAsTextFile("E:/BDS/newsparkml/src/分词结果")
     // 加载分词训练数据
     val input = sc.textFile("E:/BDS/newsparkml/src/分词结果")
       //.filter(row => row.split(" ").length>=size)
       .filter(row => regex.findFirstMatchIn(row) == None) //过滤掉无用的数字关键词
       .map(row => {
         val split = row.split(" ")
         val array : Array[String] = new Array[String](split.length)
         for(i<- 0 until split.length){
           array(i) = split(i)
         }
         new Tuple1(array)
       })

     val dataFrame = spark.sqlContext.createDataFrame(input).toDF("text")
     dataFrame.foreach(println(_))
     //创建Word2Vec对象
     val word2Vec = new Word2Vec()
       .setInputCol("text")
       .setOutputCol("result")
       .setVectorSize(50)
       .setNumPartitions(64)
     //训练模型
     val model = word2Vec.fit(dataFrame)
     //缓存模型
     model.save("E:/BDS/newsparkml/src/Word2VecModel")
     //保存词向量数据
     /*val vector = model.getVectors.map{
       case (word, vector) => Seq(word, vector)
     }
     vector.toJavaRDD.saveAsTextFile("E:/BDS/newsparkml/src/Word2VecData")*/
     //预测
     val like = model.findSynonyms("中国", 10)
     like.foreach(println(_))
     /*for((item, literacy) <- like){
       print(s"$item $literacy")
     }*/
   }
 }

分词结果：

分词结果部分数据：

模型：

结果：

分析：

　　预测结果与训练集数据紧密相关，Word2Vec会根据训练集中各词之间的紧密程度设置不同的相识度，因此，要想获得较好的预测结果，需要有合适的训练集！