lib/lda-ruby/vocabulary.rb in ealdent-lda-ruby-0.3.0 vs lib/lda-ruby/vocabulary.rb in ealdent-lda-ruby-0.3.1

@@ -1,8 +1,8 @@
 module Lda
   class Vocabulary
-    attr_reader :words
+    attr_reader :words, :indexes
 
     def initialize(words = nil)
       @words = Hash.new do |hash, key|
         if hash.member?(:MAX_VALUE)
           hash[:MAX_VALUE] = hash[:MAX_VALUE] + 1
@@ -11,13 +11,20 @@
         end
         hash[key] = hash[:MAX_VALUE]
       end
 
       words.each { |w| @words[w] } if words
+      @indexes = Hash.new
+
+      @words.each_pair do |w, i|
+        @indexes[i] = w
+      end
     end
 
     def check_word(word)
-      @words[word.dup]
+      w = @words[word.dup]
+      @indexes[w] = word.dup
+      w
     end
 
     def load_file(filename)
       txt = File.open(filename, 'r') { |f| f.read }
       txt.split(/[\n\r]+/).each { |word| check_word(word) }
\ No newline at end of file