diff --git a/database.py b/database.py index 1f6ba46..fce0805 100755 --- a/database.py +++ b/database.py @@ -168,7 +168,7 @@ class DataBase: def get_random_message(self): sql = "SELECT text FROM xxx_message ORDER BY RANDOM() LIMIT 1" result = self.execute(sql) - return(result) + return(result[0][0]) def here(self, user_id, conf_id): sql = """ diff --git a/markov/simple.py b/markov/simple.py index 42f73e5..d0e89b8 100644 --- a/markov/simple.py +++ b/markov/simple.py @@ -38,7 +38,7 @@ class Dictogram(dict): if(index > random_int): return list_of_keys[i] -def get(): +def get(text): def generate_random_start(model): if 'END' in model: seed_word = 'END' @@ -52,10 +52,13 @@ def get(): current_word = generate_random_start(markov_model) sentence = [current_word] for i in range(0, length): - current_dictogram = markov_model[current_word] - random_weighted_word = current_dictogram.return_weighted_random_word() - current_word = random_weighted_word - sentence.append(current_word) + try: + current_dictogram = markov_model[current_word] + random_weighted_word = current_dictogram.return_weighted_random_word() + current_word = random_weighted_word + sentence.append(current_word) + except: + pass sentence[0] = sentence[0].capitalize() return ' '.join(sentence) + '.' return sentence @@ -70,14 +73,6 @@ def get(): markov_model[data[i]] = Dictogram([data[i+1]]) return markov_model - text = """ - Олег Соколов, преподававший в СПбГУ, в ноябре был задержан в Петербурге, - в его рюкзаке обнаружили две отпиленные женские руки. Соколов признался, - что убил и расчленил свою бывшую студентку Анастасию Ещенко, с которой - его связывали близкие отношения. Адвокат Соколова Александр Почуев заявлял, - что не исключает «и версию самооговора» его подзащитного и иные версии - преступления, «вплоть до мистических».""" - # simple cleanup text = text.replace('—','') text = text.replace('«','') @@ -90,6 +85,6 @@ def get(): text_list = text.split() model = make_markov_model(text_list) - generated = generate_random_sentence(50, model) + generated = generate_random_sentence(30, model) generated = generated.replace(' END', '.') - print(generated) + return generated diff --git a/worker.py b/worker.py index 53155c3..ee0896d 100755 --- a/worker.py +++ b/worker.py @@ -15,6 +15,7 @@ import random from string import punctuation from urllib.parse import urlencode from database import DataBase +from markov.simple import get from pygments import highlight from pygments.lexers import PythonLexer @@ -179,6 +180,33 @@ class MessageWorker: self.send(id=conf_id, msg='```\n' + self.db.scheme + '\n```') return True + if input_message[:7] == '/markov': + conf_id = msg['message']['chat']['id'] + user_id = msg['message']['from']['id'] + if msg['message']['chat']['type'] == 'private': + chat_title = conf_id + else: + chat_title = msg['message']['chat']['title'] + self.db.add_conf(conf_id, chat_title) + rand_messages = list() + max_sen = 100 + try: + count = int(msg['message']['text'][8:]) + if count > max_sen: + count = max_sen + except: + count = 30 + for i in range(0, count): + rand_messages.append(self.db.get_random_message()) + rand_text = " ".join(rand_messages) + gen_text = get(rand_text) + try: + gen_text = gen_text.lower().capitalize() + except: + pass + self.send(id=conf_id, msg=f'{gen_text}') + return True + if input_message == '/stat': conf_id = msg['message']['chat']['id'] user_id = msg['message']['from']['id']