Initial commit

main.py

@@ -0,0 +1,60 @@
+from typing import Final
+from telegram import Update
+from telegram.ext import Application, MessageHandler, filters, ContextTypes
+from typing import Optional
+import random
+import os
+from dotenv import load_dotenv
+from models.seq2seq.model import Seq2SeqChatbot
+import torch
+
+
+load_dotenv()
+
+TOKEN: Final = os.environ.get("TOKEN")
+BOT_USERNAME: Final = os.environ.get("BOT_USERNAME")
+CHAT_ID: Final = int(os.environ.get("CHAT_ID"))
+
+CHECKPOINT_PATH: Final = "models/seq2seq/checkpoint/150_checkpoint.tar"
+
+torch.manual_seed(0)
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+chatbot = Seq2SeqChatbot(500, 2, 2, 0.1, device)
+chatbot.load_checkpoint(CHECKPOINT_PATH)
+chatbot.eval_mode()
+
+def handle_response(text: str) -> Optional[str]:
+    response_chance = 1.0
+    if random.random() < response_chance:
+        return chatbot(text)
+    return None
+
+
+async def handle_message(update: Update, context: ContextTypes.DEFAULT_TYPE):
+    if update.message.chat_id == CHAT_ID:
+        text: str = update.message.text.replace(BOT_USERNAME, '').strip().lower()
+        response: Optional[str] = handle_response(text)
+        if response:
+            await context.bot.sendMessage(update.message.chat_id, response)
+
+
+async def error(update: Update, context: ContextTypes.DEFAULT_TYPE):
+    print(f"{update.message.from_user.username} in {update.message.chat.type} "
+          f"chat caused error \"{context.error}\"\n"
+          f"{update}\"")
+
+def main() -> None:
+    """Run the bot."""
+    application = Application.builder().token(TOKEN).build()
+
+    application.add_handler(MessageHandler(filters.TEXT, handle_message))
+    application.add_error_handler(error)
+
+    application.run_polling(allowed_updates=Update.ALL_TYPES)
+
+
+if __name__ == '__main__':
+    print("Running main...")
+    # print(chatbot("test"))
+    main()

models/seq2seq/attention.py

@@ -0,0 +1,36 @@
+import torch
+import torch.nn.functional as F
+import torch.nn as nn
+from .custom_types import Method
+
+class LuongAttention(nn.Module):
+    def __init__(self, method: Method, hidden_size: int):
+        super().__init__()
+        self.hidden_size = hidden_size
+        if not isinstance(method, Method):
+            raise ValueError(method, f"should be a member of `Method` enum")
+        match method:
+            case Method.DOT:
+                self.method = self.dot
+            case Method.GENERAL:
+                self.method = self.general
+                self.Wa = nn.Linear(hidden_size, hidden_size)
+            case Method.CONCAT:
+                self.method = self.concat
+                self.Wa = nn.Linear(hidden_size * 2, hidden_size)
+                self.Va = nn.Parameter(torch.FloatTensor(1, hidden_size))
+
+    def dot(self, hidden, encoder_outputs):
+        return torch.sum(hidden * encoder_outputs, dim=2)
+
+    def general(self, hidden, encoder_outputs):
+        return torch.sum(hidden * self.Wa(encoder_outputs), dim=2)
+
+    def concat(self, hidden, encoder_outputs):
+        hidden = hidden.permute(1, 0, 2)
+        energy = self.Wa(torch.cat((hidden.permute(1, 0, 2).expand(-1, encoder_outputs.size(1), -1), encoder_outputs), 2)).tanh()
+        return torch.sum(self.Va * energy, dim=2)
+
+    def forward(self, hidden, encoder_outputs):
+        attn_weights = self.method(hidden, encoder_outputs)
+        return F.softmax(attn_weights, dim=1).unsqueeze(1)

models/seq2seq/chat_dataset.py

@@ -0,0 +1,97 @@
+import torch
+import torch.utils.data as data
+from typing import List, Union, Tuple
+from collections import OrderedDict
+from .vocab import Vocab
+from .custom_types import Message, MessageId, Conversation
+from torch.nn.utils.rnn import pad_sequence
+from .custom_types import Token
+import re
+import json
+
+class ChatDataset(data.Dataset):
+    def __init__(self, path: str, max_message_count: int = None, batch_size=5):
+        super().__init__()
+        self.path = path
+        self.batch_size = batch_size
+        self.messages: OrderedDict[MessageId, Message] = self.__load_messages_from_json(path, max_message_count)
+        self.conversations: List[Conversation] = ChatDataset.__conversations_from_messages(self.messages)
+        self.vocab = Vocab(list(self.messages.values())) # TODO: try changing this cast to something more applicable
+
+        self.batches_X, self.batches_y, self.lengths, self.mask = self.__batches_from_conversations()
+
+        self.length = len(self.batches_X)
+
+    def __batches_from_conversations(self) -> Tuple[List[torch.LongTensor], List[torch.LongTensor], List[torch.LongTensor], List[torch.BoolTensor]]: # Shape of tensor in batch: (batch_size, max_len_in_batch)
+        conversations = sorted(self.conversations, key=lambda x: len(x[0])) # Sort by input sequence length
+        batches_X: List[torch.LongTensor] = list()
+        batches_y: List[torch.LongTensor] = list()
+        lengths: List[torch.LongTensor] = list()
+        mask: List[torch.BoolTensor] = list()
+        for i in range(0, len(conversations), self.batch_size):
+            batches_X.append(pad_sequence([self.vocab.sentence_indices(conversations[i+j][0] + ["<eos>"]) for j in range(self.batch_size) if i+j < len(conversations)], batch_first=True, padding_value=0))
+            batches_y.append(pad_sequence([self.vocab.sentence_indices(conversations[i+j][1] + ["<eos>"]) for j in range(self.batch_size) if i+j < len(conversations)], batch_first=True, padding_value=0))
+            lengths.append(torch.tensor([len(conversations[i+j][0]) for j in range(self.batch_size) if i+j < len(conversations)]))
+            mask.append(batches_y[-1] != Token.PAD_TOKEN.value)
+        return batches_X, batches_y, lengths, mask
+
+    @classmethod
+    def __load_messages_from_json(cls, path: str, max_message_count: int = None) -> OrderedDict[MessageId, Message]:
+        messages: OrderedDict[MessageId, Message] = OrderedDict()
+        with open(path, "r", encoding="utf-8") as file:
+            chat_json = json.load(file)
+            for i, message in enumerate(chat_json["messages"]):
+                if max_message_count and i == max_message_count:
+                    break
+                if message["type"] != "message":
+                    continue
+                new_message = {
+                    "id": message["id"],
+                    "text": cls.__normalize(message["text"])
+                }
+                if not new_message["text"]: # Check for empty message
+                    continue
+                if "reply_to_message_id" in message.keys():
+                    new_message["reply_to_id"] = message["reply_to_message_id"]
+
+                messages[new_message["id"]] = new_message
+        return messages
+
+    @classmethod
+    def __conversations_from_messages(cls, messages: OrderedDict[MessageId, Message]) -> List[Conversation]:
+        # Search for message with `id` in the last `current_id` messages
+        def _get_message_by_id(current_id: int, id: int) -> Message:
+            for i in range(current_id - 1, -1, -1):
+                if messages[i]["id"] == id:
+                    return messages[i]
+            return None
+
+        conversations: List[Conversation] = []
+
+        messages_values = list(messages.values()) # TODO: try changing this cast to something more applicable
+        for i in range(len(messages) - 1): # There's no answer for last message so add -1
+            prev_message = messages_values[i]
+            if "reply_to_id" in messages_values[i].keys(): # Message is answer for message with `id` of `reply_to_id`
+                try:
+                    prev_message = messages[messages_values[i]["reply_to_id"]]
+                except KeyError:
+                    continue
+            conversations.append((prev_message["text"], messages_values[i+1]["text"]))
+        return conversations
+
+    @classmethod
+    def __normalize(cls, text: Union[str, List]) -> List[str]:
+        if isinstance(text, List):
+            text = " ".join([word for word in text if isinstance(word, str)])
+        text = text.lower().strip()
+        text = re.sub(r"([.!?])", r" \1 ", text)
+        text = re.sub(r"ё", r"е", text)
+        text = re.sub(r"[^а-яА-я.!?]+", r" ", text)
+        text = re.sub(r"\s+", r" ", text).strip()
+        return text.split()
+
+    def __getitem__(self, item):
+        return self.batches_X[item], self.batches_y[item], self.lengths[item], self.mask[item]
+
+    def __len__(self):
+        return self.length

models/seq2seq/custom_types.py

@@ -0,0 +1,23 @@
+from typing import List, TypedDict, NotRequired, Tuple
+from enum import Enum, auto
+
+MessageId = int
+MessageText = List[str]
+Conversation = Tuple[MessageText]
+
+
+class Message(TypedDict):
+    id: MessageId
+    text: MessageText
+    reply_to_id: NotRequired[int]
+
+class Method(Enum):
+    DOT = auto()
+    GENERAL = auto()
+    CONCAT = auto()
+
+class Token(Enum):
+    PAD_TOKEN = 0
+    BOS_TOKEN = 1
+    EOS_TOKEN = 2
+    UNK_TOKEN = 3

models/seq2seq/model.py

@@ -0,0 +1,208 @@
+import torch
+from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence
+import torch.nn.functional as F
+import torch.nn as nn
+import torch.optim as optim
+from .chat_dataset import ChatDataset
+from .attention import LuongAttention
+from .custom_types import Method, Token
+from .vocab import Vocab
+from .searchers import GreedySearch
+import os
+import random
+from tqdm import tqdm
+
+
+class Seq2SeqEncoder(nn.Module):
+    def __init__(self, input_size: int, hidden_size: int, num_layers: int, embedding: nn.Embedding):
+        super().__init__()
+        self.input_size = input_size
+        self.hidden_size = hidden_size
+        self.num_layers = num_layers
+        
+        self.embedding = embedding
+        self.rnn = nn.GRU(input_size, hidden_size, num_layers=num_layers, bidirectional=True, batch_first=True) # batch_first is True, because I don't approve self-harm
+
+    def forward(self, x, lengths):
+        x = self.embedding(x) # Output shape: (batch_size, max_len_in_batch, hidden_size)
+        packed_embedded = pack_padded_sequence(x, lengths.cpu(), batch_first=True, enforce_sorted=False)
+        outputs, hidden = self.rnn(packed_embedded)
+        outputs, _ = pad_packed_sequence(outputs, batch_first=True)
+        return outputs[:, :, :self.hidden_size] + outputs[:, :, self.hidden_size:], hidden
+
+
+class Seq2SeqDecoder(nn.Module):
+    def __init__(self, input_size: int, hidden_size: int, output_size: int, num_layers: int, attn, embedding: nn.Embedding, dropout: int = 0.1):
+        super().__init__()
+        self.input_size = input_size
+        self.hidden_size = hidden_size
+        self.output_size = output_size
+        self.num_layers = num_layers
+        
+        self.attn = attn
+        self.embedding = embedding
+        self.embedding_dropout = nn.Dropout(dropout)
+        self.rnn = nn.GRU(input_size, hidden_size, num_layers=num_layers, batch_first=True)
+
+        self.concat = nn.Linear(hidden_size * 2, hidden_size)
+        self.out = nn.Linear(hidden_size, output_size)
+
+    def forward(self, x, last_hidden, encoder_outputs):
+        embedded = self.embedding(x)
+        embedded = self.embedding_dropout(embedded)
+        decoder_outputs, hidden = self.rnn(embedded, last_hidden)
+        attn_weights = self.attn(decoder_outputs, encoder_outputs)
+
+        context = attn_weights.bmm(encoder_outputs).squeeze(1)
+
+        concat_input = torch.cat((decoder_outputs.squeeze(1), context), 1)
+        concat_output = torch.tanh(self.concat(concat_input))
+        output = self.out(concat_output)
+
+        output = F.softmax(output, dim=1)
+        return output, hidden
+
+
+class Seq2SeqChatbot(nn.Module):
+    def __init__(self, hidden_size: int, vocab_size: int, encoder_num_layers: int, decoder_num_layers: int, decoder_embedding_dropout: float, device: torch.device):
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.encoder_num_layers = encoder_num_layers
+        self.decoder_num_layers = decoder_num_layers
+        self.decoder_embedding_dropout = decoder_embedding_dropout
+        self.vocab_size = vocab_size
+        self.epoch = 0
+
+        self.device = device
+        self.vocab = Vocab([])
+        self.embedding = nn.Embedding(vocab_size, hidden_size)
+        self.attn = LuongAttention(Method.DOT, hidden_size)
+        self.encoder = Seq2SeqEncoder(hidden_size, hidden_size, encoder_num_layers, self.embedding)
+        self.decoder = Seq2SeqDecoder(hidden_size, hidden_size, vocab_size, decoder_num_layers, self.attn, self.embedding, decoder_embedding_dropout)
+        self.encoder_optimizer = optim.Adam(self.encoder.parameters())
+        self.decoder_optimizer = optim.Adam(self.decoder.parameters())
+        self.searcher = GreedySearch(self.encoder, self.decoder, self.embedding, device)
+        self.to(device)
+        self.eval_mode()
+
+    def train(self, epochs, train_data, teacher_forcing_ratio, device, save_dir, model_name, clip, save_every):
+        def maskNLLLoss(inp, target, mask):
+            crossEntropy = -torch.log(torch.gather(inp, 1, target.view(-1, 1)).squeeze(1))
+            loss = crossEntropy.masked_select(mask).mean()
+            loss = loss.to(device)
+            return loss
+        
+        epoch_progress = tqdm(range(self.epoch, self.epoch + epochs), desc="Training", unit="epoch", leave=True)
+        epoch_progress.set_description(f"maskNLLLoss: None")
+        
+        for epoch in epoch_progress:
+            for x_train, y_train, x_lengths, y_mask in train_data:
+                self.encoder_optimizer.zero_grad()
+                self.decoder_optimizer.zero_grad()
+                # Squeeze because batches are made in dataset and DataLoader is only for shuffling
+                x_train = x_train.squeeze(0).to(device)
+                y_train = y_train.squeeze(0).to(device)
+                x_lengths = x_lengths.squeeze(0) # Lengths are computed on CPU
+                y_mask = y_mask.squeeze(0).to(device)
+        
+                encoder_outputs, hidden = self.encoder(x_train, x_lengths) # Output shape: (batch_size, max_len_in_batch, hidden_size)
+                hidden = hidden[:self.decoder_num_layers]
+                loss = 0
+                decoder_input = torch.LongTensor([[Token.BOS_TOKEN.value] for _ in range(y_train.shape[0])])
+                decoder_input = decoder_input.to(device)
+                use_teacher_forcing = random.random() < teacher_forcing_ratio
+                if use_teacher_forcing:
+                    for t in range(y_train.shape[1]): # Process words in all batches for timestep t
+                        decoder_outputs, hidden = self.decoder(decoder_input, hidden, encoder_outputs)
+                        decoder_input = y_train[:, t].unsqueeze(1)
+                        mask_loss = maskNLLLoss(decoder_outputs, y_train[:, t], y_mask[:, t])
+                        loss += mask_loss
+                else:
+                    for t in range(y_train.shape[1]):
+                        decoder_outputs, hidden = self.decoder(decoder_input, hidden, encoder_outputs)
+                        decoder_input = torch.argmax(decoder_outputs, dim=1).unsqueeze(1)
+                        mask_loss = maskNLLLoss(decoder_outputs, y_train[:, t], y_mask[:, t])
+                        loss += mask_loss
+        
+                loss.backward()
+        
+                _ = nn.utils.clip_grad_norm_(self.encoder.parameters(), clip)
+                _ = nn.utils.clip_grad_norm_(self.decoder.parameters(), clip)
+        
+                self.encoder_optimizer.step()
+                self.decoder_optimizer.step()
+        
+            if (epoch % save_every == 0 and epoch != 0) or epoch == save_every - 1:
+                directory = os.path.join(save_dir, model_name, '{}-{}'.format(self.encoder_num_layers, self.decoder_num_layers, self.hidden_size))
+                if not os.path.exists(directory):
+                    os.makedirs(directory)
+                torch.save({
+                    'epoch': epoch + self.epoch,
+                    'en': self.encoder.state_dict(),
+                    'de': self.decoder.state_dict(),
+                    'en_opt': self.encoder_optimizer.state_dict(),
+                    'de_opt': self.decoder_optimizer.state_dict(),
+                    'loss': loss,
+                    'voc_dict': self.vocab.__dict__,
+                    'embedding': self.embedding.state_dict()
+                }, os.path.join(directory, '{}_{}.tar'.format(epoch, 'checkpoint')))
+            
+            epoch_progress.set_description(f"maskNLLLoss: {loss:.8f}")
+    
+    def to(self, device):
+        self.encoder = self.encoder.to(device)
+        self.decoder = self.decoder.to(device)
+        self.embedding = self.embedding.to(device)
+        self.attn = self.attn.to(device)
+
+    def train_mode(self):
+        self.encoder.train()
+        self.decoder.train()
+        self.embedding.train()
+        self.attn.train()
+
+    def eval_mode(self):
+        self.encoder.eval()
+        self.decoder.eval()
+        self.embedding.eval()
+        self.attn.eval()
+
+    def load_checkpoint(self, checkpoint_path: str):
+        checkpoint = torch.load(checkpoint_path, map_location=self.device, weights_only=False)
+        encoder_sd = checkpoint["en"]
+        decoder_sd = checkpoint["de"]
+        embedding_sd = checkpoint["embedding"]
+        self.vocab.__dict__ = checkpoint["voc_dict"]
+        encoder_optimizer_sd = checkpoint["en_opt"]
+        decoder_optimizer_sd = checkpoint["de_opt"]
+        self.epoch = checkpoint["epoch"]
+
+        self.encoder_optimizer.load_state_dict(encoder_optimizer_sd)
+        self.decoder_optimizer.load_state_dict(decoder_optimizer_sd)
+        self.embedding.load_state_dict(embedding_sd)
+        self.encoder.load_state_dict(encoder_sd)
+        self.decoder.load_state_dict(decoder_sd)
+
+    def forward(self, input_seq: str):
+        input_seq = ChatDataset._ChatDataset__normalize(input_seq)
+        input_seq = self.vocab.sentence_indices(input_seq + ["<eos>"]).unsqueeze(0).to(self.device)
+        output, _ = self.searcher(input_seq, torch.tensor(input_seq.shape[1]).view(1), 10)
+        output = [self.vocab.index2word[i.item()] for i in output]
+        output = [word for word in output if word not in ("<bos>", "<eos>", "<pad>")]
+        return " ".join(output)
+    
+
+if __name__ == "__main__": # Run as module
+    from .chat_dataset import ChatDataset
+    import torch.utils.data as data
+
+    CHAT_HISTORY_PATH = "models/seq2seq/data/train/chat_history.json"
+    batch_size = 20
+    chat_dataset = ChatDataset(CHAT_HISTORY_PATH, max_message_count=10_000, batch_size=batch_size)
+    train_data = data.DataLoader(chat_dataset, batch_size=1, shuffle=True)
+
+    device = torch.device("cpu")
+    chatbot = Seq2SeqChatbot(500, chat_dataset.vocab.size, 2, 2, 0.1, device)
+    chatbot.load_checkpoint("models/seq2seq/checkpoint/150_checkpoint.tar")
+    chatbot.train_mode()
+    chatbot.train(3, train_data, 0.5, device, "./checkpoint/temp/", "frantics_fox", 50.0, 100)

models/seq2seq/searchers.py

@@ -0,0 +1,27 @@
+import torch
+import torch.nn as nn
+from .custom_types import Token
+
+class GreedySearch(nn.Module):
+    def __init__(self, encoder, decoder, embedding, device):
+        super().__init__()
+        self.encoder = encoder
+        self.decoder = decoder
+        self.embedding = embedding
+        self.device = device
+
+    def forward(self, x, input_length, max_length):
+        encoder_outputs, hidden = self.encoder(x, input_length)
+        decoder_hidden = hidden[:self.decoder.num_layers]
+        decoder_input = torch.ones(1, 1, device=self.device, dtype=torch.long) * Token.BOS_TOKEN.value
+        all_tokens = torch.zeros([0], device=self.device, dtype=torch.long)
+        all_scores = torch.zeros([0], device=self.device)
+
+        for _ in range(max_length):
+            decoder_outputs, decoder_hidden = self.decoder(decoder_input, decoder_hidden, encoder_outputs)
+            decoder_scores, decoder_input = torch.max(decoder_outputs, dim=1)
+            all_tokens = torch.cat((all_tokens, decoder_input), dim=0)
+            all_scores = torch.cat((all_scores, decoder_scores), dim=0)
+            decoder_input.unsqueeze_(0)
+
+        return all_tokens, all_scores

models/seq2seq/vocab.py

@@ -0,0 +1,43 @@
+import torch
+import torch.nn as nn
+from typing import List, Dict
+from .custom_types import Token
+
+
+class Vocab(nn.Module):
+    def __init__(self, messages: List[Dict]):
+        super().__init__()
+        self.word2index: Dict[str, int] = {"<pad>": Token.PAD_TOKEN.value, "<bos>": Token.BOS_TOKEN.value, "<eos>": Token.EOS_TOKEN.value, "<unk>": Token.UNK_TOKEN.value}
+        self.index2word: Dict[int, str] = {Token.PAD_TOKEN.value: "<pad>", Token.BOS_TOKEN.value: "<bos>", Token.EOS_TOKEN.value: "<eos>", Token.UNK_TOKEN.value: "<unk>"}
+        self.word_count: Dict[str, int] = dict()
+        self.size = 4
+
+        for message in messages:
+            self.add_sentence(message["text"])
+
+        self.embedding = nn.Embedding(self.size, 300)
+
+    def add_sentence(self, sentence):
+        for word in sentence:
+            self.add_word(word)
+
+    def add_word(self, word):
+        if word not in self.word2index:
+            self.word2index[word] = self.size
+            self.index2word[self.size] = word
+            self.word_count[word] = 1
+            self.size += 1
+        else:
+            self.word_count[word] += 1
+
+    def sentence_indices(self, sentence: List[str]) -> torch.LongTensor:
+        indices = torch.LongTensor(len(sentence))
+        for i, word in enumerate(sentence):
+            indices[i] = self.word2index[word] if word in self.word2index else Token.UNK_TOKEN.value
+        return indices
+
+    def forward(self, indices: torch.LongTensor):
+        return self.embedding(indices)
+
+    def __len__(self):
+        return self.size