/
PoetryGenerator.py
37 lines (29 loc) · 1.01 KB
/
PoetryGenerator.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
from MarkovChain import MarkovChain
import re
class PoetryGenerator:
def __init__(self, k, n):
self.mc = MarkovChain()
self.k = k
self.n = n
def train(self, filename):
file = open(filename)
words = []
for line in file:
#words = words + re.compile("([\w]+|[.,!?;]+|[\n])").findall(line)
words = words + re.compile("([\w.,!?;']+|[\n])").findall(line)
print(words)
for i in range(len(words) - self.k - 1):
state = tuple(words[i:i+self.k])
next = words[i+self.k]
self.mc.add_state(state, next)
def generate(self):
state = list(self.mc.get_random_tuple())
for w in state:
print("{} ".format(w), end='')
for i in range(0, self.n):
gen_word = self.mc.get_state(tuple(state))
print("{} ".format(gen_word), end='')
state = self.shift(state, 1)
state[2] = gen_word
def shift(self, l, n):
return l[n:] + l[:n]