add markov to cyberlina

This commit is contained in:
bakatrouble 2019-11-27 22:23:57 +03:00
parent 6e256ae3a6
commit 61a8028ffe
2 changed files with 93 additions and 1 deletions

72
bots/modules/_mtg.py Normal file
View File

@ -0,0 +1,72 @@
import random
# import re
import json
import os # listdir, path.join, path.isdir
# from collections import Counter
# punctuation to remove
RM_PUNCT = u"'\"”„+-—:;()"
# Markov random Text Generator
class MTG(object):
def __init__(self, data: dict):
self.MC = data or {}
def train_from_seq(self, word_seq):
for w1, w2, w3 in self.triples(word_seq):
key = w1 + '+' + w2
if key not in self.MC:
self.MC[key] = []
self.MC[key].append(w3)
def train_from_text(self, text, remove=RM_PUNCT):
text = text.replace('.', ' . ').replace(',', ' , ')
word_seq = text.strip().split()
word_seq = [word.strip(remove) for word in word_seq if len(word.strip(remove)) != 0]
self.train_from_seq(word_seq)
# add data form another MTG
def __iadd__(self, other):
for key, value in other.MC.items():
if key not in self.MC:
self.MC[key] = []
self.MC[key].extend(value)
def save(self):
return self.MC
# generate all triples to build MC
def triples(self, word_seq):
if len(word_seq) < 3:
return
for i in range(len(word_seq) - 2):
yield (word_seq[i], word_seq[i + 1], word_seq[i + 2])
def choose_start_words(self, first_word=None, second_word=None):
if second_word is not None:
return first_word, second_word
elif first_word is not None:
if len(self.MC['.' + '+' + first_word]) != 0:
second_word = random.choice(self.MC['.' + '+' + first_word])
elif len(self.MC['.' + '+' + first_word.capitalize()]) != 0:
second_word = random.choice(self.MC['.' + '+' + first_word.capitalize()])
else:
return self.choose_start_words()
return first_word, second_word
else:
first_word = random.choice(self.MC[random.choice(self.MC.keys())])
return self.choose_start_words(first_word, second_word)
def generate_text(self, first_word=None, second_word=None, size=30):
w1, w2 = self.choose_start_words(first_word, second_word)
gen_words = [w1.capitalize()]
while not (len(gen_words) > size and w2 == '.'):
gen_words.append(w2.capitalize() if (w1 == '.') else w2)
w1, w2 = w2, random.choice(self.MC[w1 + '+' + w2])
gen_words.append(w2)
text = ' '.join(gen_words)
text = text.replace(' ,', ',').replace(' .', '.')
return text

View File

@ -14,6 +14,7 @@ from telegram.ext import Dispatcher, CallbackContext, MessageHandler, Filters, C
from telegram.utils.helpers import mention_html from telegram.utils.helpers import mention_html
from bots.models import TelegramBotModuleConfig from bots.models import TelegramBotModuleConfig
from bots.modules._mtg import MTG
class CyberLinaBotModuleConfig(TelegramBotModuleConfig): class CyberLinaBotModuleConfig(TelegramBotModuleConfig):
@ -24,6 +25,8 @@ class CyberLinaBotModuleConfig(TelegramBotModuleConfig):
already_ran = JSONField(default='[]') already_ran = JSONField(default='[]')
welcome_reactions = JSONField(default='[]') welcome_reactions = JSONField(default='[]')
inline_reactions = JSONField(default='[]') inline_reactions = JSONField(default='[]')
_mtg_data = JSONField(default={}, blank=True, null=True)
mtg_train = models.TextField(null=True, blank=True)
MODULE_NAME = 'Киберлиночка' MODULE_NAME = 'Киберлиночка'
CUSTOM_WIDGETS = { CUSTOM_WIDGETS = {
@ -34,8 +37,13 @@ class CyberLinaBotModuleConfig(TelegramBotModuleConfig):
'already_ran': JSONEditor(), 'already_ran': JSONEditor(),
'welcome_reactions': JSONEditor(), 'welcome_reactions': JSONEditor(),
'inline_reactions': JSONEditor(), 'inline_reactions': JSONEditor(),
'_mtg_data': JSONEditor(),
} }
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self.mtg = MTG(self._mtg_data)
def message_handler(self, update: Update, ctx: CallbackContext): def message_handler(self, update: Update, ctx: CallbackContext):
if not update.effective_chat or not update.effective_user: if not update.effective_chat or not update.effective_user:
return return
@ -89,7 +97,12 @@ class CyberLinaBotModuleConfig(TelegramBotModuleConfig):
id=uuid4(), id=uuid4(),
title='Не нажимай >_<', title='Не нажимай >_<',
input_message_content=InputTextMessageContent(choice(self.inline_reactions)) input_message_content=InputTextMessageContent(choice(self.inline_reactions))
) ),
InlineQueryResultArticle(
id=uuid4(),
title='ФлаБеПроЛейка',
input_message_content=InputTextMessageContent(choice(self.mtg.generate_text()))
),
] ]
update.inline_query.answer(results, cache_time=0) update.inline_query.answer(results, cache_time=0)
@ -99,6 +112,13 @@ class CyberLinaBotModuleConfig(TelegramBotModuleConfig):
dispatcher.add_handler(InlineQueryHandler(self.inline_query_handler)) dispatcher.add_handler(InlineQueryHandler(self.inline_query_handler))
return dispatcher return dispatcher
def save(self, force_insert=False, force_update=False, using=None, update_fields=None):
if self.mtg_train:
self.mtg.train_from_text(self.mtg_train)
self._mtg_data = self.mtg.save()
self.mtg_train = None
super().save(force_insert, force_update, using, update_fields)
class CyberLinaChat(models.Model): class CyberLinaChat(models.Model):
config = models.ForeignKey(CyberLinaBotModuleConfig, on_delete=models.CASCADE, related_name='chats') config = models.ForeignKey(CyberLinaBotModuleConfig, on_delete=models.CASCADE, related_name='chats')