lib/markov/chain.rb in markov-rb-0.1.2 vs lib/markov/chain.rb in markov-rb-0.1.3

@@ -70,19 +70,11 @@
       generated_count = 0
       generated_text = start_text.chars
       word = nil
 
       until max && (generated_count +=1)>=max
-        if new_words
-          word = @words.first
-          until !@words.include?(word.chomp) && word.chomp.match(/^[A-Z][a-z]+$/)
-            word = (generate_word!) #.chomp #(context: generated_text)
-          end
-          @words.push(word.chomp)
-        else
-          word = generate_word!(context: generated_text)
-        end
+        word = generate_word!(context: new_words ? ["\n"] : generated_text)
 
         if show
           word.chars.each do |ch|
             sleep 0.05 + rand*(0.01)
             yield ch
@@ -95,13 +87,27 @@
 
       generated_text.join
     end
 
     def generate_word!(context: ["\n"])
+      if new_words
+        word = @words.first
+        until !@words.include?(word.chomp) && word.chomp.match(/^[A-Z][a-z]+$/)
+          word = (generate_any_word!) #.chomp #(context: generated_text)
+        end
+        @words.push(word.chomp)
+        word
+      else
+        generate_any_word!(context: context) #generated_text)
+      end    
+    end
+
+    private
+    def generate_any_word!(context: ["\n"])
       generated_chars = []
       char = ''
       until char.match(/[ .,;!?\n]/)
-	char = predict_next_state(context+generated_chars)
+        char = predict_next_state(context+generated_chars)
         generated_chars.push(char)
       end
       generated_chars.join
     end
   end