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")

     }*/

   }

 }