forked from Firepup650/FireBot
52 lines
1.6 KiB
Python
52 lines
1.6 KiB
Python
|
import random
|
||
|
|
from typing import Union
|
||
|
|
|
||
|
|
|
||
|
|
class MarkovBot:
|
||
|
|
def __init__(self, text: list[list[str]]) -> None:
|
||
|
|
self.text = text
|
||
|
|
self.chains = {}
|
||
|
|
self.__build_chains()
|
||
|
|
|
||
|
|
def __build_chains(self) -> None:
|
||
|
|
for i in range(len(self.text)):
|
||
|
|
text = self.text[i]
|
||
|
|
for j in range(len(text) - 1):
|
||
|
|
current_word = text[j]
|
||
|
|
next_word = text[j + 1]
|
||
|
|
|
||
|
|
if current_word not in self.chains:
|
||
|
|
self.chains[current_word] = {}
|
||
|
|
|
||
|
|
if next_word not in self.chains[current_word]:
|
||
|
|
self.chains[current_word][next_word] = 0
|
||
|
|
|
||
|
|
self.chains[current_word][next_word] += 1
|
||
|
|
|
||
|
|
def generate_text(self, word: Union[str, None] = None) -> str:
|
||
|
|
if not word:
|
||
|
|
current_word = random.choice(list(self.chains.keys()))
|
||
|
|
else:
|
||
|
|
current_word = word
|
||
|
|
generated_text = current_word
|
||
|
|
|
||
|
|
while current_word in self.chains:
|
||
|
|
next_word = random.choices(
|
||
|
|
list(self.chains[current_word].keys()),
|
||
|
|
weights=list(self.chains[current_word].values()),
|
||
|
|
)[0]
|
||
|
|
generated_text += " " + next_word
|
||
|
|
current_word = next_word
|
||
|
|
|
||
|
|
return generated_text
|
||
|
|
|
||
|
|
def generate_from_sentence(self, msg: Union[str, None] = None) -> str:
|
||
|
|
if not msg:
|
||
|
|
word = random.choice(list(self.chains.keys()))
|
||
|
|
else:
|
||
|
|
word = random.choice(msg.split())
|
||
|
|
if (for_word := self.generate_text(word)) != word:
|
||
|
|
return for_word
|
||
|
|
else:
|
||
|
|
return self.generate_text()
|