Permalink
Browse files

Merge pull request #1 from GCorbel/master

Add the untrain function
  • Loading branch information...
2 parents 08bbc74 + 25f2c9f commit e9e48bc3dea92bafa23d4c4c4d83ed03254915ab @reddavis committed Jan 29, 2012
Showing with 16 additions and 2 deletions.
  1. +8 −1 lib/naive_bayes.rb
  2. +1 −1 spec/db/naive.nb
  3. +7 −0 spec/naive_bayes_spec.rb
View
9 lib/naive_bayes.rb
@@ -45,6 +45,13 @@ def train(klass, *features)
@klass_count[klass] += 1
end
+ def untrain(klass, *features)
+ features.uniq.each do |feature|
+ @features_count[klass][feature] -= 1
+ end
+ @klass_count[klass] -= 1
+ end
+
#P(Class | Item) = P(Item | Class) * P(Class)
def classify(*features)
scores = {}
@@ -95,4 +102,4 @@ def assumed_probability
0.5 / (total_items/2)
end
-end
+end
View
2 spec/db/naive.nb
@@ -1 +1 @@
-o:NaiveBayes :@klass_count{: spamf1:hamf1:@db_filepath"U/Users/reddavis/Documents/projects/machine_learning/naive_bayes/spec/db/naive.nb:@features_count{;}" wordf1"badf1f0;}"wef1"badf1@:
+o:NaiveBayes :@features_count{: spam}I"bad:EFf1I" word;Ff1f0:ham}I"we;Ff1I"bad;Ff1@:@klass_count{;f1; f1:
View
7 spec/naive_bayes_spec.rb
@@ -12,6 +12,13 @@
a[0].should == :spam
a[1].should == 0.5
end
+
+ it "should change information" do
+ @classifier.untrain(:spam, 'bad', 'word')
+ a = @classifier.classify('bad', 'word')
+ a[0].should == :ham
+ a[1].should == 1.0
+ end
end
describe "Saving the NB" do

0 comments on commit e9e48bc

Please sign in to comment.