Merge pull request jekyll#61 from mike-stewart/enable_stemmer

Option to Disable Stemming

Author: Ch4s3

lib/classifier-reborn/bayes.rb

@@ -17,12 +17,14 @@ class Bayes
     #   auto_categorize:  false  When true, enables ability to dynamically declare a category
     #   enable_threshold: false  When true, enables a threshold requirement for classifition
     #   threshold:        0.0    Default threshold, only used when enabled
+    #   enable_stemmer:   true   When false, disables word stemming
     def initialize(*args)
       @categories = {}
       options = { language:         'en',
                   auto_categorize:  false,
                   enable_threshold: false,
-                  threshold:        0.0
+                  threshold:        0.0,
+                  enable_stemmer:   true
                 }
       args.flatten.each do |arg|
         if arg.is_a?(Hash)
@@ -40,6 +42,7 @@ def initialize(*args)
       @auto_categorize     = options[:auto_categorize]
       @enable_threshold    = options[:enable_threshold]
       @threshold           = options[:threshold]
+      @enable_stemmer      = options[:enable_stemmer]
     end
 
     # Provides a general training method for all categories specified in Bayes#new
@@ -61,7 +64,7 @@ def train(category, text)
       end
 
       @category_counts[category] += 1
-      Hasher.word_hash(text, @language).each do |word, count|
+      Hasher.word_hash(text, @language, @enable_stemmer).each do |word, count|
         @categories[category][word] += count
         @category_word_count[category] += count
         @total_words += count
@@ -78,7 +81,7 @@ def train(category, text)
     def untrain(category, text)
       category = CategoryNamer.prepare_name(category)
       @category_counts[category] -= 1
-      Hasher.word_hash(text, @language).each do |word, count|
+      Hasher.word_hash(text, @language, @enable_stemmer).each do |word, count|
         next if @total_words < 0
         orig = @categories[category][word] || 0
         @categories[category][word] -= count
@@ -98,7 +101,7 @@ def untrain(category, text)
     # The largest of these scores (the one closest to 0) is the one picked out by #classify
     def classifications(text)
       score = {}
-      word_hash = Hasher.word_hash(text, @language)
+      word_hash = Hasher.word_hash(text, @language, @enable_stemmer)
       training_count = @category_counts.values.reduce(:+).to_f
       @categories.each do |category, category_words|
         score[category.to_s] = 0
@@ -155,6 +158,16 @@ def threshold_disabled?
       !@enable_threshold
     end
 
+    # Is word stemming enabled?
+    def stemmer_enabled?
+      @enable_stemmer
+    end
+
+    # Is word stemming disabled?
+    def stemmer_disabled?
+      !@enable_stemmer
+    end
+
     # Provides training and untraining methods for the categories specified in Bayes#new
     # For example:
     #     b = ClassifierReborn::Bayes.new 'This', 'That', 'the_other'

lib/classifier-reborn/extensions/hasher.rb

@@ -13,21 +13,26 @@ module Hasher
 
     # Return a Hash of strings => ints. Each word in the string is stemmed,
     # interned, and indexes to its frequency in the document.
-    def word_hash(str, language = 'en')
-      cleaned_word_hash = clean_word_hash(str, language)
+    def word_hash(str, language = 'en', enable_stemmer = true)
+      cleaned_word_hash = clean_word_hash(str, language, enable_stemmer)
       symbol_hash = word_hash_for_symbols(str.scan(/[^\s\p{WORD}]/))
       cleaned_word_hash.merge(symbol_hash)
     end
 
     # Return a word hash without extra punctuation or short symbols, just stemmed words
-    def clean_word_hash(str, language = 'en')
-      word_hash_for_words str.gsub(/[^\p{WORD}\s]/, '').downcase.split, language
+    def clean_word_hash(str, language = 'en', enable_stemmer = true)
+      word_hash_for_words str.gsub(/[^\p{WORD}\s]/, '').downcase.split, language, enable_stemmer
     end
 
-    def word_hash_for_words(words, language = 'en')
+    def word_hash_for_words(words, language = 'en', enable_stemmer = true)
       d = Hash.new(0)
       words.each do |word|
-        d[word.stem.intern] += 1 if word.length > 2 && !STOPWORDS[language].include?(word)
+        next unless word.length > 2 && !STOPWORDS[language].include?(word)
+        if enable_stemmer
+          d[word.stem.intern] += 1
+        else
+          d[word.intern] += 1
+        end
       end
       d
     end

test/bayes/bayesian_test.rb

@@ -14,6 +14,10 @@ def test_training_with_utf8
     assert_nothing_raised { @classifier.train_interesting 'Água' }
   end
 
+  def test_stemming_enabled_by_default
+    assert @classifier.stemmer_enabled?
+  end
+
   def test_bad_training
     assert_raise(StandardError) { @classifier.train_no_category 'words' }
   end

test/extensions/hasher_test.rb

@@ -15,6 +15,11 @@ def test_clean_word_hash
     assert_equal hash, Hasher.clean_word_hash("here are some good words of test's. I hope you love them!")
   end
 
+  def test_clean_word_hash_without_stemming
+    hash = { good: 1, words: 1, hope: 1, love: 1, them: 1, tests: 1 }
+    assert_equal hash, Hasher.clean_word_hash("here are some good words of test's. I hope you love them!", 'en', false)
+  end
+
   def test_default_stopwords
     assert_not_empty Hasher::STOPWORDS['en']
     assert_not_empty Hasher::STOPWORDS['fr']