<commit_before>/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package com.github.cloudml.zen.ml.clustering

import java.lang.ref.SoftReference
import java.util.Random

import breeze.linalg.{DenseVector => BDV, SparseVector => BSV, sum => brzSum, Vector => BV}
import com.github.cloudml.zen.ml.DBHPartitioner
import com.github.cloudml.zen.ml.clustering.LDA._
import com.github.cloudml.zen.ml.clustering.LDAUtils._
import com.github.cloudml.zen.ml.util.{XORShiftRandom, AliasTable}
import com.github.cloudml.zen.ml.util.SparkUtils._
import org.apache.spark.broadcast.Broadcast
import org.apache.spark.graphx._
import org.apache.spark.graphx.impl.GraphImpl
import org.apache.spark.mllib.linalg.distributed.{MatrixEntry, RowMatrix}
import org.apache.spark.mllib.linalg.{DenseVector => SDV, SparseVector => SSV, Vector => SV}
import org.apache.spark.mllib.regression.LabeledPoint
import org.apache.spark.mllib.util.MLUtils
import org.apache.spark.rdd.RDD
import org.apache.spark.serializer.KryoRegistrator
import org.apache.spark.storage.StorageLevel
import org.apache.spark.util.collection.AppendOnlyMap
import org.apache.spark.{HashPartitioner, Logging}

abstract class LDA private[ml](
  @transient private var corpus: Graph[VD, ED],
  private val numTopics: Int,
  private val numTerms: Int,
  private var alpha: Float,
  private var beta: Float,
  private var alphaAS: Float,
  private var storageLevel: StorageLevel,
  private var useDBHStrategy: Boolean) extends Serializable with Logging {

  /**
   * Doc number in corpus
   */
  val numDocs = docVertices.count()

  /**
   * Token number in corpus
   */
  val numTokens = corpus.edges.map(e => e.attr.size.toDouble).sum().toLong

  def setAlpha(alpha: Float): this.type = {
    this.alpha = alpha
    this
  }

  def setBeta(beta: Float): this.type = {
    this.beta = beta
    this
  }

  def setAlphaAS(alphaAS: Float): this.type = {
    this.alphaAS = alphaAS
    this
  }

  def setStorageLevel(newStorageLevel: StorageLevel): this.type = {
    this.storageLevel = newStorageLevel
    this
  }

  def setSeed(newSeed: Int): this.type = {
    this.seed = newSeed
    this
  }

  def getCorpus: Graph[VD, ED] = corpus

  @transient private var seed = new XORShiftRandom().nextInt()
  @transient private var innerIter = 1
  @transient private var totalTopicCounter: BDV[Count] = collectTotalTopicCounter(corpus)

  private def termVertices = corpus.vertices.filter(t => t._1 >= 0)

  private def docVertices = corpus.vertices.filter(t => t._1 < 0)

  private def checkpoint(corpus: Graph[VD, ED]): Unit = {
    val sc = corpus.edges.sparkContext
    if (innerIter % 12 == 0 && sc.getCheckpointDir.isDefined) {
      corpus.checkpoint()
    }
  }

  private def collectTotalTopicCounter(graph: Graph[VD, ED]): BDV[Count] = {
    val globalTopicCounter = collectGlobalCounter(graph, numTopics)
    // avoid the numTokens is larger than Int.MaxValue
    var count = 0L
    globalTopicCounter.activeValuesIterator.foreach(v => count += v.toLong)
    assert(count == numTokens)
    globalTopicCounter
  }

  private def gibbsSampling(iter: Int): Unit = {
    val previousCorpus = corpus
    val sampledCorpus = sampleTokens(corpus, totalTopicCounter, innerIter + seed,
      numTokens, numTopics, numTerms, alpha, alphaAS, beta)
    sampledCorpus.persist(storageLevel)
    sampledCorpus.edges.setName(s"sampledEdges-$iter")
    sampledCorpus.vertices.setName(s"sampledVertices-$iter")

    corpus = updateCounter(sampledCorpus, numTopics)
    checkpoint(corpus)
    corpus.persist(storageLevel)
    corpus.edges.setName(s"edges-$iter").count()
    corpus.vertices.setName(s"vertices-$iter")
    // completeCorpus.vertices.count()
    totalTopicCounter = collectTotalTopicCounter(corpus)

    previousCorpus.edges.unpersist(blocking = false)
    previousCorpus.vertices.unpersist(blocking = false)
    sampledCorpus.edges.unpersist(blocking = false)
    sampledCorpus.vertices.unpersist(blocking = false)
    innerIter += 1
  }

  private def collectGlobalCounter(graph: Graph[VD, ED], numTopics: Int): BDV[Count] = {
    graph.vertices.filter(t => t._1 >= 0).map(_._2).aggregate(BDV.zeros[Count](numTopics))(_ :+= _, _ :+= _)
  }

  protected def sampleTokens(
    graph: Graph[VD, ED],
    totalTopicCounter: BDV[Count],
    innerIter: Long,
    numTokens: Long,
    numTopics: Int,
    numTerms: Int,
    alpha: Float,
    alphaAS: Float,
    beta: Float): Graph[VD, ED]

  /**
   * Save the term-topic related model
   * @param totalIter
   */
  def saveModel(totalIter: Int = 1): DistributedLDAModel = {
    var termTopicCounter: RDD[(VertexId, VD)] = null
    for (iter <- 1 to totalIter) {
      logInfo(s"Save TopicModel (Iteration $iter/$totalIter)")
      var previousTermTopicCounter = termTopicCounter
      gibbsSampling(iter)
      val newTermTopicCounter = termVertices
      termTopicCounter = Option(termTopicCounter).map(_.join(newTermTopicCounter).map {
        case (term, (a, b)) =>
          (term, a :+ b)
      }).getOrElse(newTermTopicCounter)

      termTopicCounter.persist(storageLevel).count()
      Option(previousTermTopicCounter).foreach(_.unpersist(blocking = false))
      previousTermTopicCounter = termTopicCounter
    }
    val rand = new Random()
    val ttc = termTopicCounter.mapValues(c => {
      val nc = new BSV[Count](c.index.slice(0, c.used), c.data.slice(0, c.used).map(v => {
        val mid = v.toDouble / totalIter
        val l = math.floor(mid)
        if (rand.nextDouble() > mid - l) {
          l
        } else {
          l + 1
        }
      }.asInstanceOf[Count]), c.length)
      nc
    })
    ttc.persist(storageLevel)
    val gtc = ttc.map(_._2).aggregate(BDV.zeros[Count](numTopics))(_ :+= _, _ :+= _)
    new DistributedLDAModel(gtc, ttc, numTopics, numTerms, alpha, beta, alphaAS)
  }

  def runGibbsSampling(iterations: Int): Unit = {
    for (iter <- 1 to iterations) {
      logInfo(s"Start Gibbs sampling (Iteration $iter/$iterations)")
      val startedAt = System.nanoTime()
      gibbsSampling(iter)
      val elapsedSeconds = (System.nanoTime() - startedAt) / 1e9
      // logInfo(s"Gibbs sampling (Iteration $iter/$iterations):  ${perplexity()}")
      logInfo(s"End Gibbs sampling  (Iteration $iter/$iterations) takes:  $elapsedSeconds")
    }
  }<commit_msg>still fix<commit_after>/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package com.github.cloudml.zen.ml.clustering

import java.lang.ref.SoftReference
import java.util.Random

import breeze.linalg.{DenseVector => BDV, SparseVector => BSV, sum => brzSum, Vector => BV}
import com.github.cloudml.zen.ml.DBHPartitioner
import com.github.cloudml.zen.ml.clustering.LDA._
import com.github.cloudml.zen.ml.clustering.LDAUtils._
import com.github.cloudml.zen.ml.util.{XORShiftRandom, AliasTable}
import com.github.cloudml.zen.ml.util.SparkUtils._
import org.apache.spark.broadcast.Broadcast
import org.apache.spark.graphx._
import org.apache.spark.graphx.impl.GraphImpl
import org.apache.spark.mllib.linalg.distributed.{MatrixEntry, RowMatrix}
import org.apache.spark.mllib.linalg.{DenseVector => SDV, SparseVector => SSV, Vector => SV}
import org.apache.spark.mllib.regression.LabeledPoint
import org.apache.spark.mllib.util.MLUtils
import org.apache.spark.rdd.RDD
import org.apache.spark.serializer.KryoRegistrator
import org.apache.spark.storage.StorageLevel
import org.apache.spark.util.collection.AppendOnlyMap
import org.apache.spark.{HashPartitioner, Logging}

abstract class LDA private[ml](
  @transient private var corpus: Graph[VD, ED],
  private val numTopics: Int,
  private val numTerms: Int,
  private var alpha: Float,
  private var beta: Float,
  private var alphaAS: Float,
  private var storageLevel: StorageLevel,
  private var useDBHStrategy: Boolean) extends Serializable with Logging {

  /**
   * Doc number in corpus
   */
  val numDocs = docVertices.count()

  /**
   * Token number in corpus
   */
  val numTokens = corpus.edges.map(e => e.attr.size.toDouble).sum().toLong

  def setAlpha(alpha: Float): this.type = {
    this.alpha = alpha
    this
  }

  def setBeta(beta: Float): this.type = {
    this.beta = beta
    this
  }

  def setAlphaAS(alphaAS: Float): this.type = {
    this.alphaAS = alphaAS
    this
  }

  def setStorageLevel(newStorageLevel: StorageLevel): this.type = {
    this.storageLevel = newStorageLevel
    this
  }

  def setSeed(newSeed: Int): this.type = {
    this.seed = newSeed
    this
  }

  def getCorpus: Graph[VD, ED] = corpus

  @transient private var seed = new XORShiftRandom().nextInt()
  @transient private var innerIter = 1
  @transient private var totalTopicCounter: BDV[Count] = collectTotalTopicCounter(corpus)

  private def termVertices = corpus.vertices.filter(t => t._1 >= 0)

  private def docVertices = corpus.vertices.filter(t => t._1 < 0)

  private def checkpoint(corpus: Graph[VD, ED]): Unit = {
    val sc = corpus.edges.sparkContext
    if (innerIter % 12 == 0 && sc.getCheckpointDir.isDefined) {
      corpus.checkpoint()
    }
  }

  private def collectTotalTopicCounter(graph: Graph[VD, ED]): BDV[Count] = {
    val globalTopicCounter = collectGlobalCounter(graph, numTopics)
    // avoid the numTokens is larger than Int.MaxValue
    var count = 0L
    globalTopicCounter.activeValuesIterator.foreach(v => count += v.toLong)
    assert(count == numTokens)
    globalTopicCounter
  }

  private def gibbsSampling(iter: Int): Unit = {
    val previousCorpus = corpus
    val sampledCorpus = sampleTokens(corpus, totalTopicCounter, innerIter + seed,
      numTokens, numTopics, numTerms, alpha, alphaAS, beta)
    sampledCorpus.persist(storageLevel)
    sampledCorpus.edges.setName(s"sampledEdges-$iter")
    sampledCorpus.vertices.setName(s"sampledVertices-$iter")

    corpus = updateCounter(sampledCorpus, numTopics)
    checkpoint(corpus)
    corpus.persist(storageLevel)
    corpus.edges.setName(s"edges-$iter").count()
    corpus.vertices.setName(s"vertices-$iter")
    // completeCorpus.vertices.count()
    totalTopicCounter = collectTotalTopicCounter(corpus)

    previousCorpus.edges.unpersist(blocking = false)
    previousCorpus.vertices.unpersist(blocking = false)
    sampledCorpus.edges.unpersist(blocking = false)
    sampledCorpus.vertices.unpersist(blocking = false)
    innerIter += 1
  }

  private def collectGlobalCounter(graph: Graph[VD, ED], numTopics: Int): BDV[Count] = {
    graph.vertices.filter(t => t._1 >= 0).map(_._2).aggregate(BDV.zeros[Count](numTopics))(_ :+= _, _ :+= _)
  }

  protected def sampleTokens(
    graph: Graph[VD, ED],
    totalTopicCounter: BDV[Count],
    innerIter: Long,
    numTokens: Long,
    numTopics: Int,
    numTerms: Int,
    alpha: Float,
    alphaAS: Float,
    beta: Float): Graph[VD, ED]

  /**
   * Save the term-topic related model
   * @param totalIter
   */
  def saveModel(totalIter: Int = 1): DistributedLDAModel = {
    var termTopicCounter: RDD[(VertexId, VD)] = null
    for (iter <- 1 to totalIter) {
      logInfo(s"Save TopicModel (Iteration $iter/$totalIter)")
      var previousTermTopicCounter = termTopicCounter
      gibbsSampling(iter)
      val newTermTopicCounter = termVertices
      termTopicCounter = Option(termTopicCounter).map(_.join(newTermTopicCounter).map {
        case (term, (a, b)) =>
          (term, a :+ b)
      }).getOrElse(newTermTopicCounter)

      termTopicCounter.persist(storageLevel).count()
      Option(previousTermTopicCounter).foreach(_.unpersist(blocking = false))
      previousTermTopicCounter = termTopicCounter
    }
    val rand = new Random()
    val ttc = termTopicCounter.mapValues(c => {
      val nc = new BSV[Count](c.index.slice(0, c.used), c.data.slice(0, c.used).map(v => {
        val mid = v.toDouble / totalIter
        val l = math.floor(mid)
        if (rand.nextDouble() > mid - l) {
          l
        } else {
          l + 1
        }
      }.toInt), c.length)
      nc
    })
    ttc.persist(storageLevel)
    val gtc = ttc.map(_._2).aggregate(BDV.zeros[Count](numTopics))(_ :+= _, _ :+= _)
    new DistributedLDAModel(gtc, ttc, numTopics, numTerms, alpha, beta, alphaAS)
  }

  def runGibbsSampling(iterations: Int): Unit = {
    for (iter <- 1 to iterations) {
      logInfo(s"Start Gibbs sampling (Iteration $iter/$iterations)")
      val startedAt = System.nanoTime()
      gibbsSampling(iter)
      val elapsedSeconds = (System.nanoTime() - startedAt) / 1e9
      // logInfo(s"Gibbs sampling (Iteration $iter/$iterations):  ${perplexity()}")
      logInfo(s"End Gibbs sampling  (Iteration $iter/$iterations) takes:  $elapsedSeconds")
    }
  }
{'fraction_non_alphanumeric': 0.11321291423695065, 'fraction_numerical': 0.0036979092589958756, 'mean_word_length': 3.6535407015221706, 'pattern_counts': {'":': 0, '<': 9, '<?xml version=': 0, '>': 39, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4823276', 'n_tokens_mistral': 4632, 'n_tokens_neox': 4459, 'n_words': 1085}