Applying patch

tlm/model.py

@@ -5,21 +5,20 @@ from .helpers import rolling_window, join
 
 class TinyLanguageModel:
     def __init__(self, n=2):
-        "Create a new model that looks at n words at a time."
+        "Create a new model that looks at a n-word context window."
         self.n = n
         self.model = {}
 
     def train(self, words):
         "Learn word patterns from a list of words."
         for sequence in rolling_window(words, self.n + 1):
-            pattern = tuple(sequence[:-1])
-            next_word = sequence[-1]
-            if pattern not in self.model:
-                self.model[pattern] = []
-            self.model[pattern].append(next_word)
+            context, next_word = tuple(sequence[:-1]), sequence[-1]
+            if context not in self.model:
+                self.model[context] = []
+            self.model[context].append(next_word)
                 
-    def generate(self, length, prompt=None, join_fn=None, step_callback=None):
-        "Create new words based on what the model learned."
+    def generate(self, length, prompt=None, join_fn=join, step_callback=None):
+        "Generate new text based on what the model learned."
         if not self.model:
             raise Exception("The model has not been trained")
         output = list(prompt or self.get_random_pattern())
@@ -32,10 +31,10 @@ class TinyLanguageModel:
             if step_callback:
                 step_callback(pattern, options, chosen)
             output.append(chosen)
-        return (join_fn or join)(output)
+        return join_fn(output)
 
     def get_random_pattern(self):
-        "Randomly chooses one of the observed patterns"
+        "Randomly chooses one of the observed contexts"
         return random.choice(list(self.model.keys()))
 
     def save(self, filepath):