73 lines
2.5 KiB
Python
73 lines
2.5 KiB
Python
import random
|
|
# import re
|
|
import json
|
|
import os # listdir, path.join, path.isdir
|
|
|
|
# from collections import Counter
|
|
|
|
# punctuation to remove
|
|
RM_PUNCT = u"'\"”„+-—:;()"
|
|
|
|
|
|
# Markov random Text Generator
|
|
class MTG(object):
|
|
def __init__(self, data: dict):
|
|
self.MC = data or {}
|
|
|
|
def train_from_seq(self, word_seq):
|
|
for w1, w2, w3 in self.triples(word_seq):
|
|
key = w1 + '+' + w2
|
|
if key not in self.MC:
|
|
self.MC[key] = []
|
|
self.MC[key].append(w3)
|
|
|
|
def train_from_text(self, text, remove=RM_PUNCT):
|
|
text = text.replace('.', ' . ').replace(',', ' , ')
|
|
word_seq = text.strip().split()
|
|
word_seq = [word.strip(remove) for word in word_seq if len(word.strip(remove)) != 0]
|
|
self.train_from_seq(word_seq)
|
|
|
|
# add data form another MTG
|
|
def __iadd__(self, other):
|
|
for key, value in other.MC.items():
|
|
if key not in self.MC:
|
|
self.MC[key] = []
|
|
self.MC[key].extend(value)
|
|
|
|
def save(self):
|
|
return self.MC
|
|
|
|
# generate all triples to build MC
|
|
|
|
def triples(self, word_seq):
|
|
if len(word_seq) < 3:
|
|
return
|
|
for i in range(len(word_seq) - 2):
|
|
yield (word_seq[i], word_seq[i + 1], word_seq[i + 2])
|
|
|
|
def choose_start_words(self, first_word=None, second_word=None):
|
|
if second_word is not None:
|
|
return first_word, second_word
|
|
elif first_word is not None:
|
|
if len(self.MC['.' + '+' + first_word]) != 0:
|
|
second_word = random.choice(self.MC['.' + '+' + first_word])
|
|
elif len(self.MC['.' + '+' + first_word.capitalize()]) != 0:
|
|
second_word = random.choice(self.MC['.' + '+' + first_word.capitalize()])
|
|
else:
|
|
return self.choose_start_words()
|
|
return first_word, second_word
|
|
else:
|
|
first_word = random.choice(self.MC[random.choice(list(self.MC.keys()))])
|
|
return self.choose_start_words(first_word, second_word)
|
|
|
|
def generate_text(self, first_word=None, second_word=None, size=30):
|
|
w1, w2 = self.choose_start_words(first_word, second_word)
|
|
gen_words = [w1.capitalize()]
|
|
while not (len(gen_words) > size and w2 == '.'):
|
|
gen_words.append(w2.capitalize() if (w1 == '.') else w2)
|
|
w1, w2 = w2, random.choice(self.MC[w1 + '+' + w2])
|
|
gen_words.append(w2)
|
|
text = ' '.join(gen_words)
|
|
text = text.replace(' ,', ',').replace(' .', '.')
|
|
return text
|