Merge 0dae4b0 into 6968649

jsvine · Sep 26, 2018 · 7a45441 · 7a45441
2 parents 6968649 + 0dae4b0
commit 7a45441
Showing 1 changed file with 7 additions and 2 deletions.
diff --git a/markovify/chain.py b/markovify/chain.py
@@ -43,6 +43,7 @@ def __init__(self, corpus, state_size, model=None):
         self.state_size = state_size
         self.model = model or self.build(corpus, self.state_size)
         self.precompute_begin_state()
+        self.mem = {}
 
     def build(self, corpus, state_size):
         """
@@ -90,8 +91,12 @@ def move(self, state):
             choices = self.begin_choices
             cumdist = self.begin_cumdist
         else:
-            choices, weights = zip(*self.model[state].items())
-            cumdist = list(accumulate(weights))
+            try:
+                choices, cumdist = self.mem[state]
+            except:
+                choices, weights = zip(*self.model[state].items())
+                cumdist = list(accumulate(weights))
+                self.mem[state] = (choices, cumdist)
         r = random.random() * cumdist[-1]
         selection = choices[bisect.bisect(cumdist, r)]
         return selection