bayes-probas/lib/naive_bayes.js

Naive Bayes Classifier for node.js

/*
    Expose our naive-bayes generator function
 */
module.exports = function(options) {
  return new Naivebayes(options);
};

// keys we use to serialize a classifier's state
const STATE_KEYS = (module.exports.STATE_KEYS = [
  'categories',
  'docCount',
  'totalDocuments',
  'vocabulary',
  'vocabularySize',
  'wordCount',
  'wordFrequencyCount',
  'options',
]);

/**
 * Initializes a NaiveBayes instance from a JSON state representation.
 * Use this with classifier.toJson().
 *
 * @param  {String} jsonStr   state representation obtained by classifier.toJson()
 * @return {NaiveBayes}       Classifier
 */
module.exports.fromJson = jsonStr => {
  let parsed;

  try {
    parsed = JSON.parse(jsonStr);
  } catch (e) {
    console.error(e);
    throw new Error('Naivebayes.fromJson expects a valid JSON string.');
  }

  // init a new classifier
  let classifier = new Naivebayes(parsed.options);

  // override the classifier's state
  STATE_KEYS.forEach(k => {
    if (!parsed[k]) {
      throw new Error(
        `Naivebayes.fromJson: JSON string is missing an expected property: [${k}].`
      );
    }
    classifier[k] = parsed[k];
  });

  return classifier;
};

/**
 * Given an input string, tokenize it into an array of word tokens.
 * This is the default tokenization function used if user does not provide one in `options`.
 *
 * @param  {String} text
 * @return {Array}
 */
const defaultTokenizer = text => {
  //remove punctuation from text - remove anything that isn't a word char or a space
  let rgxPunctuation = /[^(a-zA-ZA-Яa-я0-9_)+\s]/g;

  let sanitized = text.replace(rgxPunctuation, ' ');

  return sanitized.split(/\s+/);
};

/**
 * Naive-Bayes Classifier
 *
 * This is a naive-bayes classifier that uses Laplace Smoothing.
 *
 * Takes an (optional) options object containing:
 *   - `tokenizer`  => custom tokenization function
 *
 */
function Naivebayes(options) {
  // set options object
  this.options = {};
  if (typeof options !== 'undefined') {
    if (!options || typeof options !== 'object' || Array.isArray(options)) {
      throw TypeError(
        `NaiveBayes got invalid 'options': ${options}'. Pass in an object.`
      );
    }
    this.options = options;
  }

  this.tokenizer = this.options.tokenizer || defaultTokenizer;

  //initialize our vocabulary and its size
  this.vocabulary = {};
  this.vocabularySize = 0;

  //number of documents we have learned from
  this.totalDocuments = 0;

  //document frequency table for each of our categories
  //=> for each category, how often were documents mapped to it
  this.docCount = {};

  //for each category, how many words total were mapped to it
  this.wordCount = {};

  //word frequency table for each category
  //=> for each category, how frequent was a given word mapped to it
  this.wordFrequencyCount = {};

  //hashmap of our category names
  this.categories = {};
}

/**
 * Initialize each of our data structure entries for this new category
 *
 * @param  {String} categoryName
 */
Naivebayes.prototype.initializeCategory = function(categoryName) {
  if (!this.categories[categoryName]) {
    this.docCount[categoryName] = 0;
    this.wordCount[categoryName] = 0;
    this.wordFrequencyCount[categoryName] = {};
    this.categories[categoryName] = true;
  }
  return this;
};

/**
 * train our naive-bayes classifier by telling it what `category`
 * the `text` corresponds to.
 *
 * @param  {String} text
 * @param  {String} class
 */
Naivebayes.prototype.learn = function(text, category) {
  //initialize category data structures if we've never seen this category
  this.initializeCategory(category);

  //update our count of how many documents mapped to this category
  this.docCount[category]++;

  //update the total number of documents we have learned from
  this.totalDocuments++;

  //normalize the text into a word array
  let tokens = this.tokenizer(text);

  //get a frequency count for each token in the text
  let frequencyTable = this.frequencyTable(tokens);

  Object.keys(frequencyTable).forEach(token => {
    //add this word to our vocabulary if not already existing
    if (!this.vocabulary[token] || this.vocabulary[token] === 0) {
      this.vocabularySize++;
      this.vocabulary[token] = 1;
    } else if (this.vocabulary[token] > 0) {
      this.vocabulary[token]++;
    }

    let frequencyInText = frequencyTable[token];

    //update the frequency information for this word in this category
    if (!this.wordFrequencyCount[category][token])
      this.wordFrequencyCount[category][token] = frequencyInText;
    else this.wordFrequencyCount[category][token] += frequencyInText;

    //update the count of all words we have seen mapped to this category
    this.wordCount[category] += frequencyInText;
  });

  return this;
};

/**
 * untrain our naive-bayes classifier by telling it what `category`
 * the `text` to remove corresponds to.
 *
 * @param  {String} text
 * @param  {String} class
 */
Naivebayes.prototype.unlearn = function(text, category){
  //update our count of how many documents mapped to this category
  this.docCount[category]--;

  //update the total number of documents we have learned from
  this.totalDocuments--;

  //normalize the text into a word array
  let tokens = this.tokenizer(text);

  //get a frequency count for each token in the text
  let frequencyTable = this.frequencyTable(tokens);

  /*
   Update our vocabulary and our word frequency count for this category
   */

  Object.keys(frequencyTable).forEach(token => {
    //add this word to our vocabulary if not already existing
    if (this.vocabulary[token] && this.vocabulary[token] > 0) {
      this.vocabulary[token]--;
      if (this.vocabulary[token] === 0) this.vocabularySize--;
    }

    let frequencyInText = frequencyTable[token];

    this.wordFrequencyCount[category][token] -= frequencyInText;

    //update the count of all words we have seen mapped to this category
    this.wordCount[category] -= frequencyInText;
  });

  return this;
};

/**
 * Determine and rank categories `text` belongs to.
 *
 * @param  {String} text
 * @return {Object} categories sorted by pertinence
 */
Naivebayes.prototype.categorizeObj = function(text){
  let maxProbability = -Infinity,
    chosenCategory = null;

  let tokens = this.tokenizer(text);
  let frequencyTable = this.frequencyTable(tokens);
  let categoryProbas = [];

  //iterate thru our categories to find the one with max probability for this text
  Object.keys(this.categories).forEach(category => {
    //start by calculating the overall probability of this category
    //=>  out of all documents we've ever looked at, how many were
    //    mapped to this category
    let categoryProbability = this.docCount[category] / this.totalDocuments;

    //take the log to avoid underflow
    let logProbability = Math.log(categoryProbability);

    //now determine P( w | c ) for each word `w` in the text
    Object.keys(frequencyTable).forEach(token => {
      let frequencyInText = frequencyTable[token];
      let tokenProbability = this.tokenProbability(token, category);

      //determine the log of the P( w | c ) for this word
      logProbability += frequencyInText * Math.log(tokenProbability);
    });
    categoryProbas.push({
      category: category,
      proba: logProbability,
    });

    if (logProbability > maxProbability) {
      maxProbability = logProbability;
      chosenCategory = category;
    }
  });
  // sort categoryProbas;
  categoryProbas.sort((a, b) => {
    return a.proba - b.proba;
  });

  return categoryProbas;
};

/**
 * Determine what category `text` belongs to.
 *
 * @param  {String} text
 * @return {String} category
 */
Naivebayes.prototype.categorize = function(text){
  let maxProbability = -Infinity,
    minProbability = Infinity,
    chosenCategory = null;

  let tokens = this.tokenizer(text);
  let frequencyTable = this.frequencyTable(tokens);
  let categoryProbas = [];

  //iterate thru our categories to find the one with max probability for this text
  Object.keys(this.categories).forEach(category => {
    //start by calculating the overall probability of this category
    //=>  out of all documents we've ever looked at, how many were
    //    mapped to this category
    let categoryProbability = this.docCount[category] / this.totalDocuments;

    //take the log to avoid underflow
    let logProbability = Math.log(categoryProbability);

    //now determine P( w | c ) for each word `w` in the text
    Object.keys(frequencyTable).forEach(token => {
      let frequencyInText = frequencyTable[token];
      let tokenProbability = this.tokenProbability(token, category);

      // determine the log of the P( w | c ) for this word
      logProbability += frequencyInText * Math.log(tokenProbability);
    });

    if (logProbability == Number.NEGATIVE_INFINITY) {
      console.error(
        `category ${category} had -Infinity odds, didn't add it to the categoryProbas array`
      );
    } else {
      categoryProbas.push({ category: category, proba: logProbability });
      if (logProbability > maxProbability) {
        maxProbability = logProbability;
        chosenCategory = category;
      }
      if (logProbability != null && logProbability < minProbability) {
        minProbability = logProbability;
      }
    }
  });

  // sort categoryProbas;
  categoryProbas.sort((a, b) => {
    return a.proba - b.proba;
  });

  // from log-likelihood probas get more human readable values (0 = less probable -100 most probable)
  let totalNum = categoryProbas.length;
  let bigDelta = maxProbability - minProbability;

  categoryProbas.forEach((value, index) => {
    let deltaValue = value.proba - minProbability;
    value.probaH = Math.round(deltaValue * 100 / bigDelta, 10); // Round handles values like 99.999999..9 and 100.00000...1
  });

  return { probas: categoryProbas, chosenCategory: chosenCategory };
};

/**
 * Calculate probability that a `token` belongs to a `category`
 *
 * @param  {String} token
 * @param  {String} category
 * @return {Number} probability
 */
Naivebayes.prototype.tokenProbability = function(token, category){
  //how many times this word has occurred in documents mapped to this category
  let wordFrequencyCount = this.wordFrequencyCount[category][token] || 0;

  //what is the count of all words that have ever been mapped to this category
  let wordCount = this.wordCount[category];

  //use laplace Add-1 Smoothing equation
  return (wordFrequencyCount + 1) / (wordCount + this.vocabularySize);
};

/**
 * Build a frequency hashmap where
 * - the keys are the entries in `tokens`
 * - the values are the frequency of each entry in `tokens`
 *
 * @param  {Array} tokens  Normalized word array
 * @return {Object}
 */
Naivebayes.prototype.frequencyTable = function(tokens) {
  let frequencyTable = Object.create(null);

  tokens.forEach(token => {
    if (!frequencyTable[token]) frequencyTable[token] = 1;
    else frequencyTable[token]++;
  });

  return frequencyTable;
};

/**
 * Dump the classifier's state as a JSON string.
 * @return {String} Representation of the classifier.
 */
Naivebayes.prototype.toJson = function() {
  let state = {};

  STATE_KEYS.forEach(k => (state[k] = this[k]));

  return JSON.stringify(state);
};