spark-mllib-TFIDF实现

TF就是词在一篇文章中的词频,IDF就是逆词频,IFIDF就是两者乘积,常用来表示词在文章中重要性,公式表示为:

官网上给出使用IF-IDF的例子代码:
object TfIdfTest {
  def main(args:Array[String]){
    val conf = new SparkConf().setAppName("TfIdfTest")
    val sc = new SparkContext(conf)

    // Load documents (one per line).
    val documents: RDD[Seq[String]] = sc.textFile("...").map(_.split(" ").toSeq)

    val hashingTF = new HashingTF()
    val tf: RDD[Vector] = hashingTF.transform(documents)

    tf.cache()
    val idf = new IDF().fit(tf)
    val tfidf: RDD[Vector] = idf.transform(tf)
  }
}

可以看到重要的就是三个类, HashingTF,IDF,IDFModel,其中val idf 的类型就是 IDFModel。
首先明确的是,它要求源数据为一篇文章一行
先看 val tf: RDD[Vector] = hashingTF.transform(documents),将调用 HashingTF类的如下方法

/**
   * Transforms the input document to term frequency vectors.
   */
  def transform[D <: Iterable[_]](dataset: RDD[D]): RDD[Vector] = {
    dataset.map(this.transform)
  }

该方法对每一篇文章调用:
/**
   * Transforms the input document into a sparse term frequency vector.
   */
  def transform(document: Iterable[_]): Vector = {
    val termFrequencies = mutable.HashMap.empty[Int, Double]
    document.foreach { term =>
      val i = indexOf(term)
      termFrequencies.put(i, termFrequencies.getOrElse(i, 0.0) + 1.0)
    }
    Vectors.sparse(numFeatures, termFrequencies.toSeq)
  }

该方法会将词特征映射到一个很大维度的向量中去,每篇文章一个向量,向量长度为 numFeatures,这是 HashingTF类的成员变量,默认为2的20次方。映射时先使用indexOf(term)获取词项term对应的向量位置,该方法如下:

/**
   * Returns the index of the input term.
   */
  def indexOf(term: Any): Int = Utils.nonNegativeMod(term.##, numFeatures)

term.##是得到一个int型的hash值,然后对numFeatures取模,得到位置后,在对应位置上进行加1操作,即:termFrequencies.put(i, termFrequencies.getOrElse(i, 0.0) + 1.0),最后把转化为一个稀疏的向量。
这样就得到了所有文章的词频向量,一篇文章一个向量。

val idf = new IDF().fit(tf)将调用IDF类的fit方法并返回一个 IDFModel 对象。源码如下:

/**
   * Computes the inverse document frequency.
   * @param dataset an RDD of term frequency vectors
   */
  def fit(dataset: RDD[Vector]): IDFModel = {
    val idf = dataset.treeAggregate(new IDF.DocumentFrequencyAggregator)(
      seqOp = (df, v) => df.add(v),
      combOp = (df1, df2) => df1.merge(df2)
    ).idf()
    new IDFModel(idf)
  }

该方法将完成IDF的所有计算工作。dataset.treeAggregate(zeroValue: U)(seqOp, combOp)的作用是使用方法 seqOp加入新的需计算聚合对象和 combOp聚合合并各自的已聚合结果(因为 treeAggregate是使用mapPartition进行计算的,最后需要合并结果 。该调用返回一个类型为U的对象,在这里就是 IDF.DocumentFrequencyAggregator的对象了,该返回对象再调用方法 idf()完成 IDF 的除法计算,最后通过new IDFModel(idf) 得到一个IDFModel对象。来看看  IDF.DocumentFrequencyAggregator 的源码:

/** Document frequency aggregator. */
  class DocumentFrequencyAggregator extends Serializable {
    /** number of documents */
    private var m = 0L
    /** document frequency vector */
    private var df: BDV[Long] = _
    /** Adds a new document. */
    def add(doc: Vector): this.type = {
      if (isEmpty) {
        df = BDV.zeros(doc.size)
      }
      doc match {
        case sv: SparseVector =>
          val nnz = sv.indices.size
          var k = 0
          while (k < nnz) {
            if (sv.values(k) > 0) {
              df(sv.indices(k)) += 1L
            }
            k += 1
          }
        case dv: DenseVector =>
          val n = dv.size
          var j = 0
          while (j < n) {
            if (dv.values(j) > 0.0) {
              df(j) += 1L
            }
            j += 1
          }
        case other =>
          throw new UnsupportedOperationException(
            s"Only sparse and dense vectors are supported but got ${other.getClass}.")
      }
      m += 1L
      this
    }
    /** Merges another. */
    def merge(other: DocumentFrequencyAggregator): this.type = {
      if (!other.isEmpty) {
        m += other.m
        if (df == null) {
          df = other.df.copy
        } else {
          df += other.df
        }
      }
      this
    }
    private def isEmpty: Boolean = m == 0L
    /** Returns the current IDF vector. */
    def idf(): Vector = {
      if (isEmpty) {
        throw new IllegalStateException("Haven't seen any document yet.")
      }
      val n = df.length
      val inv = new Array[Double](n)
      var j = 0
      while (j < n) {
        inv(j) = math.log((m + 1.0)/ (df(j) + 1.0))
        j += 1
      }
      Vectors.dense(inv)
    }
  }

一个DocumentFrequencyAggregator对象代表了一个partition里面的聚合结果,成员 m 表示该partition里有多少篇文章,df: BDV[Long] 表示词出现的文档频率向量,向量的每个元素表示这个词在多少篇文章中出现过。
看add方法,对每一个加入进来的文档Vector,会循环取出Vector中的term,如果该term对应的位置的值是大于0的,即表示该词在这篇文档中出现了,那么 df(sv.indices(k)) += 1L。
看merge方法,这是两个partition进行聚合的时候使用,即把对应的成员值相加即可。
看idf方法,也只是简单的进行最终的log计算,最后返回结果向量。

得到IDFModel对象后,进行最后的val tfidf: RDD[Vector] = idf.transform(tf)操作,得到最终结果:

/**
   * Transforms term frequency (TF) vectors to TF-IDF vectors.
   * @param dataset an RDD of term frequency vectors
   * @return an RDD of TF-IDF vectors
   */
  def transform(dataset: RDD[Vector]): RDD[Vector] = {
    val bcIdf = dataset.context.broadcast(idf)
    dataset.mapPartitions { iter =>
      val thisIdf = bcIdf.value
      iter.map { v =>
        val n = v.size
        v match {
          case sv: SparseVector =>
            val nnz = sv.indices.size
            val newValues = new Array[Double](nnz)
            var k = 0
            while (k < nnz) {
              newValues(k) = sv.values(k) * thisIdf(sv.indices(k))
              k += 1
            }
            Vectors.sparse(n, sv.indices, newValues)
          case dv: DenseVector =>
            val newValues = new Array[Double](n)
            var j = 0
            while (j < n) {
              newValues(j) = dv.values(j) * thisIdf(j)
              j += 1
            }
            Vectors.dense(newValues)
          case other =>
            throw new UnsupportedOperationException(
              s"Only sparse and dense vectors are supported but got ${other.getClass}.")
        }
      }
    }
  }

先将fit方法中得到的逆文档词频向量idf进行广播,然后进行简单的相乘操作而已。

顺便提一下mllib中的Vector,分为DenseVector和SparseVector, DenseVector就是扎扎实实的顺序存储向量中的每一个元素,SparseVector略有不同,成员size是真正的向量长度,indices是记录所有值不为0的元素的索引,values是记录所有不为0的元素的值,比如一个向量为(0,1,0,0,2,0,0),那么indices则为数组(1,4),values则为数组(1,2),这两个数组必须长度相同。

你可能感兴趣的:(spark-mllib源码走读)