Merge pull request #14 from accraze/custom-state-size

Allow custom hidden state size

src/ia_markov/markov.py

@@ -15,10 +15,11 @@ class MarkovModel(object):
     on an Internet Archive text file.
     """
 
-    def __init__(self, exclude=None):
+    def __init__(self, exclude=None, state_size=2):
         self.archive_name = None
         self.exclude = exclude
         self.model = None
+        self.state_size = state_size
 
     def train_model(self, archive_name):
         """
@@ -78,7 +79,7 @@ def _create_markov(self, text):
         Assign markovify.Text
         as model
         """
-        self.model = markovify.Text(text)
+        self.model = markovify.Text(text, state_size=self.state_size)
 
 
 class POSMarkov(MarkovModel):
@@ -88,7 +89,7 @@ def _create_markov(self, text):
         Assign part of speech tagged markov
         to model.
         """
-        self.model = POSifiedText(text)
+        self.model = POSifiedText(text, state_size=self.state_size)
 
 
 class POSifiedText(markovify.Text):