mirror of
https://gitlab.com/vylion/velascobot.git
synced 2025-04-19 21:46:35 +02:00
7de67c0238
- Added ability to change Velasco's probability to answer replies to his messages, per chat - Rewritten save-to-file functions for slightly more readability, and rewritten load-from-file functions accordingly
105 lines
2.8 KiB
Python
105 lines
2.8 KiB
Python
#!/usr/bin/env python3
|
|
|
|
import random
|
|
import json
|
|
|
|
HEAD = "\n!kvl"
|
|
TAIL = "!kvl"
|
|
|
|
def trim_and_split(text):
|
|
words = text.replace('\n', '\n ').split(' ')
|
|
i = 0
|
|
while i < len(words):
|
|
w = words[i].strip(' \t')
|
|
if len(w) > 0:
|
|
words[i] = w
|
|
else:
|
|
del words[i]
|
|
i -= 1
|
|
i += 1
|
|
return words
|
|
|
|
def getkey(w1, w2):
|
|
key = (w1.strip().casefold(), w2.strip().casefold())
|
|
return str(key)
|
|
|
|
def getwords(key):
|
|
words = key.strip('()').split(', ')
|
|
for i in range(len(words)):
|
|
words[i].strip('\'')
|
|
return words
|
|
|
|
def triples(wordlist):
|
|
""" Generates triples from the given data string. So if our string were
|
|
"What a lovely day", we'd generate (What, a, lovely) and then
|
|
(a, lovely, day).
|
|
"""
|
|
|
|
if len(wordlist) < 3:
|
|
return
|
|
|
|
for i in range(len(wordlist) - 2):
|
|
yield (wordlist[i], wordlist[i+1], wordlist[i+2])
|
|
|
|
class Markov(object):
|
|
def __init__(self, text=None, from_json=False):
|
|
self.cache = {}
|
|
if not from_json:
|
|
if text is not None:
|
|
for line in text:
|
|
self.add_text(line)
|
|
else:
|
|
self.cache = json.loads(text)
|
|
|
|
def to_json(self):
|
|
return json.dumps(self.cache)
|
|
|
|
def from_json(string):
|
|
return Markov(string, True)
|
|
|
|
def add_text(self, text):
|
|
words = [HEAD]
|
|
words.extend(trim_and_split(text))
|
|
self.database(words)
|
|
|
|
def database(self, wordlist):
|
|
for w1, w2, w3 in triples(wordlist):
|
|
if w1 == HEAD:
|
|
if w1 in self.cache:
|
|
self.cache[HEAD].append(w2)
|
|
else:
|
|
self.cache[HEAD] = [w2]
|
|
key = getkey(w1, w2)
|
|
if key in self.cache:
|
|
self.cache[key].append(w3)
|
|
else:
|
|
self.cache[key] = [w3]
|
|
|
|
def generate_markov_text(self, size=50):
|
|
w1 = random.choice(self.cache[HEAD])
|
|
w2 = random.choice(self.cache[getkey(HEAD, w1)])
|
|
gen_words = []
|
|
for i in range(size):
|
|
gen_words.append(w1)
|
|
if w2 == TAIL or not getkey(w1, w2) in self.cache:
|
|
# print("Generated text")
|
|
break
|
|
else:
|
|
w1, w2 = w2, random.choice(self.cache[getkey(w1, w2)])
|
|
return ' '.join(gen_words)
|
|
|
|
def fuse_with(self, gen):
|
|
d = gen.cache
|
|
for key in gen.cache:
|
|
if key in self.cache:
|
|
self.cache[key].extend(d[key])
|
|
else:
|
|
self.cache[key] = list(d[key])
|
|
|
|
def new_count(self):
|
|
count = 0
|
|
for key in self.cache:
|
|
for word in self.cache[key]:
|
|
if word == "!kvl":
|
|
count += 1
|
|
return count
|