lib/bae/classifier.rb in bae-0.0.9 vs lib/bae/classifier.rb in bae-0.1.0

@@ -85,24 +85,17 @@
 
       normalize(posterior)
     end
 
     def save_state(path)
-      state = {}
-      state['frequency_table'] = frequency_table
-      state['label_instance_count'] = label_instance_count
-      state['label_index'] = label_index
-      state['label_index_sequence'] = label_index_sequence
-      state['total_terms'] = total_terms
-
       ::File.open(::File.expand_path(path), 'w') do |handle|
-        handle.write(state.to_json)
+        handle.write(to_json)
       end
     end
 
-    def load_state(path)
-      state = ::JSON.parse(::File.read(::File.expand_path(path)))
+    def load_from_json(json)
+      state = ::JSON.parse(json)
 
       fail 'Missing frequency_table' unless state['frequency_table']
       fail 'Missing label_instance_count' unless state['label_instance_count']
       fail 'Missing label_index' unless state['label_index']
       fail 'Missing label_index_sequence' unless state['label_index_sequence']
@@ -113,9 +106,24 @@
       @label_index = state['label_index']
       @label_index_sequence = state['label_index_sequence']
       @total_terms = state['total_terms']
 
       finish_training!
+    end
+
+    def load_state(path)
+      state_json = ::File.read(::File.expand_path(path))
+      load_from_json(state_json)
+    end
+
+    def to_json
+      state = {}
+      state['frequency_table'] = frequency_table
+      state['label_instance_count'] = label_instance_count
+      state['label_index'] = label_index
+      state['label_index_sequence'] = label_index_sequence
+      state['total_terms'] = total_terms
+      state.to_json
     end
 
   private
 
     def calculate_likelihoods!