From 60e35b3e5ef38845c8b0931b10cfb64750fc8d07 Mon Sep 17 00:00:00 2001
From: accraze <accraze@gmail.com>
Date: Sun, 29 Jul 2018 08:36:25 -0700
Subject: [PATCH] allow custom hidden state size

---
 src/ia_markov/markov.py | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/src/ia_markov/markov.py b/src/ia_markov/markov.py
index 80b744f..537ee05 100644
--- a/src/ia_markov/markov.py
+++ b/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):