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会根据训练集中各词之间的紧密程度设置不同的相识度，因此，要想获得较好的预测结果，需要有合适的训练集！