imaginaryfriend/src/service/data_learner.py

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()