Add lots of features

tlm/model.py

@@ -18,7 +18,7 @@ class TinyLanguageModel:
                 self.model[pattern] = []
             self.model[pattern].append(next_word)
                 
-    def generate(self, length, prompt=None):
+    def generate(self, length, prompt=None, join_fn=None, step_callback=None):
         "Create new words based on what the model learned."
         if not self.model:
             raise Exception("The model has not been trained")
@@ -27,9 +27,12 @@ class TinyLanguageModel:
             pattern = tuple(output[-self.n:])
             if pattern not in self.model:
                 break
-            next_word = random.choice(self.model[pattern])
-            output.append(next_word)
-        return join(output)
+            options = self.model[pattern]
+            chosen = random.choice(options)
+            if step_callback:
+                step_callback(pattern, options, chosen)
+            output.append(chosen)
+        return (join_fn or join)(output)
 
     def get_random_pattern(self):
         "Randomly chooses one of the observed patterns"