diff --git a/modules/markov.py b/modules/markov.py new file mode 100644 index 00000000..3ecfc628 --- /dev/null +++ b/modules/markov.py @@ -0,0 +1,75 @@ +import random +from src import ModuleManager, utils + +class Module(ModuleManager.BaseModule): + def _on_load(self): + if not self.database.has_table("markov"): + self.database.execute("""CREATE TABLE markov + (channel_id INTEGER, first_word TEXT, second_word TEXT, + third_word TEXT, frequency INT, + FOREIGN KEY (channel_id) REFERENCES channels(channel_id), + PRIMARY KEY (channel_id, first_word, second_word)""") + + @utils.hook("received.message.channel") + def channel_message(self, event): + words = event["message_split"] + words_n = len(words) + if words_n > 2 and event["channel"].get_setting("markov", False): + + inserts = [] + inserts.append([None, None, words[0]]) + inserts.append([None, words[0], words[1]]) + + for i in range(words_n-2): + inserts.append(words[i:i+3]) + + inserts.append([words[-2], words[-1], None]) + inserts.append([words[-1], None, None]) + + for insert in inserts: + frequency = self.database.execute("""SELECT frequency + FROM markov WHERE channel_id=? AND first_word=? + AND second_word=? AND third_word=?""", + [event["channel"].id]+insert) + frequency = (frequency or [0])[0]+1 + + self.database.execute( + "INSERT OR REPLACE INTO markov VALUES (?, ?, ?, ?, ?)", + [event["channel"].id]+insert+[frequency]) + + print(self.generate(event["channel"])) + + def _choose(self, words): + words, frequencies = list(zip(*words)) + return random.choices(words, weights=frequencies, k=1)[0] + + def generate(self, channel_id): + first_words = self.database.execute("""SELECT third_word, frequency + FROM markov WHERE channel_id=? AND first_word IS NULL AND + second_word IS NULL AND third_word NOT NULL""", + [channel_id]) + if not first_words: + return None + first_word = self._choose(first_words) + + second_words = self.database.execute("""SELECT third_word, frequency + FROM markov WHERE channel_id=? AND first_word IS NULL AND + second_word=? AND third_word NOT NULL""", + [channel_id, first_word]) + if not second_words: + return None + second_word = self._choose(second_words) + + words = [first_word, second_word] + for i in range(30): + two_words = words[-2:] + third_words = self.database.execute("""SELECT third_word, frequency + FROM markov WHERE channel_id=? AND first_word=? AND + second_word=?""", [channel_id]+two_words) + + third_word = self._choose(third_words) + if third_word == None: + break + words.append(third_word) + + return " ".join(words)