<commit_before>package BIDMach.networks

import BIDMat.{Mat,SBMat,CMat,DMat,FMat,IMat,LMat,HMat,GMat,GDMat,GIMat,GLMat,GSMat,GSDMat,SMat,SDMat}
import BIDMat.MatFunctions._
import BIDMat.SciFunctions._
import BIDMach.datasources._
import BIDMach.updaters._
import BIDMach.mixins._
import BIDMach.models._
import BIDMach._
import edu.berkeley.bid.CUMACH
import edu.berkeley.bid.CPUMACH
import jcuda.runtime.JCuda._
import scala.util.hashing.MurmurHash3
import scala.collection.mutable.ArrayBuffer
import java.text.SimpleDateFormat
import java.util.Calendar

/**
 * 
 */

class Word2Vec(override val opts:Word2Vec.Opts = new Word2Vec.Options) extends Model(opts) {
  
  var firstPos = -1L;
  var wordtab:Mat = null;
  var randpermute:Mat = null;
  var ubound:Mat = null;
  var minusone:Mat = null;
  var wordmask:Mat = null;
  var allones:Mat = null;
  var randwords:Mat = null;
  var randsamp:Mat = null;
  var retEvalPos:GMat = null;
  var retEvalNeg:GMat = null;
  var nfeats = 0;
  var ncols = 0;
  var expt = 0f;
  var vexp = 0f;
  var salpha = 0f;
  
  var ntimes = 12;
  var times:FMat = null;
  var delays:FMat = null;
  var log:ArrayBuffer[String] = null
  val dateFormat = new SimpleDateFormat("hh:mm:ss:SSS")

  
  def addTime(itime:Int, lasti:Int = -1) = {
    val t = toc
    times(itime) = t;
    if (itime > 0) {
    	delays(itime) += times(itime) - times(itime + lasti);
    } 
    val today = Calendar.getInstance().getTime()
    log += "Log: %s, GPU %d, event %d" format (dateFormat.format(today), if (useGPU) getGPU else 0, itime);
  }
  
  var test1:Mat = null;
  var test2:Mat = null;
  var test3:Mat = null;
  var test4:Mat = null;
  

  override def init() = {
    val mats = datasource.next;
	  nfeats = opts.vocabSize;
	  ncols = mats(0).ncols;
	  datasource.reset;
    if (refresh) {
    	setmodelmats(new Array[Mat](2));
    	modelmats(0) = (rand(opts.dim, nfeats) - 0.5f)/opts.dim;              // syn0 - context model
    	modelmats(1) = zeros(opts.dim, nfeats);                               // syn1neg - target word model
    }
    modelmats(0) = convertMat(modelmats(0));
    modelmats(1) = convertMat(modelmats(1));
    val nskip = opts.nskip;
    val nwindow = nskip * 2 + 1;
    val skipcol = icol((-nskip) to -1) on icol(1 to nskip)
    expt = 1f / (1f - opts.wexpt);
    wordtab = convertMat(max(0, min(ncols+1, iones(nwindow-1, 1) * irow(1 -> (ncols+1)) + skipcol)));  // Indices for convolution matrix
    wordmask = convertMat(skipcol * iones(1, ncols));                                   // columns = distances from center word
    randpermute = convertMat(zeros(nwindow-1, ncols));                                    // holds random values for permuting negative context words
    ubound = convertMat(zeros(1, ncols));                                               // upper bound random matrix
    minusone = convertMat(irow(-1));
    allones = convertMat(iones(1, ncols));
    randwords = convertMat(zeros(1, (1.01 * opts.nneg * nskip * ncols / opts.nreuse).toInt)); // generates random negative words
    randsamp = convertMat(zeros(1, ncols));                                             // For sub-sampling frequent words
    val gopts = opts.asInstanceOf[ADAGrad.Opts];
    vexp = gopts.vexp.v;
    salpha = opts.wsample * math.log(nfeats).toFloat;
    if (useGPU) {
      retEvalPos = GMat(1,1);
      retEvalNeg = GMat(1,1);
    }
    times = zeros(1, ntimes);
    delays = zeros(1, ntimes);
    log = ArrayBuffer();
  }
  
  def dobatch(gmats:Array[Mat], ipass:Int, pos:Long):Unit = {
    addTime(0);
    if (gmats(0).ncols == ncols) {
    	if (firstPos < 0) firstPos = pos;
    	val nsteps = 1f * pos / firstPos;
    	val gopts = opts.asInstanceOf[ADAGrad.Opts];
    	val lrate = gopts.lrate.dv.toFloat * math.pow(nsteps, - gopts.texp.dv).toFloat;
    	val (words, lb, ub, trandwords, goodwords) = wordMats(gmats, ipass, pos);

    	val lrpos = lrate.dv.toFloat;
    	val lrneg = lrpos/opts.nneg;  
//    	val lrneg = lrpos;
    	procPositives(opts.nskip, words, lb, ub, modelmats(1), modelmats(0), lrpos, vexp);
    	addTime(8);
    	procNegatives(opts.nneg, opts.nreuse, trandwords, goodwords, modelmats(1), modelmats(0), lrneg, vexp); 
    	addTime(9);
    }
  }
  
  def evalbatch(gmats:Array[Mat], ipass:Int, pos:Long):FMat = {
  	addTime(0);
  	if (gmats(0).ncols == ncols) {
  	val (words, lb, ub, trandwords, goodwords) = wordMats(gmats, ipass, pos);
  	val epos = evalPositives(opts.nskip, words, lb, ub, modelmats(1), modelmats(0));
  	addTime(10,-3);
    val eneg = evalNegatives(opts.nneg, opts.nreuse, trandwords, goodwords, modelmats(1), modelmats(0));
    addTime(11)
//  	val score = ((epos + eneg / opts.nneg) /opts.nskip / words.ncols);
  	val score = ((epos + eneg / opts.nneg) / goodwords.length);
  	row(score)
  	} else row(0);
  }
  
  def wordMats(mats:Array[Mat], ipass:Int, pos:Long):(Mat, Mat, Mat, Mat, Mat) = {
  
    val wordsens = mats(0);
    val words = if (opts.iflip) wordsens(1,?) else wordsens(0,?);
    val wgood = words < opts.vocabSize;                                        // Find OOV words 
    addTime(1);
    
    rand(randsamp);                                                            // Take a random sample
    val wrat = float(words+1) * salpha;
    wrat ~ sqrt(wrat) + wrat;
    wgood ~ wgood ∘ int(randsamp < wrat);
    words ~ (wgood ∘ (words + 1)) - 1;                                        // Set OOV or skipped samples to -1
    addTime(2);
        
    rand(ubound);                                                              // get random upper and lower bounds   
    val ubrand = min(opts.nskip, int(ubound * opts.nskip) + 1);
//    println(ubrand.toString)
    val lbrand = - ubrand;
    addTime(3);
    
    val sentencenum = if (opts.iflip) wordsens(0,?) else wordsens(1,?);        // Get the nearest sentence boundaries
    val lbsentence = - cumsumByKey(allones, sentencenum) + 1;
    val ubsentence = reverse(cumsumByKey(allones, reverse(sentencenum))) - 1;
    val lb = max(lbrand, lbsentence);                                          // Combine the bounds
    val ub = min(ubrand, ubsentence);
    test3 = lb
    test4 = ub
    addTime(4);
    
    val (trandwords, contextwords) = (words, lb, ub) match {
      case (giwords:GIMat, gilb:GIMat, giub:GIMat) => {

      	val iwords = minusone \ words \ minusone;                                  // Build a convolution matrix.
      	val cwords = iwords(wordtab);
      	val pgoodwords = (wordmask >= lb) ∘ (wordmask <= ub) ∘ (cwords >= 0) ∘ (words >= 0);  // Find context words satisfying the bound
      	                                                                           // and check that context and center word are good.
      	val fgoodwords = float(pgoodwords);
      	addTime(5);
      	
      	test1 = cwords;

      	rand(randpermute);                                                         // Prepare a random permutation of context words for negative sampling
      	randpermute ~ (fgoodwords ∘ (randpermute + 1f)) - 1f;                      // set the values for bad words to -1.
      	val (vv, ii) = sortdown2(randpermute.view(randpermute.length, 1));         // Permute the good words
      	val ngood = sum(vv >= 0f).dv.toInt;                                        // Count of the good words
      	val ngoodcols = ngood / opts.nreuse;                                       // Number of good columns
      	val cwi = cwords(ii);
      	
      	test2 = cwi
      	addTime(6);

      	rand(randwords);                                                           // Compute some random negatives
      	val irandwords = min(nfeats-1, int(nfeats * (randwords ^ expt)));    
      	val trandwords0 = irandwords.view(opts.nneg, ngoodcols);                   // shrink the matrices to the available data
      	val contextwords0 = cwi.view(opts.nreuse, ngoodcols);
      	addTime(7);
      	(trandwords0, contextwords0)
      }
      case (iwords:IMat, ilb:IMat, iub:IMat) => {
        getnegs(iwords, ilb, iub, Mat.numThreads);
      }
    }
    
    (words, lb, ub, trandwords, contextwords);
  }
  
  def getnegs(words:IMat, lb:IMat, ub:IMat, nthreads:Int):(IMat, IMat) = {
    val ncols = words.ncols;
                                                                               // First count the good context words
    val cwcounts = irow((0 until nthreads).par.map((ithread:Int) => {          // work on blocks
      val istart = ((1L * ncols * ithread)/nthreads).toInt;
      val iend = ((1L * ncols * (ithread + 1))/nthreads).toInt;
      var i = istart;
      var icount = 0;
      while (i < iend) {                                                       // iterate over center words
        if (words.data(i) >= 0) {                                              // check center word is good
        	var j = lb.data(i);                                                  // get lower and upper bounds
        	var jend = ub.data(i);
        	while (j <= jend) {
        		if (j != 0 && words.data(i + j) >= 0) {                            // if not center word and context word is good, count it. 
        		  icount += 1;        		  
        		}
        		j += 1;
        	}
        }
        i += 1;
      }
      icount
    }).toArray)
                                                                               // Now we know how many good words in each block
    val ccc = cumsum(cwcounts);                                                // so size the context word and neg word matrices
    val ngroups = ccc(ccc.length - 1) / opts.nreuse;
    val contextwords0 = izeros(opts.nreuse, ngroups);
    val trandwords0 = izeros(opts.nneg, ngroups);
    
    (0 until nthreads).par.map((ithread:Int) => {                              // Copy the good words into a dense matrix (contextwords0)
      val istart = ((1L * ncols * ithread)/nthreads).toInt;
      val iend = ((1L * ncols * (ithread + 1))/nthreads).toInt;
      var i = istart;
      var icount = 0;
      val mptr = ccc(ithread) - ccc(0);
      while (i < iend) {
        if (words.data(i) >= 0) {
        	var j = lb.data(i);
        	var jend = ub.data(i);
        	while (j <= jend && mptr + icount < contextwords0.length) {
        		if (j != 0 && words.data(i + j) >= 0) {
        		  contextwords0.data(mptr + icount) = words.data(i + j)
        		  icount += 1;        		  
        		}
        		j += 1;
        	}
        }
        i += 1;
      }
      icount
    })
    
    addTime(5);
    
    val prand = drand(opts.nreuse, ngroups);                             // Rands for permutation
    
    var i = 0;                                                           // Permute the good context words randomly
    val n = prand.length;
    while (i < n) {
      val indx = math.min(n-1, i + math.floor(prand.data(i) * (n - i)).toInt);
      if (indx > i) {
        val tmp = contextwords0.data(i);
        contextwords0.data(i) = contextwords0.data(indx);
        contextwords0.data(indx) = tmp;
      }
      i += 1;
    }
    addTime(6);
    
    val randneg = rand(opts.nneg, ngroups);                                // Compute some random negatives

    (0 until nthreads).par.map((ithread:Int) => {                          // Work in blocks over the negs
      val istart = ((1L * ngroups * opts.nneg * ithread)/nthreads).toInt;
      val iend = ((1L * ngroups * opts.nneg * (ithread + 1))/nthreads).toInt;
      var i = istart;
      while (i < iend) {
      	trandwords0.data(i) = math.min(nfeats-1, (nfeats * math.pow(randneg.data(i), expt)).toInt);
        i += 1;
      }
    })
    addTime(7);
    
    (trandwords0, contextwords0)    
  }
  
  def procPositives(nskip:Int, words:Mat, lbound:Mat, ubound:Mat, model1:Mat, model2:Mat, lrate:Float, vexp:Float) = {
    val nrows = model1.nrows;
    val ncols = model1.ncols;
    val nwords = words.ncols;
    Mat.nflops += 6L * nwords * nskip * nrows;
    (words, lbound, ubound, model1, model2) match {
      case (w:GIMat, lb:GIMat, ub:GIMat, m1:GMat, m2:GMat) => {
        val err = CUMACH.word2vecPos(nrows, nwords, nskip, w.data, lb.data, ub.data, m1.data, m2.data, lrate, vexp);
        if (err != 0)  throw new RuntimeException("CUMACH.word2vecPos error " + cudaGetErrorString(err));
      }
      case (w:IMat, lb:IMat, ub:IMat, m1:FMat, m2:FMat) => if (Mat.useMKL) {
        val gw = GIMat(w);
        val glb = GIMat(lb);
        val gub = GIMat(ub);
        val gm1 = GMat(m1);
        val gm2 = GMat(m2);
        val err = CUMACH.word2vecPos(nrows, nwords, nskip, gw.data, glb.data, gub.data, gm1.data, gm2.data, lrate, vexp);
        m1 <-- gm1
        m2 <-- gm2
        if (err != 0) throw new RuntimeException("CUMACH.word2vecEvalPos error " + cudaGetErrorString(err));
//        CPUMACH.word2vecPos(nrows, nwords, nskip, w.data, lb.data, ub.data, m1.data, m2.data, lrate, vexp, Mat.numThreads);
      } else {
        Word2Vec.procPosCPU(nrows, nwords, nskip, w.data, lb.data, ub.data, m1.data, m2.data, lrate, vexp, Mat.numThreads);
      }
    }
  }
  
  def procNegatives(nwa:Int, nwb:Int, wordsa:Mat, wordsb:Mat, modela:Mat, modelb:Mat, lrate:Float, vexp:Float) = {
    val nrows = modela.nrows;
    val ncols = modela.ncols;
    val nwords = wordsa.ncols;
    Mat.nflops += 6L * nwords * nwa * nwb * nrows;
    (wordsa, wordsb, modela, modelb) match {
      case (wa:GIMat, wb:GIMat, ma:GMat, mb:GMat) => {
        val err = CUMACH.word2vecNeg(nrows, nwords, nwa, nwb, wa.data, wb.data, ma.data, mb.data, lrate, vexp);
        if (err != 0) throw new RuntimeException("CUMACH.word2vecNeg error " + cudaGetErrorString(err));
      }
      case (wa:IMat, wb:IMat, ma:FMat, mb:FMat) => if (Mat.useMKL) {
        CPUMACH.word2vecNeg(nrows, nwords, nwa, nwb, wa.data, wb.data, ma.data, mb.data, lrate, vexp, Mat.numThreads);
      } else {
      	Word2Vec.procNegCPU(nrows, nwords, nwa, nwb, wa.data, wb.data, ma.data, mb.data, lrate, vexp, Mat.numThreads);
      }
    }
  }
    
  def evalPositives(nskip:Int, words:Mat, lbound:Mat, ubound:Mat, model1:Mat, model2:Mat):Double = {
    val nrows = model1.nrows;
    val ncols = model1.ncols;
    val nwords = words.ncols;
    Mat.nflops += 2L * nwords * nskip * nrows;
    (words, lbound, ubound, model1, model2) match {
      case (w:GIMat, lb:GIMat, ub:GIMat, m1:GMat, m2:GMat) => {
        retEvalPos.clear
        val err = CUMACH.word2vecEvalPos(nrows, nwords, nskip, w.data, lb.data, ub.data, m1.data, m2.data, retEvalPos.data);
    		if (err != 0) throw new RuntimeException("CUMACH.word2vecEvalPos error " + cudaGetErrorString(err));
        retEvalPos.dv;
      }
      case (w:IMat, lb:IMat, ub:IMat, m1:FMat, m2:FMat) => 
      if (Mat.useMKL) {
      	CPUMACH.word2vecEvalPos(nrows, nwords, nskip, w.data, lb.data, ub.data, m1.data, m2.data, Mat.numThreads);
      } else {
        Word2Vec.evalPosCPU(nrows, nwords, nskip, w.data, lb.data, ub.data, m1.data, m2.data, Mat.numThreads);
      }
    }
  }
  
  def evalNegatives(nwa:Int, nwb:Int, wordsa:Mat, wordsb:Mat, modela:Mat, modelb:Mat):Double = {
    val nrows = modela.nrows;
    val ncols = modela.ncols;
    val nwords = wordsa.ncols;
    Mat.nflops += 2L * nwords * nwa * nwb * nrows;
    (wordsa, wordsb, modela, modelb) match {
      case (wa:GIMat, wb:GIMat, ma:GMat, mb:GMat) => {
        retEvalNeg.clear
        val err = CUMACH.word2vecEvalNeg(nrows, nwords, nwa, nwb, wa.data, wb.data, ma.data, mb.data, retEvalNeg.data);
    		if (err != 0) throw new RuntimeException("CUMACH.word2vecEvalNeg error " + cudaGetErrorString(err));
        retEvalNeg.dv;      
      }
      case (wa:IMat, wb:IMat, ma:FMat, mb:FMat) => 
      if (Mat.useMKL) {
      	CPUMACH.word2vecEvalNeg(nrows, nwords, nwa, nwb, wa.data, wb.data, ma.data, mb.data, Mat.numThreads); 
      } else {
      	Word2Vec.evalNegCPU(nrows, nwords, nwa, nwb, wa.data, wb.data, ma.data, mb.data, Mat.numThreads);
      }
    }
  }
  
  def trailingZeros(a:Long):Int = {
    var aa = a;
    var nz = 0;
    while ((aa & 1L) == 0) {
      aa = aa >> 1;
      nz += 1;
    }
    nz
  }
  
  override def mergeModelFn(models:Array[Model], mm:Array[Mat], um:Array[Mat], istep:Long):Unit = {
    val headlen = if (istep > 0) math.max(opts.headlen, opts.headlen << trailingZeros(istep)) else 0;
    val mlen = models(0).modelmats.length;
    val thisGPU = getGPU;
    val modj = new Array[Mat](models.length);
    for (j <- 0 until mlen) {
      val mmj = if (headlen > 0) mm(j).view(mm(j).nrows, math.min(mm(j).ncols, headlen)) else mm(j);
      mmj.clear
      for (i <- 0 until models.length) {
        if (useGPU && i < Mat.hasCUDA) setGPU(i);
        modj(i) = if (headlen > 0) models(i).modelmats(j).view(models(i).modelmats(j).nrows, math.min(models(i).modelmats(j).ncols, headlen)) else models(i).modelmats(j);
        val umj = if (headlen > 0) um(j).view(um(j).nrows, math.min(um(j).ncols, headlen)) else um(j);
        umj <-- modj(i)
        mmj ~ mmj + umj;
      }
      mmj ~ mmj * (1f/models.length);
      for (i <- 0 until models.length) {
        modj(i) <-- mmj;
      }
    }
    setGPU(thisGPU);
  }

}

object Word2Vec  {
  trait Opts extends Model.Opts {
    var aopts:ADAGrad.Opts = null;
    var nskip = 5;
    var nneg = 5;
    var nreuse = 5;  
    var vocabSize = 100000;
    var wexpt = 0.75f;
    var wsample = 1e-4f;
    var headlen = 10000;
    var iflip = false;
  }
  
  class Options extends Opts {}
  
  
  def procPosCPU(nrows:Int, ncols:Int, skip:Int, W:Array[Int], LB:Array[Int], UB:Array[Int],
      A:Array[Float], B:Array[Float], lrate:Float, vexp:Float, nthreads:Int):Int = {

    (0 until nthreads).par.map((ithread:Int) => {
    	val istart = ((1L * ithread * ncols)/nthreads).toInt;
    	val iend = ((1L * (ithread+1) * ncols)/nthreads).toInt;
    	val daa = new Array[Float](nrows);
    	var i = istart;
    	while (i < iend) {
    	  var j = 0;
    	  var k = 0;
    	  var c = 0;
    	  var cv = 0f;

    	  val iac = W(i);
    	  val ascale = math.pow(iac, vexp).toFloat;
    	  val ia = nrows * iac;                                        // Get the current word (as a model matrix offset). 
    	  if (ia >= 0) {                                               // Check for OOV words<commit_msg>added eqPosNeg option to word2vec:<commit_after>package BIDMach.networks

import BIDMat.{Mat,SBMat,CMat,DMat,FMat,IMat,LMat,HMat,GMat,GDMat,GIMat,GLMat,GSMat,GSDMat,SMat,SDMat}
import BIDMat.MatFunctions._
import BIDMat.SciFunctions._
import BIDMach.datasources._
import BIDMach.updaters._
import BIDMach.mixins._
import BIDMach.models._
import BIDMach._
import edu.berkeley.bid.CUMACH
import edu.berkeley.bid.CPUMACH
import jcuda.runtime.JCuda._
import scala.util.hashing.MurmurHash3
import scala.collection.mutable.ArrayBuffer
import java.text.SimpleDateFormat
import java.util.Calendar

/**
 * 
 */

class Word2Vec(override val opts:Word2Vec.Opts = new Word2Vec.Options) extends Model(opts) {
  
  var firstPos = -1L;
  var wordtab:Mat = null;
  var randpermute:Mat = null;
  var ubound:Mat = null;
  var minusone:Mat = null;
  var wordmask:Mat = null;
  var allones:Mat = null;
  var randwords:Mat = null;
  var randsamp:Mat = null;
  var retEvalPos:GMat = null;
  var retEvalNeg:GMat = null;
  var nfeats = 0;
  var ncols = 0;
  var expt = 0f;
  var vexp = 0f;
  var salpha = 0f;
  
  var ntimes = 12;
  var times:FMat = null;
  var delays:FMat = null;
  var log:ArrayBuffer[String] = null
  val dateFormat = new SimpleDateFormat("hh:mm:ss:SSS")

  
  def addTime(itime:Int, lasti:Int = -1) = {
    val t = toc
    times(itime) = t;
    if (itime > 0) {
    	delays(itime) += times(itime) - times(itime + lasti);
    } 
    val today = Calendar.getInstance().getTime()
    log += "Log: %s, GPU %d, event %d" format (dateFormat.format(today), if (useGPU) getGPU else 0, itime);
  }
  
  var test1:Mat = null;
  var test2:Mat = null;
  var test3:Mat = null;
  var test4:Mat = null;
  

  override def init() = {
    val mats = datasource.next;
	  nfeats = opts.vocabSize;
	  ncols = mats(0).ncols;
	  datasource.reset;
    if (refresh) {
    	setmodelmats(new Array[Mat](2));
    	modelmats(0) = (rand(opts.dim, nfeats) - 0.5f)/opts.dim;              // syn0 - context model
    	modelmats(1) = zeros(opts.dim, nfeats);                               // syn1neg - target word model
    }
    modelmats(0) = convertMat(modelmats(0));
    modelmats(1) = convertMat(modelmats(1));
    val nskip = opts.nskip;
    val nwindow = nskip * 2 + 1;
    val skipcol = icol((-nskip) to -1) on icol(1 to nskip)
    expt = 1f / (1f - opts.wexpt);
    wordtab = convertMat(max(0, min(ncols+1, iones(nwindow-1, 1) * irow(1 -> (ncols+1)) + skipcol)));  // Indices for convolution matrix
    wordmask = convertMat(skipcol * iones(1, ncols));                                   // columns = distances from center word
    randpermute = convertMat(zeros(nwindow-1, ncols));                                    // holds random values for permuting negative context words
    ubound = convertMat(zeros(1, ncols));                                               // upper bound random matrix
    minusone = convertMat(irow(-1));
    allones = convertMat(iones(1, ncols));
    randwords = convertMat(zeros(1, (1.01 * opts.nneg * nskip * ncols / opts.nreuse).toInt)); // generates random negative words
    randsamp = convertMat(zeros(1, ncols));                                             // For sub-sampling frequent words
    val gopts = opts.asInstanceOf[ADAGrad.Opts];
    vexp = gopts.vexp.v;
    salpha = opts.wsample * math.log(nfeats).toFloat;
    if (useGPU) {
      retEvalPos = GMat(1,1);
      retEvalNeg = GMat(1,1);
    }
    times = zeros(1, ntimes);
    delays = zeros(1, ntimes);
    log = ArrayBuffer();
  }
  
  def dobatch(gmats:Array[Mat], ipass:Int, pos:Long):Unit = {
    addTime(0);
    if (gmats(0).ncols == ncols) {
    	if (firstPos < 0) firstPos = pos;
    	val nsteps = 1f * pos / firstPos;
    	val gopts = opts.asInstanceOf[ADAGrad.Opts];
    	val lrate = gopts.lrate.dv.toFloat * math.pow(nsteps, - gopts.texp.dv).toFloat;
    	val (words, lb, ub, trandwords, goodwords) = wordMats(gmats, ipass, pos);

    	val lrpos = lrate.dv.toFloat;
    	val lrneg = if (opts.eqPosNeg) lrpos else lrpos/opts.nneg;  
    	procPositives(opts.nskip, words, lb, ub, modelmats(1), modelmats(0), lrpos, vexp);
    	addTime(8);
    	procNegatives(opts.nneg, opts.nreuse, trandwords, goodwords, modelmats(1), modelmats(0), lrneg, vexp); 
    	addTime(9);
    }
  }
  
  def evalbatch(gmats:Array[Mat], ipass:Int, pos:Long):FMat = {
  	addTime(0);
  	if (gmats(0).ncols == ncols) {
  	val (words, lb, ub, trandwords, goodwords) = wordMats(gmats, ipass, pos);
  	val epos = evalPositives(opts.nskip, words, lb, ub, modelmats(1), modelmats(0));
  	addTime(10,-3);
    val eneg = evalNegatives(opts.nneg, opts.nreuse, trandwords, goodwords, modelmats(1), modelmats(0));
    addTime(11)
//  	val score = ((epos + eneg / opts.nneg) /opts.nskip / words.ncols);
  	val score = ((epos + eneg / (if (opts.eqPosNeg) 1 else opts.nneg)) / goodwords.length);
  	row(score)
  	} else row(0);
  }
  
  def wordMats(mats:Array[Mat], ipass:Int, pos:Long):(Mat, Mat, Mat, Mat, Mat) = {
  
    val wordsens = mats(0);
    val words = if (opts.iflip) wordsens(1,?) else wordsens(0,?);
    val wgood = words < opts.vocabSize;                                        // Find OOV words 
    addTime(1);
    
    rand(randsamp);                                                            // Take a random sample
    val wrat = float(words+1) * salpha;
    wrat ~ sqrt(wrat) + wrat;
    wgood ~ wgood ∘ int(randsamp < wrat);
    words ~ (wgood ∘ (words + 1)) - 1;                                        // Set OOV or skipped samples to -1
    addTime(2);
        
    rand(ubound);                                                              // get random upper and lower bounds   
    val ubrand = min(opts.nskip, int(ubound * opts.nskip) + 1);
//    println(ubrand.toString)
    val lbrand = - ubrand;
    addTime(3);
    
    val sentencenum = if (opts.iflip) wordsens(0,?) else wordsens(1,?);        // Get the nearest sentence boundaries
    val lbsentence = - cumsumByKey(allones, sentencenum) + 1;
    val ubsentence = reverse(cumsumByKey(allones, reverse(sentencenum))) - 1;
    val lb = max(lbrand, lbsentence);                                          // Combine the bounds
    val ub = min(ubrand, ubsentence);
    test3 = lb
    test4 = ub
    addTime(4);
    
    val (trandwords, contextwords) = (words, lb, ub) match {
      case (giwords:GIMat, gilb:GIMat, giub:GIMat) => {

      	val iwords = minusone \ words \ minusone;                                  // Build a convolution matrix.
      	val cwords = iwords(wordtab);
      	val pgoodwords = (wordmask >= lb) ∘ (wordmask <= ub) ∘ (cwords >= 0) ∘ (words >= 0);  // Find context words satisfying the bound
      	                                                                           // and check that context and center word are good.
      	val fgoodwords = float(pgoodwords);
      	addTime(5);
      	
      	test1 = cwords;

      	rand(randpermute);                                                         // Prepare a random permutation of context words for negative sampling
      	randpermute ~ (fgoodwords ∘ (randpermute + 1f)) - 1f;                      // set the values for bad words to -1.
      	val (vv, ii) = sortdown2(randpermute.view(randpermute.length, 1));         // Permute the good words
      	val ngood = sum(vv >= 0f).dv.toInt;                                        // Count of the good words
      	val ngoodcols = ngood / opts.nreuse;                                       // Number of good columns
      	val cwi = cwords(ii);
      	
      	test2 = cwi
      	addTime(6);

      	rand(randwords);                                                           // Compute some random negatives
      	val irandwords = min(nfeats-1, int(nfeats * (randwords ^ expt)));    
      	val trandwords0 = irandwords.view(opts.nneg, ngoodcols);                   // shrink the matrices to the available data
      	val contextwords0 = cwi.view(opts.nreuse, ngoodcols);
      	addTime(7);
      	(trandwords0, contextwords0)
      }
      case (iwords:IMat, ilb:IMat, iub:IMat) => {
        getnegs(iwords, ilb, iub, Mat.numThreads);
      }
    }
    
    (words, lb, ub, trandwords, contextwords);
  }
  
  def getnegs(words:IMat, lb:IMat, ub:IMat, nthreads:Int):(IMat, IMat) = {
    val ncols = words.ncols;
                                                                               // First count the good context words
    val cwcounts = irow((0 until nthreads).par.map((ithread:Int) => {          // work on blocks
      val istart = ((1L * ncols * ithread)/nthreads).toInt;
      val iend = ((1L * ncols * (ithread + 1))/nthreads).toInt;
      var i = istart;
      var icount = 0;
      while (i < iend) {                                                       // iterate over center words
        if (words.data(i) >= 0) {                                              // check center word is good
        	var j = lb.data(i);                                                  // get lower and upper bounds
        	var jend = ub.data(i);
        	while (j <= jend) {
        		if (j != 0 && words.data(i + j) >= 0) {                            // if not center word and context word is good, count it. 
        		  icount += 1;        		  
        		}
        		j += 1;
        	}
        }
        i += 1;
      }
      icount
    }).toArray)
                                                                               // Now we know how many good words in each block
    val ccc = cumsum(cwcounts);                                                // so size the context word and neg word matrices
    val ngroups = ccc(ccc.length - 1) / opts.nreuse;
    val contextwords0 = izeros(opts.nreuse, ngroups);
    val trandwords0 = izeros(opts.nneg, ngroups);
    
    (0 until nthreads).par.map((ithread:Int) => {                              // Copy the good words into a dense matrix (contextwords0)
      val istart = ((1L * ncols * ithread)/nthreads).toInt;
      val iend = ((1L * ncols * (ithread + 1))/nthreads).toInt;
      var i = istart;
      var icount = 0;
      val mptr = ccc(ithread) - ccc(0);
      while (i < iend) {
        if (words.data(i) >= 0) {
        	var j = lb.data(i);
        	var jend = ub.data(i);
        	while (j <= jend && mptr + icount < contextwords0.length) {
        		if (j != 0 && words.data(i + j) >= 0) {
        		  contextwords0.data(mptr + icount) = words.data(i + j)
        		  icount += 1;        		  
        		}
        		j += 1;
        	}
        }
        i += 1;
      }
      icount
    })
    
    addTime(5);
    
    val prand = drand(opts.nreuse, ngroups);                             // Rands for permutation
    
    var i = 0;                                                           // Permute the good context words randomly
    val n = prand.length;
    while (i < n) {
      val indx = math.min(n-1, i + math.floor(prand.data(i) * (n - i)).toInt);
      if (indx > i) {
        val tmp = contextwords0.data(i);
        contextwords0.data(i) = contextwords0.data(indx);
        contextwords0.data(indx) = tmp;
      }
      i += 1;
    }
    addTime(6);
    
    val randneg = rand(opts.nneg, ngroups);                                // Compute some random negatives

    (0 until nthreads).par.map((ithread:Int) => {                          // Work in blocks over the negs
      val istart = ((1L * ngroups * opts.nneg * ithread)/nthreads).toInt;
      val iend = ((1L * ngroups * opts.nneg * (ithread + 1))/nthreads).toInt;
      var i = istart;
      while (i < iend) {
      	trandwords0.data(i) = math.min(nfeats-1, (nfeats * math.pow(randneg.data(i), expt)).toInt);
        i += 1;
      }
    })
    addTime(7);
    
    (trandwords0, contextwords0)    
  }
  
  def procPositives(nskip:Int, words:Mat, lbound:Mat, ubound:Mat, model1:Mat, model2:Mat, lrate:Float, vexp:Float) = {
    val nrows = model1.nrows;
    val ncols = model1.ncols;
    val nwords = words.ncols;
    Mat.nflops += 6L * nwords * nskip * nrows;
    (words, lbound, ubound, model1, model2) match {
      case (w:GIMat, lb:GIMat, ub:GIMat, m1:GMat, m2:GMat) => {
        val err = CUMACH.word2vecPos(nrows, nwords, nskip, w.data, lb.data, ub.data, m1.data, m2.data, lrate, vexp);
        if (err != 0)  throw new RuntimeException("CUMACH.word2vecPos error " + cudaGetErrorString(err));
      }
      case (w:IMat, lb:IMat, ub:IMat, m1:FMat, m2:FMat) => if (Mat.useMKL) {
        val gw = GIMat(w);
        val glb = GIMat(lb);
        val gub = GIMat(ub);
        val gm1 = GMat(m1);
        val gm2 = GMat(m2);
        val err = CUMACH.word2vecPos(nrows, nwords, nskip, gw.data, glb.data, gub.data, gm1.data, gm2.data, lrate, vexp);
        m1 <-- gm1
        m2 <-- gm2
        if (err != 0) throw new RuntimeException("CUMACH.word2vecEvalPos error " + cudaGetErrorString(err));
//        CPUMACH.word2vecPos(nrows, nwords, nskip, w.data, lb.data, ub.data, m1.data, m2.data, lrate, vexp, Mat.numThreads);
      } else {
        Word2Vec.procPosCPU(nrows, nwords, nskip, w.data, lb.data, ub.data, m1.data, m2.data, lrate, vexp, Mat.numThreads);
      }
    }
  }
  
  def procNegatives(nwa:Int, nwb:Int, wordsa:Mat, wordsb:Mat, modela:Mat, modelb:Mat, lrate:Float, vexp:Float) = {
    val nrows = modela.nrows;
    val ncols = modela.ncols;
    val nwords = wordsa.ncols;
    Mat.nflops += 6L * nwords * nwa * nwb * nrows;
    (wordsa, wordsb, modela, modelb) match {
      case (wa:GIMat, wb:GIMat, ma:GMat, mb:GMat) => {
        val err = CUMACH.word2vecNeg(nrows, nwords, nwa, nwb, wa.data, wb.data, ma.data, mb.data, lrate, vexp);
        if (err != 0) throw new RuntimeException("CUMACH.word2vecNeg error " + cudaGetErrorString(err));
      }
      case (wa:IMat, wb:IMat, ma:FMat, mb:FMat) => if (Mat.useMKL) {
        CPUMACH.word2vecNeg(nrows, nwords, nwa, nwb, wa.data, wb.data, ma.data, mb.data, lrate, vexp, Mat.numThreads);
      } else {
      	Word2Vec.procNegCPU(nrows, nwords, nwa, nwb, wa.data, wb.data, ma.data, mb.data, lrate, vexp, Mat.numThreads);
      }
    }
  }
    
  def evalPositives(nskip:Int, words:Mat, lbound:Mat, ubound:Mat, model1:Mat, model2:Mat):Double = {
    val nrows = model1.nrows;
    val ncols = model1.ncols;
    val nwords = words.ncols;
    Mat.nflops += 2L * nwords * nskip * nrows;
    (words, lbound, ubound, model1, model2) match {
      case (w:GIMat, lb:GIMat, ub:GIMat, m1:GMat, m2:GMat) => {
        retEvalPos.clear
        val err = CUMACH.word2vecEvalPos(nrows, nwords, nskip, w.data, lb.data, ub.data, m1.data, m2.data, retEvalPos.data);
    		if (err != 0) throw new RuntimeException("CUMACH.word2vecEvalPos error " + cudaGetErrorString(err));
        retEvalPos.dv;
      }
      case (w:IMat, lb:IMat, ub:IMat, m1:FMat, m2:FMat) => 
      if (Mat.useMKL) {
      	CPUMACH.word2vecEvalPos(nrows, nwords, nskip, w.data, lb.data, ub.data, m1.data, m2.data, Mat.numThreads);
      } else {
        Word2Vec.evalPosCPU(nrows, nwords, nskip, w.data, lb.data, ub.data, m1.data, m2.data, Mat.numThreads);
      }
    }
  }
  
  def evalNegatives(nwa:Int, nwb:Int, wordsa:Mat, wordsb:Mat, modela:Mat, modelb:Mat):Double = {
    val nrows = modela.nrows;
    val ncols = modela.ncols;
    val nwords = wordsa.ncols;
    Mat.nflops += 2L * nwords * nwa * nwb * nrows;
    (wordsa, wordsb, modela, modelb) match {
      case (wa:GIMat, wb:GIMat, ma:GMat, mb:GMat) => {
        retEvalNeg.clear
        val err = CUMACH.word2vecEvalNeg(nrows, nwords, nwa, nwb, wa.data, wb.data, ma.data, mb.data, retEvalNeg.data);
    		if (err != 0) throw new RuntimeException("CUMACH.word2vecEvalNeg error " + cudaGetErrorString(err));
        retEvalNeg.dv;      
      }
      case (wa:IMat, wb:IMat, ma:FMat, mb:FMat) => 
      if (Mat.useMKL) {
      	CPUMACH.word2vecEvalNeg(nrows, nwords, nwa, nwb, wa.data, wb.data, ma.data, mb.data, Mat.numThreads); 
      } else {
      	Word2Vec.evalNegCPU(nrows, nwords, nwa, nwb, wa.data, wb.data, ma.data, mb.data, Mat.numThreads);
      }
    }
  }
  
  def trailingZeros(a:Long):Int = {
    var aa = a;
    var nz = 0;
    while ((aa & 1L) == 0) {
      aa = aa >> 1;
      nz += 1;
    }
    nz
  }
  
  override def mergeModelFn(models:Array[Model], mm:Array[Mat], um:Array[Mat], istep:Long):Unit = {
    val headlen = if (istep > 0) math.max(opts.headlen, opts.headlen << trailingZeros(istep)) else 0;
    val mlen = models(0).modelmats.length;
    val thisGPU = getGPU;
    val modj = new Array[Mat](models.length);
    for (j <- 0 until mlen) {
      val mmj = if (headlen > 0) mm(j).view(mm(j).nrows, math.min(mm(j).ncols, headlen)) else mm(j);
      mmj.clear
      for (i <- 0 until models.length) {
        if (useGPU && i < Mat.hasCUDA) setGPU(i);
        modj(i) = if (headlen > 0) models(i).modelmats(j).view(models(i).modelmats(j).nrows, math.min(models(i).modelmats(j).ncols, headlen)) else models(i).modelmats(j);
        val umj = if (headlen > 0) um(j).view(um(j).nrows, math.min(um(j).ncols, headlen)) else um(j);
        umj <-- modj(i)
        mmj ~ mmj + umj;
      }
      mmj ~ mmj * (1f/models.length);
      for (i <- 0 until models.length) {
        modj(i) <-- mmj;
      }
    }
    setGPU(thisGPU);
  }

}

object Word2Vec  {
  trait Opts extends Model.Opts {
    var aopts:ADAGrad.Opts = null;
    var nskip = 5;
    var nneg = 5;
    var nreuse = 5;  
    var vocabSize = 100000;
    var wexpt = 0.75f;
    var wsample = 1e-4f;
    var headlen = 10000;
    var iflip = false;
    var eqPosNeg = false;
  }
  
  class Options extends Opts {}
  
  
  def procPosCPU(nrows:Int, ncols:Int, skip:Int, W:Array[Int], LB:Array[Int], UB:Array[Int],
      A:Array[Float], B:Array[Float], lrate:Float, vexp:Float, nthreads:Int):Int = {

    (0 until nthreads).par.map((ithread:Int) => {
    	val istart = ((1L * ithread * ncols)/nthreads).toInt;
    	val iend = ((1L * (ithread+1) * ncols)/nthreads).toInt;
    	val daa = new Array[Float](nrows);
    	var i = istart;
    	while (i < iend) {
    	  var j = 0;
    	  var k = 0;
    	  var c = 0;
    	  var cv = 0f;

    	  val iac = W(i);
    	  val ascale = math.pow(iac, vexp).toFloat;
    	  val ia = nrows * iac;                                        // Get the current word (as a model matrix offset). 
    	  if (ia >= 0) {                                               // Check for OOV words
{'fraction_non_alphanumeric': 0.134154707351206, 'fraction_numerical': 0.015311582471256482, 'mean_word_length': 2.004947229551451, 'pattern_counts': {'":': 0, '<': 47, '<?xml version=': 0, '>': 67, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 110, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2971628', 'n_tokens_mistral': 13526, 'n_tokens_neox': 13122, 'n_words': 3093}