<commit_before>package org.clulab.lm

import com.typesafe.config.ConfigFactory
import edu.cmu.dynet.{Dim, LstmBuilder, ParameterCollection}
import org.clulab.sequences.LstmUtils
import org.clulab.struct.Counter
import org.clulab.utils.Serializer
import org.slf4j.{Logger, LoggerFactory}

import scala.io.Source

/**
 * Constructs the RnnLM model
 * This starts similarly to Lample et al. (2016), but it includes a LM training component as well
 * This contains static word embeddings (GloVe) + biLSTM character embeddings that are learned with the task,
 *   which then feed into a word-level biLSTM
 */
object RnnLMTrain {
  val logger:Logger = LoggerFactory.getLogger(classOf[RnnLMTrain])

  def main(args: Array[String]): Unit = {
    LstmUtils.initializeDyNet() // autoBatch = true, mem = "1660,1664,2496,1400")
    val configName = "rnnlm-en"
    val config = new FlairConfig(ConfigFactory.load(configName))

    val charEmbeddingSize = config.getArgInt("rnnlm.train.charEmbeddingSize", Some(32))
    val charRnnStateSize = config.getArgInt("rnnlm.train.charRnnStateSize", Some(16))
    val wordRnnStateSize = config.getArgInt("rnnlm.train.wordRnnStateSize", Some(256))
    val posTagEmbeddingSize = config.getArgInt("rnnlm.train.posTagEmbeddingSize", Some(32))
    val positionEmbeddingSize = config.getArgInt("rnnlm.train.positionEmbeddingSize", Some(16))
    val positionWindowSize = config.getArgInt("rnnlm.train.positionWindowSize", Some(20))

    //
    // This stores all parameters in the RnnLM
    //
    val parameters = new ParameterCollection()

    //
    // Load the character map
    //
    logger.debug(s"Loading the character map...")
    val c2iFilename = config.getArgString("rnnlm.train.c2i", None)
    val c2i = Serializer.using(LstmUtils.newSource(c2iFilename)) { source =>
      val byLineCharMapBuilder = new LstmUtils.ByLineCharIntMapBuilder()
      val lines = source.getLines()
      val c2i = byLineCharMapBuilder.build(lines)
      c2i
    }
    logger.debug(s"Loaded a character map with ${c2i.keySet.size} entries.")

    //
    // Load the word embeddings
    //
    logger.debug("Loading word embeddings...")
    val embedFilename = config.getArgString("rnnlm.train.embed", None)
    val docFreqFilename = config.getArgString("rnnlm.train.docFreq", None)
    val minFreq = config.getArgInt("rnnlm.train.minWordFreq", Some(100))
    val minMandatoryFreq = config.getArgInt("rnnlm.train.minMandatoryWordFreq", Some(1))
    val w2v = LstmUtils.loadEmbeddings(
      Some(docFreqFilename), minFreq, embedFilename,
      Some(config.getArgString("rnnlm.train.mandatoryWords", None)), minMandatoryFreq)
    val w2i = LstmUtils.mkWordVocab(w2v)

    //
    // Construct the vocab from the words in the training dataset(s)
    //
    val tw2i = mkTrainWordVocab(config.getArgString("rnnlm.train.taskTrainCorpus", None), minMandatoryFreq)
    val trainWordEmbeddingSize = config.getArgInt("rnnlm.train.trainWordEmbeddingSize", Some(32))
    val trainWordLookupParameters = parameters.addLookupParameters(tw2i.size, Dim(trainWordEmbeddingSize))

    //
    // POS tag embeddings
    //
    val pos2is = LstmUtils.readString2Ids("org/clulab/lm/pos2i-en.txt")
    val posEmbeddings = parameters.addLookupParameters(pos2is.size, Dim(posTagEmbeddingSize))

    //
    // Position embeddings [-20, 20] + < 20 + > 20. total = 43
    //
    val positionEmbeddings = parameters.addLookupParameters(positionWindowSize * 2 + 3, Dim(positionEmbeddingSize))

    //
    // Convert the word embeddings we loaded above into DyNet LookupParameters
    //
    val wordLookupParameters = parameters.addLookupParameters(w2i.size, Dim(w2v.dimensions))
    LstmUtils.initializeEmbeddings(w2v, w2i, wordLookupParameters)
    logger.debug("Completed loading word embeddings.")

    //
    // Character embeddings and character biLSTM, initialized randomly
    //
    val charLookupParameters = parameters.addLookupParameters(c2i.size, Dim(charEmbeddingSize))
    val charFwRnnBuilder = new LstmBuilder(1, charEmbeddingSize, charRnnStateSize, parameters)
    val charBwRnnBuilder = new LstmBuilder(1, charEmbeddingSize, charRnnStateSize, parameters)

    val embeddingSize = 2 * charRnnStateSize + w2v.dimensions + trainWordEmbeddingSize + 1 + posTagEmbeddingSize + positionEmbeddingSize // 1 for isPredFeature
    val fwBuilder = new LstmBuilder(4, embeddingSize, wordRnnStateSize, parameters)
    val bwBuilder = new LstmBuilder(4, embeddingSize, wordRnnStateSize, parameters)

    //
    // Feed forward networks, used for LM prediction, only if training is enabled
    //
    val lmLabelCount = config.getArgInt("rnnlm.train.lmLabelCount", Some(40000)) + 2 // + 2 for UNK and EOS
    val fwO = parameters.addParameters(Dim(lmLabelCount, wordRnnStateSize))
    val bwO = parameters.addParameters(Dim(lmLabelCount, wordRnnStateSize))

    //
    // Create the LM object
    //
    val lm = new RnnLM(w2i, tw2i, c2i, pos2is,
      wordRnnStateSize, charRnnStateSize, lmLabelCount,
      positionEmbeddingSize, positionWindowSize,
      parameters,
      wordLookupParameters, trainWordLookupParameters, charLookupParameters,
      posEmbeddings, positionEmbeddings,
      charFwRnnBuilder, charBwRnnBuilder,<commit_msg>bug fix<commit_after>package org.clulab.lm

import com.typesafe.config.ConfigFactory
import edu.cmu.dynet.{Dim, LstmBuilder, ParameterCollection}
import org.clulab.sequences.LstmUtils
import org.clulab.struct.Counter
import org.clulab.utils.Serializer
import org.slf4j.{Logger, LoggerFactory}

import scala.io.Source

/**
 * Constructs the RnnLM model
 * This starts similarly to Lample et al. (2016), but it includes a LM training component as well
 * This contains static word embeddings (GloVe) + biLSTM character embeddings that are learned with the task,
 *   which then feed into a word-level biLSTM
 */
object RnnLMTrain {
  val logger:Logger = LoggerFactory.getLogger(classOf[RnnLMTrain])

  def main(args: Array[String]): Unit = {
    LstmUtils.initializeDyNet() // autoBatch = true, mem = "1660,1664,2496,1400")
    val configName = "rnnlm-en"
    val config = new FlairConfig(ConfigFactory.load(configName))

    val charEmbeddingSize = config.getArgInt("rnnlm.train.charEmbeddingSize", Some(32))
    val charRnnStateSize = config.getArgInt("rnnlm.train.charRnnStateSize", Some(16))
    val wordRnnStateSize = config.getArgInt("rnnlm.train.wordRnnStateSize", Some(256))
    val posTagEmbeddingSize = config.getArgInt("rnnlm.train.posTagEmbeddingSize", Some(32))
    val positionEmbeddingSize = config.getArgInt("rnnlm.train.positionEmbeddingSize", Some(16))
    val positionWindowSize = config.getArgInt("rnnlm.train.positionWindowSize", Some(20))

    //
    // This stores all parameters in the RnnLM
    //
    val parameters = new ParameterCollection()

    //
    // Load the character map
    //
    logger.debug(s"Loading the character map...")
    val c2iFilename = config.getArgString("rnnlm.train.c2i", None)
    val c2i = Serializer.using(LstmUtils.newSource(c2iFilename)) { source =>
      val byLineCharMapBuilder = new LstmUtils.ByLineCharIntMapBuilder()
      val lines = source.getLines()
      val c2i = byLineCharMapBuilder.build(lines)
      c2i
    }
    logger.debug(s"Loaded a character map with ${c2i.keySet.size} entries.")

    //
    // Load the word embeddings
    //
    logger.debug("Loading word embeddings...")
    val embedFilename = config.getArgString("rnnlm.train.embed", None)
    val docFreqFilename = config.getArgString("rnnlm.train.docFreq", None)
    val minFreq = config.getArgInt("rnnlm.train.minWordFreq", Some(100))
    val minMandatoryFreq = config.getArgInt("rnnlm.train.minMandatoryWordFreq", Some(1))
    val w2v = LstmUtils.loadEmbeddings(
      Some(docFreqFilename), minFreq, embedFilename,
      Some(config.getArgString("rnnlm.train.mandatoryWords", None)), minMandatoryFreq)
    val w2i = LstmUtils.mkWordVocab(w2v)

    //
    // Construct the vocab from the words in the training dataset(s)
    //
    val tw2i = mkTrainWordVocab(config.getArgString("rnnlm.train.taskTrainCorpus", None), minMandatoryFreq)
    val trainWordEmbeddingSize = config.getArgInt("rnnlm.train.trainWordEmbeddingSize", Some(32))
    val trainWordLookupParameters = parameters.addLookupParameters(tw2i.size, Dim(trainWordEmbeddingSize))

    //
    // POS tag embeddings
    //
    val pos2is = LstmUtils.readString2Ids("org/clulab/lm/pos2i-en.txt")
    val posEmbeddings = parameters.addLookupParameters(pos2is.size, Dim(posTagEmbeddingSize))

    //
    // Position embeddings [-20, 20] + < 20 + > 20. total = 43
    //
    val positionEmbeddings = parameters.addLookupParameters(positionWindowSize * 2 + 3, Dim(positionEmbeddingSize))

    //
    // Convert the word embeddings we loaded above into DyNet LookupParameters
    //
    val wordLookupParameters = parameters.addLookupParameters(w2i.size, Dim(w2v.dimensions))
    LstmUtils.initializeEmbeddings(w2v, w2i, wordLookupParameters)
    logger.debug("Completed loading word embeddings.")

    //
    // Character embeddings and character biLSTM, initialized randomly
    //
    val charLookupParameters = parameters.addLookupParameters(c2i.size, Dim(charEmbeddingSize))
    val charFwRnnBuilder = new LstmBuilder(1, charEmbeddingSize, charRnnStateSize, parameters)
    val charBwRnnBuilder = new LstmBuilder(1, charEmbeddingSize, charRnnStateSize, parameters)

    val embeddingSize = 2 * charRnnStateSize + w2v.dimensions + // trainWordEmbeddingSize +
      1 + posTagEmbeddingSize + positionEmbeddingSize // 1 for isPredFeature
    val fwBuilder = new LstmBuilder(4, embeddingSize, wordRnnStateSize, parameters)
    val bwBuilder = new LstmBuilder(4, embeddingSize, wordRnnStateSize, parameters)

    //
    // Feed forward networks, used for LM prediction, only if training is enabled
    //
    val lmLabelCount = config.getArgInt("rnnlm.train.lmLabelCount", Some(40000)) + 2 // + 2 for UNK and EOS
    val fwO = parameters.addParameters(Dim(lmLabelCount, wordRnnStateSize))
    val bwO = parameters.addParameters(Dim(lmLabelCount, wordRnnStateSize))

    //
    // Create the LM object
    //
    val lm = new RnnLM(w2i, tw2i, c2i, pos2is,
      wordRnnStateSize, charRnnStateSize, lmLabelCount,
      positionEmbeddingSize, positionWindowSize,
      parameters,
      wordLookupParameters, trainWordLookupParameters, charLookupParameters,
      posEmbeddings, positionEmbeddings,
      charFwRnnBuilder, charBwRnnBuilder,
{'fraction_non_alphanumeric': 0.10064096431646417, 'fraction_numerical': 0.01760260212379221, 'mean_word_length': 4.0673776054289865, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, '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': '2645568', 'n_tokens_mistral': 3482, 'n_tokens_neox': 3290, 'n_words': 720}