Spark Shell中Scala实现Naive Bayes推特位置预测任务咨询

嘿，看来你已经搞定了TSV文件读取和RDD转换这一步，很棒！接下来咱们一步步完成基于Naive Bayes的推特位置预测任务，我会把每一步的思路和可运行的Scala代码都给你理清楚～

1. 数据预处理：把原始数据转成模型能认的格式 #

Naive Bayes需要的是LabeledPoint类型的数据（数字标签+特征向量），所以咱们得先完成两步转换：把位置字符串转成数字标签，把文本转成TF-IDF特征向量。

1.1 位置标签化 #

首先把字符串类型的位置映射成数字（模型只认数值型标签）：

// 假设t是你已经读取好的原始RDD，每行按制表符分割
val tweetsRDD = t.map(line => {
  val parts = line.split("\t")
  (parts(0), parts(1)) // 这里假设第0列是位置，第1列是推特文本，结构不对的话自行调整
})

// 收集所有唯一位置，建立「位置->数字标签」的映射
val locationToIndex = tweetsRDD.map(_._1).distinct().collect().zipWithIndex.toMap
// 把原始数据转成（数字标签，文本）的格式
val indexedTweets = tweetsRDD.map{ case (loc, text) => (locationToIndex(loc).toDouble, text) }

1.2 文本转TF-IDF特征向量 #

文本分类里TF-IDF是最常用的特征提取方式，咱们用Spark ML库的工具来实现：

import org.apache.spark.ml.feature.{HashingTF, IDF, Tokenizer, StopWordsRemover}
import org.apache.spark.mllib.linalg.Vector
import org.apache.spark.mllib.regression.LabeledPoint

// 先把RDD转成DataFrame，方便用ML库的工具链
val df = indexedTweets.toDF("label", "text")

// 第一步：分词，把整段文本拆成单个单词
val tokenizer = new Tokenizer().setInputCol("text").setOutputCol("words")
val wordsData = tokenizer.transform(df)

// 可选但推荐：去掉停用词（比如the、a这类无意义的词），提升模型效果
val stopWordsRemover = new StopWordsRemover().setInputCol("words").setOutputCol("filteredWords")
val filteredWordsData = stopWordsRemover.transform(wordsData)

// 第二步：计算词频（TF）
val hashingTF = new HashingTF()
  .setInputCol("filteredWords") // 如果没做停用词过滤，这里改成"words"
  .setOutputCol("rawFeatures")
  .setNumFeatures(10000) // 哈希桶数量，可根据文本复杂度调整，比如5000/20000
val tfData = hashingTF.transform(filteredWordsData)

// 第三步：计算逆文档频率（IDF），得到最终的TF-IDF特征向量
val idf = new IDF().setInputCol("rawFeatures").setOutputCol("features")
val idfModel = idf.fit(tfData)
val tfIdfData = idfModel.transform(tfData)

// 最后转成RDD[LabeledPoint]，供Mllib的Naive Bayes使用
val trainingData = tfIdfData.select("label", "features").rdd.map{ row =>
  LabeledPoint(row.getDouble(0), row.getAs[Vector]("features"))
}

2. 训练Naive Bayes模型 #

Spark Mllib里有现成的Multinomial Naive Bayes实现（刚好适合文本分类场景），直接训练就行：

import org.apache.spark.mllib.classification.NaiveBayes
import org.apache.spark.mllib.classification.NaiveBayesModel

// 训练模型，lambda是平滑参数，防止出现零概率问题，默认1.0就够用
val model = NaiveBayes.train(trainingData, lambda = 1.0)

3. 处理输入文本并输出预测结果 #

现在就可以处理用户输入的文本，比如"please help me out stackoverflow"，先把它转换成和训练数据一样的特征向量，然后预测概率，最后把数字标签转回位置字符串：

// 定义一个工具函数，把输入文本转换成TF-IDF特征向量
def convertTextToFeatures(text: String): Vector = {
  // 把输入文本包装成DataFrame，复用之前的分词/TF-IDF流程
  val inputDF = spark.createDataFrame(Seq((0.0, text))).toDF("label", "text")
  val words = tokenizer.transform(inputDF)
  val filteredWords = stopWordsRemover.transform(words)
  val tf = hashingTF.transform(filteredWords)
  idfModel.transform(tf).select("features").first().getAs[Vector]("features")
}

// 测试输入文本
val inputText = "please help me out stackoverflow"
val inputFeatures = convertTextToFeatures(inputText)

// 获取所有位置的预测概率
val probArray = model.predictProbabilities(inputFeatures).toArray
// 找到概率最高的位置的标签索引
val (maxProb, maxIndex) = probArray.zipWithIndex.maxBy(_._1)

// 把数字标签转回位置字符串（需要反转之前的映射）
val indexToLocation = locationToIndex.map(_.swap)
val predictedLocation = indexToLocation(maxIndex)

// 输出结果，保留两位小数
println(s"$predictedLocation, ${maxProb.round(2)}")

一些小提示 #

• 如果你的TSV结构不是「位置+文本」，一定要调整split("\t")后的字段索引，比如文本在第3列就改成parts(2)。
• setNumFeatures的数值可以调整：太大容易过拟合，太小可能丢失关键特征，10000是比较稳妥的初始值。
• 300万条数据量不小，建议先做数据清洗：比如去掉空文本、特殊符号，或者过滤掉出现次数极少的位置（比如只有几条推特的位置，模型很难学好）。

内容的提问来源于stack exchange，提问作者OneQuestionAwayFromNirvana