1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
|
import random
import logging
from velascobot2.database import Database
from velascobot2.utils import GraphicType, Response, ChatData
logger = logging.getLogger(__name__)
def _rewrite(text: str) -> list[str]:
"""
Splits strings into lists of words delimited by whitespaces.
Extra whitespaces are inserted around newline characters so they are
included as their own Markov chain element, to avoid polluting with
very similar words that would only differ in having a whitespace
attached or not
"""
words = text.replace("\n", " \n ").split(" ")
return [word.strip(" ") for word in words if len(word.strip(" ")) > 0]
def _normalize_key(*words: str) -> str:
"""Returns a case insensitive n-gram key by joining them with whitespaces"""
return " ".join(word.casefold() for word in words)
def _triplets(wordlist: list[str]):
"""Generates triplets of words 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 Parrot:
def __init__(self, bot_name: str):
self.db = Database(bot_name)
def parse_text(self, chat_id: str, text: str) -> None:
logger.debug(f"{chat_id} - Processing text: {text}")
words = [w for w in _rewrite(text) if w]
words.append("") # end of message
words = ["", ""] + words # start of message
for key, word, count in self.parse_words(words):
self.db.add_words(chat_id, key, word, count)
def parse_words(self, words: list[str]):
word_counts: dict[tuple[str, str], int] = {}
for w1, w2, w3 in _triplets(words):
key = _normalize_key(w1, w2)
word_counts[(key, w3)] = word_counts.get((key, w3), 0) + 1
for (key, word), count in word_counts.items():
yield (key, word, count)
def parse_graphic(
self, chat_id: str, graph_type: GraphicType, graph_id: str
):
logger.debug(f"{chat_id} - Parsing graphic {graph_type}: {graph_id}")
assert graph_type in GraphicType
self.db.add_graphic(chat_id, graph_type, graph_id)
def give_text(self, chat_id: str) -> str:
w1, w2 = "", ""
w3 = "invalid"
words = []
while w3:
key = _normalize_key(w1, w2)
weighted_next_words = self.db.get_next_words(chat_id, key)
w3 = self.choose_word(weighted_next_words)
if w3 is None:
raise Exception(f"Got a nonexistent word. Generated text so far: {' '.join(words)}. Dropping text")
words.append(w3)
w1, w2 = w2, w3
return " ".join(words)
def give_graphic(self, chat_id: str) -> tuple[GraphicType, str]:
weighted_graphics = self.db.get_graphics(chat_id)
return self.choose_graphic(weighted_graphics)
def choose_word(self, weighted_words: list[tuple[str, int]]) -> str | None:
if len(weighted_words) == 0:
return None
words, weights = zip(*weighted_words)
return random.choices(words, weights=weights)[0]
def choose_graphic(
self, weighted_graphics: list[tuple[str, str, int]]
) -> tuple[GraphicType, str]:
graphics = [
(graph_type, graph_id)
for graph_type, graph_id, _ in weighted_graphics
]
weights = [count for _, _, count in weighted_graphics]
chosen_type, chosen_id = random.choices(graphics, weights=weights)[0]
return GraphicType(chosen_type), chosen_id
def choose_response_type(self, chat_id: str) -> Response:
graphics_count = self.db.count_graphics(chat_id)
words_count = self.db.count_words(chat_id, _normalize_key("", ""))
total = graphics_count + words_count
if total > 0 and random.random() < graphics_count / total:
return Response.GRAPHIC
else:
return Response.TEXT
def get_all_chats(self) -> list[ChatData]:
return self.db.get_all_chats()
def get_chat_data(self, chat_id: str) -> ChatData | None:
return self.db.get_chat_data(chat_id)
def ensure_chat_exists(
self, chat_id: str, chat_type: str, chat_name: str
) -> bool:
return self.db.ensure_chat_exists(chat_id, chat_type, chat_name)
def update_chat_name(self, chat_id: str, chat_name: str):
self.db.update_chat_name(chat_id, chat_name)
def update_chat_answer_chance(self, chat_id: str, answer_chance: float):
self.db.update_chat_answer_chance(chat_id, answer_chance)
def update_chat_message_period(self, chat_id: str, period: int):
self.db.update_chat_message_period(chat_id, period)
def update_chat_set_restricted(self, chat_id: str, restricted: bool):
self.db.update_chat_set_restricted(chat_id, restricted)
def update_chat_set_silenced(self, chat_id: str, silenced: bool):
self.db.update_chat_set_silenced(chat_id, silenced)
def get_processed_messages(self, chat_id: str) -> int:
graph_count = self.db.get_chat_graph_count(chat_id)
start_words = self.db.get_next_words(chat_id, " ")
start_word_count = sum(weight for _, weight in start_words)
return graph_count + start_word_count
|
