{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# This Python 3 environment comes with many helpful analytics libraries installed\n# It is defined by the kaggle/python Docker image: https://github.com/kaggle/docker-python\n# For example, here's several helpful packages to load\n\nimport numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\n\n# Input data files are available in the read-only \"../input/\" directory\n# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory\n\nimport os\nfor dirname, _, filenames in os.walk('/kaggle/input'):\n    for filename in filenames:\n        print(os.path.join(dirname, filename))\n\n# You can write up to 20GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using \"Save & Run All\" \n# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import numpy as np\nimport json\nimport torch\nfrom torch import nn, Tensor\nimport torch.nn.functional as F\nfrom torch.utils.data import Dataset, DataLoader\nfrom tqdm import tqdm\nfrom torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence\nfrom typing import Optional\nimport matplotlib.pyplot as plt\n\nfrom transformers import AutoTokenizer, AutoModel, AutoModelForCausalLM","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Сделаем датасет\nВ этом задании мне удобнее задать путь к файлу, а внутри пусть будет чтение файла. К тому же пусть датасет отдает пары строк, токенизирую я их потом.","metadata":{}},{"cell_type":"code","source":"class PairsDataset(Dataset):\n    \n    def __init__(self, file_path, sep='\\t'):\n        super().__init__()\n        \n        self.file_path = file_path\n        self.sep = sep\n        \n        self.data = self.load_data()\n        \n    def load_data(self):\n        \n        data = list()\n        \n        with open(self.file_path) as file_object:\n            for line in file_object:\n                data.append(line.strip().split(self.sep))\n                \n        return data\n    \n    def __len__(self):\n        return len(self.data)\n    \n    def __getitem__(self, index):\n        \n        return self.data[index]","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Давайте сделаем токенизатор\nПусть у него будет методы `tokenize` для того, чтобы мы могли подать список текстов и получить нужную нам информацию (индексы слов, пад маску) и `collate` для того, чтобы обрабатывать батчи, эту функция нужна для `collate_fn` в `DataLoader`.","metadata":{}},{"cell_type":"code","source":"class Collater:\n    \n    def __init__(self, \n                 tokenizer_name='sberbank-ai/rugpt3small_based_on_gpt2', \n                 max_length=24):\n        \n        # будем использовать предобученный токенизатор\n        self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)\n        # устанавливаем пад токен, он будет иметь индекс \"длина словаря + 1\"\n        self.tokenizer.pad_token = '<pad>'\n        self.max_length = max_length\n        \n    def tokenize(self, texts):\n        tokenized = self.tokenizer(texts, padding=True, \n                                   truncation=True,\n                                   max_length=self.max_length, \n                                   return_tensors='pt')\n        # за пару минут я не нашел как вставить пад на нужное мне место, поэтому прибегнул к такому хаку\n        tokenized['input_ids'] = tokenized['input_ids'] * tokenized['attention_mask']\n        return tokenized\n    \n    def __call__(self, batch):\n        \n        questions, responses = list(), list()\n        \n        for question, response in batch:\n            questions.append(question)\n            responses.append(response)\n            \n        tokenized_questions = self.tokenize(questions)\n        tokenized_responses = self.tokenize(responses)\n        \n        return tokenized_questions, tokenized_responses","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Сделаем модель","metadata":{}},{"cell_type":"markdown","source":"### Модуль, который делает `max/mean` пулинг с учетом падов\nВедь усреднить\n[1, 2, 3, 0, 0]  \nНе все равно, что усреднить\n[1, 2, 3]  \nПады нам учитывать не нужно","metadata":{}},{"cell_type":"code","source":"def embedding_masking(x: Tensor,\n                      pad_mask: Tensor,\n                      value: float = 0.) -> Tensor:\n    x = x.masked_fill((~(pad_mask.bool())).unsqueeze(-1), value)\n    return x\n\n\nclass GlobalMaskedPooling(nn.Module):\n\n    POOLING_TYPES = ('mean', 'max')\n\n    def __init__(self,\n                 pooling_type: str = 'mean',\n                 dim: int = 1,\n                 normalize: bool = False,\n                 length_scaling: bool = True,\n                 scaling_square_root: bool = True):\n        super().__init__()\n\n        self.pooling_type = pooling_type\n        self.dim = dim\n\n        self.normalize = normalize\n        self.length_scaling = length_scaling\n        self.scaling_square_root = scaling_square_root\n\n        if self.pooling_type == 'max':\n            self.mask_value = -float('inf')\n        else:\n            self.mask_value = 0.\n\n        if self.pooling_type not in self.POOLING_TYPES:\n            raise ValueError(f'Available types: {\", \".join(self.POOLING_TYPES)}')\n\n    def forward(self, x: Tensor, pad_mask: Tensor) -> Tensor:\n        lengths = pad_mask.sum(self.dim).float()\n\n        x = embedding_masking(x=x, pad_mask=pad_mask, value=self.mask_value)\n\n        if self.pooling_type == 'mean':\n            scaling = x.size(self.dim) / lengths\n        else:\n            scaling = torch.ones(x.size(0))\n\n        if self.length_scaling:\n            lengths_factor = lengths\n            if self.scaling_square_root:\n                lengths_factor = lengths_factor ** 0.5\n            scaling /= lengths_factor\n\n        scaling = scaling.masked_fill(lengths == 0, 1.).unsqueeze(-1)\n\n        if self.pooling_type == 'mean':\n            x = x.mean(self.dim)\n        else:\n            x, _ = x.max(self.dim)\n\n        x *= scaling\n\n        if self.normalize:\n            x = F.normalize(x)\n\n        return x","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Модуль для модели\nВнутри у нас есть наш backbone - это lstm. Эмбеддинги будем использовать предобученные, от gpt (той версии, которую использовали для токенизатора).","metadata":{}},{"cell_type":"code","source":"class LSTMBiEncoder(nn.Module):\n    \n    def __init__(self, \n                 embedding_matrix, \n                 num_layers=3, \n                 model_dim=512, \n                 dropout=0.25, \n                 freeze_embeddings=True):\n        super().__init__()\n\n        self.embedding_layer = nn.Embedding.from_pretrained(embeddings=embedding_matrix, \n                                                            freeze=freeze_embeddings,\n                                                            padding_idx=0)\n        \n        self.rnn = nn.LSTM(input_size=embedding_matrix.size(1), \n                           hidden_size=model_dim, \n                           num_layers=num_layers,\n                           batch_first=True,\n                           dropout=dropout, \n                           bidirectional=True)\n        \n        self.rnn_projection = nn.Linear(model_dim * 2, model_dim)\n        \n        self.pooling = GlobalMaskedPooling()\n        \n        self.question_encoder = self.get_head(model_dim, dropout)\n        self.response_encoder = self.get_head(model_dim, dropout)\n        \n    def get_head(self, model_dim, dropout):\n        return nn.Sequential(\n            nn.BatchNorm1d(model_dim), \n            nn.Dropout(dropout),\n            nn.Linear(model_dim, model_dim),\n            nn.ReLU(),\n            nn.BatchNorm1d(model_dim), \n            nn.Dropout(dropout),\n            nn.Linear(model_dim, model_dim),\n            nn.ReLU(),\n            nn.BatchNorm1d(model_dim), \n            nn.Dropout(dropout),\n            nn.Linear(model_dim, model_dim)\n        )\n    \n    def backbone_encoding(self, batch):\n        \n        x = self.embedding_layer(batch['input_ids'])\n        \n        pad_mask = batch['attention_mask']\n        \n        x = pack_padded_sequence(x, \n                                 lengths=pad_mask.sum(dim=-1).cpu(),\n                                 batch_first=True, \n                                 enforce_sorted=False)\n        \n        x, _ = self.rnn(x)\n        \n        x, _ = pad_packed_sequence(x, batch_first=True)\n        \n        x = self.pooling(x, pad_mask)\n        \n        x = self.rnn_projection(x)\n        \n        return x\n    \n    def question_encoding(self, batch):\n        return self.question_encoder(self.backbone_encoding(batch))\n    \n    def response_encoding(self, batch):\n        return self.response_encoder(self.backbone_encoding(batch))\n        \n    def forward(self, question, response):\n        \n        question_embeddings = self.question_encoding(question)\n        response_embeddings = self.response_encoding(response)\n        \n        return question_embeddings, response_embeddings","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Сделаем лосс","metadata":{}},{"cell_type":"code","source":"class SoftmaxLoss(nn.Module):\n    \n    def __init__(self):\n        super().__init__()\n        \n        self.criterion = nn.CrossEntropyLoss()\n        \n    def forward(self, question_embeddings, response_embeddings):\n        \n        similarity_matrix = question_embeddings @ response_embeddings.T\n        \n        targets = torch.arange(similarity_matrix.size(0)).to(question_embeddings.device)\n        \n        loss = self.criterion(similarity_matrix, targets) + self.criterion(similarity_matrix.T, targets)\n        \n        return loss","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# функция, которая перекидывает список из тензоров на нужный девайс\ndef batch_to_device(batch, device):\n    for key, value in batch.items():\n        batch[key] = value.to(device)\n    return batch","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"pretrained_model = 'sberbank-ai/rugpt3small_based_on_gpt2'","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_dataset = PairsDataset('/kaggle/input/chit-chat-encoders/train.txt')\ncollater = Collater(tokenizer_name=pretrained_model)\n# да, у нас большой батч сайз\ntrain_loader = DataLoader(train_dataset, batch_size=512, shuffle=True, collate_fn=collater)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# от gpt нам потребуется только матрица эмбеддингов\ngpt_model = AutoModelForCausalLM.from_pretrained(pretrained_model)\nembedding_matrix = gpt_model.transformer.wte.weight.detach()\nembedding_matrix[0] = torch.zeros(embedding_matrix.size(1))\n\n# сделаем модель\nmodel = LSTMBiEncoder(embedding_matrix)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"device = torch.device('cuda')","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model.to(device)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# зададим оптимизатор\noptimizer = torch.optim.AdamW(model.parameters())","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# зададим лосс\ncriterion = SoftmaxLoss()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Обучение\nДобавьте еще валидацию, чтобы не переобучиться","metadata":{}},{"cell_type":"code","source":"n_epochs = 3","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"for n_epoch in range(1, n_epochs + 1):\n    \n    progress_bar = tqdm(total=len(train_loader), desc=f'Epoch {n_epoch}')\n    \n    losses = list()\n    \n    for n_iter, (question, response) in enumerate(train_loader):\n        \n        question = batch_to_device(question, device)\n        response = batch_to_device(response, device)\n        \n        question_embeddings, response_embeddings = model(question, response)\n        \n        loss = criterion(question_embeddings, response_embeddings)\n        \n        optimizer.zero_grad()\n        loss.backward()\n        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=3.)\n        optimizer.step()\n        \n        losses.append(loss.item())\n        \n        progress_bar.update()\n        progress_bar.set_postfix(loss=np.mean(losses[-100:]))\n        \n    torch.save(model.state_dict(), f'lstm_bi_encoder_state_dict_{n_epoch}.pth')\n        \n    progress_bar.close()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"plt.figure(figsize=(20, 10))\nplt.plot(losses)\nplt.title('Loss Values Dynamic')\nplt.xlabel('Step')\nplt.ylabel('Value');","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Сделаем удобнее инференс\nТакой класс, куда подаешь `collater` и `model` и по вызову этого класса можно векторизовать тексты.","metadata":{}},{"cell_type":"code","source":"class Encoder:\n    \n    def __init__(self, collater, model, batch_size=512, device=None):\n        super().__init__()\n        \n        self.collater = collater\n        self.model = model\n        \n        self.batch_size = batch_size\n        self.device = torch.device('cpu') if device is None else device\n        \n        self.model.to(self.device)\n        self.model.eval()\n        \n    def encoding(self, texts, is_questions=True):\n        \n        tokenized = self.collater.tokenize(texts)\n        \n        tokenized = batch_to_device(tokenized, device=self.device)\n        \n        if is_questions:\n            embeddings = self.model.question_encoding(tokenized)\n        else:\n            embeddings = self.model.response_encoding(tokenized)\n            \n        embeddings = F.normalize(embeddings).detach().cpu()\n            \n        return embeddings\n    \n    def __call__(self, texts, is_questions=True):\n        \n        if len(texts) <= self.batch_size:\n            return self.encoding(texts, is_questions=is_questions)\n        \n        embeddings = list()\n        \n        for batch_index in range(0, len(texts), self.batch_size):\n            batch = texts[batch_index:batch_index + self.batch_size]\n            embeddings.append(self.encoding(batch, is_questions=is_questions))\n            \n        embeddings = torch.cat(embeddings)\n            \n        return embeddings","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"encoder = Encoder(collater, model, batch_size=4096, device=device)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"data = list()\n\nwith open('/kaggle/input/chit-chat-encoders/valid.jsonl') as file_object:\n    for line in file_object:\n        data.append(json.loads(line.strip()))","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Соберем уникальные ответы\nЧтобы быстрее посчитать эмбеддинги","metadata":{}},{"cell_type":"code","source":"unique_responses = set()\n\nfor sample in data:\n    unique_responses.update(sample['candidates'])\n    \nunique_responses = list(unique_responses)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"unique_responses_mapper = {text: index for index, text in enumerate(unique_responses)}","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Векторизуем тексты","metadata":{}},{"cell_type":"code","source":"questions_embeddings = encoder([sample['question'] for sample in data])\nresponses_embeddings = encoder(unique_responses, is_questions=False)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Получим предсказания\nИндекс самого близкого ответа из кандидатов","metadata":{}},{"cell_type":"code","source":"prediction = list()\n\nfor i_sample in range(len(data)):\n    question_emb = questions_embeddings[i_sample].unsqueeze(0)\n    candidates_emb = torch.stack(\n        [responses_embeddings[unique_responses_mapper[data[i_sample]['candidates'][i_cand]]]\n         for i_cand in range(len(data[i_sample]['candidates']))])\n    similarity = question_emb @ candidates_emb.T\n    prediction.append(similarity.argmax().item())","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Соберем итоговые предсказания в файл","metadata":{}},{"cell_type":"code","source":"with open('submission.csv', 'w') as file_object:\n    file_object.write('Id,Category\\n')\n    for n, p in enumerate(prediction):\n        file_object.write(f'{n},{p}\\n')","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}