imaginaryfriend/src/service/data_learner.py

20 lines
515 B
Python

from src.config import redis
class DataLearner:
def __init__(self, tokenizer):
self.redis = redis
self.tokenizer = tokenizer
def learn(self, message):
pipe = self.redis.instance().pipeline()
words = self.tokenizer.extract_words(message)
for trigram in self.tokenizer.split_to_trigrams(words):
key = self.tokenizer.to_key(message.chat_id, trigram[:-1])
last_word = trigram[-1]
pipe.sadd(key, last_word)
pipe.execute()