From 39f7f118e503c6011c479ae45acdef32ba7371f5 Mon Sep 17 00:00:00 2001 From: Firepup Sixfifty Date: Sat, 11 May 2024 22:27:24 -0500 Subject: [PATCH] Oops. This should have been added earlier --- markov.py | 51 +++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 51 insertions(+) create mode 100644 markov.py diff --git a/markov.py b/markov.py new file mode 100644 index 0000000..756f6bd --- /dev/null +++ b/markov.py @@ -0,0 +1,51 @@ +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()