{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "available-trauma",
   "metadata": {
    "id": "following-action"
   },
   "source": [
    "## Описание задачи"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "express-warrant",
   "metadata": {
    "id": "vocal-corner"
   },
   "source": [
    "Задача: Сгенерировать youtube-комментарии по _ссылке_ на видео\n",
    "\n",
    "Всё просто, юзер постит ссылку на видео - вы его комментируете. Можно заранее обусловиться что видео только на английском или на русском. Нужно сочинить _несколько_ комментариев. Kudos если вместе с основным комментарием вы порождаете юзернеймы и-или ответы на него.\n",
    "\n",
    "\n",
    "Датасет для файнтюна можно [взять с kaggle](https://www.kaggle.com/tanmay111/youtube-comments-sentiment-analysis/data?select=UScomments.csv) или [собрать самостоятельно](https://towardsdatascience.com/how-to-build-your-own-dataset-of-youtube-comments-39a1e57aade).\n",
    "\n",
    "В качестве основной модели можно использовать [GPT-2 large](https://huggingface.co/gpt2-large). Вот как её файнтюнить: https://tinyurl.com/gpt2-finetune-colab. \n",
    "\n",
    "Если хотите больше - можно взять что-то из творчества https://huggingface.co/EleutherAI. Например, вот [тут](https://tinyurl.com/gpt-j-8bit) есть пример как файнтюнить GPT-J-6B (в 8 раз больше gpt2-large). Однако, этим стоит заниматься уже после того, как у вас заработал базовый сценарий с GPT2-large или даже base.\n",
    "\n",
    "В итоговом сервисе можно дать пользователю вариировать параметры генерации: \n",
    "- температура или top-p, если сэмплинг; \n",
    "- beam size и length penalty, если beam search; \n",
    "- сколько комментариев сгенерировать, etc. \n",
    "\n",
    "Отдельный респект если ваш код будет выводить комментарий по одному слову, прямо в процессе генерёжки - чтобы пользователь не ждал пока вы настругаете абзац целиком.\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "deadly-sensitivity",
   "metadata": {
    "id": "RvtIWT89zQWX"
   },
   "source": [
    "# Подготовка датасета"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "unnecessary-shame",
   "metadata": {
    "id": "final-liquid",
    "outputId": "07c50913-ce77-47c2-ec42-729016b020b9"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Archive:  USvideos.csv.zip\n",
      "\n",
      "  inflating: USvideos.csv            \n"
     ]
    }
   ],
   "source": [
    "# !unzip UScomments.csv.zip\n",
    "# !unzip USvideos.csv.zip\n",
    "\n",
    "# !pip install datasets transformers\n",
    "# !pip install bitsandbytes\n",
    "# !pip install nltk\n",
    "# !pip install langdetect"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 118,
   "id": "clinical-beauty",
   "metadata": {
    "id": "excessive-ownership"
   },
   "outputs": [],
   "source": [
    "# import nltk\n",
    "# nltk.download('stopwords')\n",
    "# nltk.download('wordnet')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "attended-friend",
   "metadata": {
    "id": "needed-bikini"
   },
   "outputs": [],
   "source": [
    "from tqdm import tqdm\n",
    "import pandas as pd\n",
    "import string\n",
    "import re\n",
    "\n",
    "import datasets\n",
    "import transformers\n",
    "from transformers import pipeline, set_seed\n",
    "\n",
    "from nltk.stem import WordNetLemmatizer\n",
    "from nltk.corpus import stopwords\n",
    "\n",
    "import torch\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "import bitsandbytes as bnb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "id": "decreased-processing",
   "metadata": {
    "id": "fifteen-instrument"
   },
   "outputs": [],
   "source": [
    "# generator = pipeline('text-generation', model='gpt2-large')\n",
    "# set_seed(42)\n",
    "# generator(\"The man worked as a\", max_length=10, num_return_sequences=5)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "monthly-growing",
   "metadata": {
    "id": "disturbed-banner"
   },
   "source": [
    "### Готовим комментарии"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 190,
   "id": "classified-olive",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "bronze-distributor",
    "outputId": "20c0252b-28c1-4365-e813-b7df3df010de"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "<ipython-input-190-7f1bcf368402>:1: DtypeWarning: Columns (2,3) have mixed types. Specify dtype option on import or set low_memory=False.\n",
      "\n",
      "  df = pd.read_csv(\"UScomments.csv\", on_bad_lines='skip')\n",
      "\n",
      "<ipython-input-190-7f1bcf368402>:4: SettingWithCopyWarning: \n",
      "\n",
      "A value is trying to be set on a copy of a slice from a DataFrame.\n",
      "\n",
      "Try using .loc[row_indexer,col_indexer] = value instead\n",
      "\n",
      "\n",
      "\n",
      "See the caveats in the documentation: https://pandas.pydata.org/pandas-docs/stable/user_guide/indexing.html#returning-a-view-versus-a-copy\n",
      "\n",
      "  df['likes'] = df['likes'].astype(int)\n",
      "\n",
      "<ipython-input-190-7f1bcf368402>:5: SettingWithCopyWarning: \n",
      "\n",
      "A value is trying to be set on a copy of a slice from a DataFrame.\n",
      "\n",
      "Try using .loc[row_indexer,col_indexer] = value instead\n",
      "\n",
      "\n",
      "\n",
      "See the caveats in the documentation: https://pandas.pydata.org/pandas-docs/stable/user_guide/indexing.html#returning-a-view-versus-a-copy\n",
      "\n",
      "  df['replies'] = df['replies'].astype(int)\n"
     ]
    }
   ],
   "source": [
    "df = pd.read_csv(\"UScomments.csv\", on_bad_lines='skip')\n",
    "df.video_id = df.video_id.astype(\"str\")\n",
    "df = df[df['likes'] != \"likes\"]\n",
    "df['likes'] = df['likes'].astype(int)\n",
    "df['replies'] = df['replies'].astype(int)\n",
    "df.dropna(inplace=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 191,
   "id": "distant-solid",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "cognitive-registration",
    "outputId": "aede1178-b0c5-42d4-9a43-55eb06250d12"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 6/6 [00:24<00:00,  4.02s/it]\n"
     ]
    }
   ],
   "source": [
    "df['cleaned_txt'] = df['comment_text'].str.replace('#+', '#', regex=True).str.strip()\n",
    "normal_puncs = [\"!\", \"\\.\", \"?\", \",\", \":\", \";\"]\n",
    "for char in tqdm([\"!\", \"\\.\", \"?\", \",\", \":\", \";\"]):\n",
    "  char1 = char if char != \"\\.\" else \".\"\n",
    "  df['cleaned_txt'] = df['cleaned_txt'].str.replace(f'[\\s]*[{char}]+', f'{char1}', regex=True)\n",
    "\n",
    "df['cleaned_txt'] = df['cleaned_txt'].str.replace(\"[^a-zA-Z#'!?,\\.:;]\", \" \", regex=True)\n",
    "df['cleaned_txt'] = df['cleaned_txt'].str.replace(' +', ' ', regex=True)\n",
    "remove_short = lambda x: ' '.join([w for w in x.split() if len(w) > 1 or w in [\"I\", \"a\", \"A\", \"u\", \"U\"]])\n",
    "df['cleaned_txt'] = df['cleaned_txt'].apply(remove_short)\n",
    "df = df[~df['cleaned_txt'].str.startswith(\"https\")]\n",
    "df['cleaned_txt'] = df['cleaned_txt'].apply(lambda x: ' '.join([w for w in x.split() if len(w) < 30]))\n",
    "df = df[df['cleaned_txt'].str.len() > 1]\n",
    "\n",
    "df['num_of_words'] = df['cleaned_txt'].str.split().apply(lambda x: len(x))\n",
    "df = df[df.num_of_words < 50]\n",
    "df = df[df.num_of_words > 3].iloc[:, :-1]\n",
    "################################################################################################################\n",
    "# df['cleaned_txt'] = df['cleaned_txt'].apply(lambda x:x.lower())\n",
    "\n",
    "# wnl = WordNetLemmatizer()\n",
    "\n",
    "# tokenized_tweet = df['cleaned_txt'].apply(lambda x: x.split())\n",
    "# tokenized_tweet.apply(lambda x: [wnl.lemmatize(i) for i in x if i not in set(stopwords.words('english'))]) \n",
    "# for i in range(len(tokenized_tweet)):\n",
    "#     tokenized_tweet[i] = ' '.join(tokenized_tweet[i])\n",
    "# df['cleaned_txt'] = tokenized_tweet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 165,
   "id": "looking-broadcast",
   "metadata": {
    "id": "according-midwest"
   },
   "outputs": [],
   "source": [
    "# from langdetect import detect, DetectorFactory\n",
    "\n",
    "# DetectorFactory.seed = 0\n",
    "\n",
    "# df_lang = df[df.groupby('video_id').cumcount().isin([0,1,2])]\n",
    "# texts = df_lang.cleaned_txt\n",
    "# results = []\n",
    "# for text in tqdm(texts):\n",
    "#     try:\n",
    "#         results.append(detect(text))\n",
    "#     except Exception:\n",
    "#         break\n",
    "# df_lang['lang'] = results"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "secondary-floor",
   "metadata": {
    "id": "incorporated-count"
   },
   "source": [
    "### Приклеиваем инфо о видео"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 193,
   "id": "reserved-ultimate",
   "metadata": {
    "id": "checked-boxing"
   },
   "outputs": [],
   "source": [
    "videos = pd.read_csv(\"USvideos.csv\", on_bad_lines='skip')\n",
    "videos.video_id = videos.video_id.astype(str)\n",
    "videos.date = videos.date.astype(str).str.replace(\".1\", \".10\", regex=False)\n",
    "videos.date = pd.to_datetime(videos.date, format=\"%d.%m\")\n",
    "videos = videos.sort_values(by=[\"video_id\", \"date\"], ascending=[1, 0])\n",
    "videos = videos[videos.groupby(\"video_id\").cumcount(\"date\") == 0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 194,
   "id": "specialized-asset",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 206
    },
    "id": "numeric-groove",
    "outputId": "3c077cd6-1551-4b5d-f958-99711239d600"
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "  <div id=\"df-df806cae-3387-4f03-bdae-738a45dd3e06\">\n",
       "    <div class=\"colab-df-container\">\n",
       "      <div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>video_id</th>\n",
       "      <th>title</th>\n",
       "      <th>channel_title</th>\n",
       "      <th>comment_text</th>\n",
       "      <th>likes</th>\n",
       "      <th>replies</th>\n",
       "      <th>cleaned_txt</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>--JinobXWPk</td>\n",
       "      <td>DANGEROUS Jungle Spider!</td>\n",
       "      <td>Brave Wilderness</td>\n",
       "      <td>I saw this wandering spider in our bathroom se...</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>I saw this wandering spider in our bathroom se...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>--JinobXWPk</td>\n",
       "      <td>DANGEROUS Jungle Spider!</td>\n",
       "      <td>Brave Wilderness</td>\n",
       "      <td>Can't you just stick to small ants, and bees? ...</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>Can't you just stick to small ants, and bees? ...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>--JinobXWPk</td>\n",
       "      <td>DANGEROUS Jungle Spider!</td>\n",
       "      <td>Brave Wilderness</td>\n",
       "      <td>Brazilian wandering spider is the deadliest sp...</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>Brazilian wandering spider is the deadliest sp...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>--JinobXWPk</td>\n",
       "      <td>DANGEROUS Jungle Spider!</td>\n",
       "      <td>Brave Wilderness</td>\n",
       "      <td>Nothing a can of hairspray and lighter couldn'...</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>Nothing a can of hairspray and lighter couldn'...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>--JinobXWPk</td>\n",
       "      <td>DANGEROUS Jungle Spider!</td>\n",
       "      <td>Brave Wilderness</td>\n",
       "      <td>Hey Coyote! Can you do an episode on the Japan...</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>Hey Coyote! Can you do an episode on the Japan...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>\n",
       "      <button class=\"colab-df-convert\" onclick=\"convertToInteractive('df-df806cae-3387-4f03-bdae-738a45dd3e06')\"\n",
       "              title=\"Convert this dataframe to an interactive table.\"\n",
       "              style=\"display:none;\">\n",
       "        \n",
       "  <svg xmlns=\"http://www.w3.org/2000/svg\" height=\"24px\"viewBox=\"0 0 24 24\"\n",
       "       width=\"24px\">\n",
       "    <path d=\"M0 0h24v24H0V0z\" fill=\"none\"/>\n",
       "    <path d=\"M18.56 5.44l.94 2.06.94-2.06 2.06-.94-2.06-.94-.94-2.06-.94 2.06-2.06.94zm-11 1L8.5 8.5l.94-2.06 2.06-.94-2.06-.94L8.5 2.5l-.94 2.06-2.06.94zm10 10l.94 2.06.94-2.06 2.06-.94-2.06-.94-.94-2.06-.94 2.06-2.06.94z\"/><path d=\"M17.41 7.96l-1.37-1.37c-.4-.4-.92-.59-1.43-.59-.52 0-1.04.2-1.43.59L10.3 9.45l-7.72 7.72c-.78.78-.78 2.05 0 2.83L4 21.41c.39.39.9.59 1.41.59.51 0 1.02-.2 1.41-.59l7.78-7.78 2.81-2.81c.8-.78.8-2.07 0-2.86zM5.41 20L4 18.59l7.72-7.72 1.47 1.35L5.41 20z\"/>\n",
       "  </svg>\n",
       "      </button>\n",
       "      \n",
       "  <style>\n",
       "    .colab-df-container {\n",
       "      display:flex;\n",
       "      flex-wrap:wrap;\n",
       "      gap: 12px;\n",
       "    }\n",
       "\n",
       "    .colab-df-convert {\n",
       "      background-color: #E8F0FE;\n",
       "      border: none;\n",
       "      border-radius: 50%;\n",
       "      cursor: pointer;\n",
       "      display: none;\n",
       "      fill: #1967D2;\n",
       "      height: 32px;\n",
       "      padding: 0 0 0 0;\n",
       "      width: 32px;\n",
       "    }\n",
       "\n",
       "    .colab-df-convert:hover {\n",
       "      background-color: #E2EBFA;\n",
       "      box-shadow: 0px 1px 2px rgba(60, 64, 67, 0.3), 0px 1px 3px 1px rgba(60, 64, 67, 0.15);\n",
       "      fill: #174EA6;\n",
       "    }\n",
       "\n",
       "    [theme=dark] .colab-df-convert {\n",
       "      background-color: #3B4455;\n",
       "      fill: #D2E3FC;\n",
       "    }\n",
       "\n",
       "    [theme=dark] .colab-df-convert:hover {\n",
       "      background-color: #434B5C;\n",
       "      box-shadow: 0px 1px 3px 1px rgba(0, 0, 0, 0.15);\n",
       "      filter: drop-shadow(0px 1px 2px rgba(0, 0, 0, 0.3));\n",
       "      fill: #FFFFFF;\n",
       "    }\n",
       "  </style>\n",
       "\n",
       "      <script>\n",
       "        const buttonEl =\n",
       "          document.querySelector('#df-df806cae-3387-4f03-bdae-738a45dd3e06 button.colab-df-convert');\n",
       "        buttonEl.style.display =\n",
       "          google.colab.kernel.accessAllowed ? 'block' : 'none';\n",
       "\n",
       "        async function convertToInteractive(key) {\n",
       "          const element = document.querySelector('#df-df806cae-3387-4f03-bdae-738a45dd3e06');\n",
       "          const dataTable =\n",
       "            await google.colab.kernel.invokeFunction('convertToInteractive',\n",
       "                                                     [key], {});\n",
       "          if (!dataTable) return;\n",
       "\n",
       "          const docLinkHtml = 'Like what you see? Visit the ' +\n",
       "            '<a target=\"_blank\" href=https://colab.research.google.com/notebooks/data_table.ipynb>data table notebook</a>'\n",
       "            + ' to learn more about interactive tables.';\n",
       "          element.innerHTML = '';\n",
       "          dataTable['output_type'] = 'display_data';\n",
       "          await google.colab.output.renderOutput(dataTable, element);\n",
       "          const docLink = document.createElement('div');\n",
       "          docLink.innerHTML = docLinkHtml;\n",
       "          element.appendChild(docLink);\n",
       "        }\n",
       "      </script>\n",
       "    </div>\n",
       "  </div>\n",
       "  "
      ],
      "text/plain": [
       "      video_id                     title     channel_title  \\\n",
       "0  --JinobXWPk  DANGEROUS Jungle Spider!  Brave Wilderness   \n",
       "1  --JinobXWPk  DANGEROUS Jungle Spider!  Brave Wilderness   \n",
       "2  --JinobXWPk  DANGEROUS Jungle Spider!  Brave Wilderness   \n",
       "3  --JinobXWPk  DANGEROUS Jungle Spider!  Brave Wilderness   \n",
       "4  --JinobXWPk  DANGEROUS Jungle Spider!  Brave Wilderness   \n",
       "\n",
       "                                        comment_text  likes  replies  \\\n",
       "0  I saw this wandering spider in our bathroom se...      0        0   \n",
       "1  Can't you just stick to small ants, and bees? ...      0        0   \n",
       "2  Brazilian wandering spider is the deadliest sp...      0        0   \n",
       "3  Nothing a can of hairspray and lighter couldn'...      0        0   \n",
       "4  Hey Coyote! Can you do an episode on the Japan...      0        0   \n",
       "\n",
       "                                         cleaned_txt  \n",
       "0  I saw this wandering spider in our bathroom se...  \n",
       "1  Can't you just stick to small ants, and bees? ...  \n",
       "2  Brazilian wandering spider is the deadliest sp...  \n",
       "3  Nothing a can of hairspray and lighter couldn'...  \n",
       "4  Hey Coyote! Can you do an episode on the Japan...  "
      ]
     },
     "execution_count": 194,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df_full = videos[[\"video_id\", \"title\", \"channel_title\"]].merge(df, \n",
    "                                                               left_on='video_id', right_on='video_id')\n",
    "df_full[\"title\"] = df_full[\"title\"]#.str.lower()\n",
    "df_full[\"channel_title\"] = df_full[\"channel_title\"]#.str.lower()\n",
    "df_full[\"cleaned_txt\"] = df_full[\"cleaned_txt\"]#.str.lower()\n",
    "df_full.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "foreign-thomas",
   "metadata": {
    "id": "completed-constitution"
   },
   "source": [
    "Будем использовать zero-shot обучение со следующим prompt'ом:\n",
    "\n",
    "Train:  \n",
    "\\<BOS>CHANNEL: _channel_title_   \n",
    "VIDEO: _video_title_  \n",
    "COMMENTARY: _comment text_ \\<EOS>\n",
    "\n",
    "Test:  \n",
    "\\<BOS>CHANNEL: _channel_title_   \n",
    "VIDEO: _video_title_  \n",
    "COMMENTARY:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 195,
   "id": "innocent-input",
   "metadata": {
    "id": "fleet-whole"
   },
   "outputs": [],
   "source": [
    "# prompts = \"<BOS>CHANNEL: \" + df_full[\"channel_title\"] + \\\n",
    "#           \" TITLE: \" + df_full[\"title\"] + \\\n",
    "#           \" COMMENT: \" + df_full[\"cleaned_txt\"] + \" <EOS>\"\n",
    "# print(prompts[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 196,
   "id": "christian-bridge",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "feyKViRnqJc0",
    "outputId": "2cc22150-c92e-4533-92da-1c9ffafbe7b6"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<BOS> TOPIC: DANGEROUS Jungle Spider! COMMENT: I saw this wandering spider in our bathroom seriously I'm not lying. nPS I will never go to Costa Rica!: <EOS>\n"
     ]
    }
   ],
   "source": [
    "prompts = \"<BOS> TOPIC: \" + df_full[\"title\"] + \\\n",
    "          \" COMMENT: \" + df_full[\"cleaned_txt\"] + \" <EOS>\"\n",
    "print(prompts[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 197,
   "id": "coordinate-dependence",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "subsequent-elevation",
    "outputId": "7f4c87db-9ec3-40eb-f930-ca11c069dfe7"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<BOS> TOPIC: Best Tom Petty Interview Ever COMMENT: Where are the other interviews Gary Chandling did? <EOS>\n"
     ]
    }
   ],
   "source": [
    "print(prompts[140])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 198,
   "id": "speaking-province",
   "metadata": {
    "id": "french-accountability"
   },
   "outputs": [],
   "source": [
    "df_full[\"prompt\"] = prompts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 199,
   "id": "dynamic-desperate",
   "metadata": {
    "id": "trained-seattle"
   },
   "outputs": [],
   "source": [
    "df_full.reset_index(inplace=True, drop=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 200,
   "id": "polished-strike",
   "metadata": {
    "id": "bottom-knowing"
   },
   "outputs": [],
   "source": [
    "train_index = df_full.sample(int(df_full.shape[0] * 0.9), random_state=42).index\n",
    "df_full.loc[train_index].reset_index(drop=True).to_csv(\"prompts_train_02.csv\")\n",
    "\n",
    "test_index = list(set(df_full.index) - set(train_index))\n",
    "df_full.loc[test_index].reset_index(drop=True).to_csv(\"prompts_test_02.csv\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "executive-coalition",
   "metadata": {
    "id": "exterior-cache"
   },
   "source": [
    "# Подготовка к обучению"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "simplified-tactics",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "execution": {
     "iopub.execute_input": "2023-04-12T20:14:52.485955Z",
     "iopub.status.busy": "2023-04-12T20:14:52.485487Z",
     "iopub.status.idle": "2023-04-12T20:14:52.517276Z",
     "shell.execute_reply": "2023-04-12T20:14:52.516185Z",
     "shell.execute_reply.started": "2023-04-12T20:14:52.485920Z"
    },
    "id": "3bV3N11hc0Al",
    "outputId": "93c03b10-28f5-4bdd-d6d7-20e71a9eb8e7"
   },
   "outputs": [],
   "source": [
    "# from google.colab import drive\n",
    "# drive.mount('/content/gdrive')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "editorial-arbor",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "execution": {
     "iopub.execute_input": "2023-04-12T20:14:52.565111Z",
     "iopub.status.busy": "2023-04-12T20:14:52.564822Z",
     "iopub.status.idle": "2023-04-12T20:14:52.570092Z",
     "shell.execute_reply": "2023-04-12T20:14:52.568747Z",
     "shell.execute_reply.started": "2023-04-12T20:14:52.565083Z"
    },
    "id": "MWde5q5jdZc2",
    "outputId": "7865e7ff-45bc-474c-fa65-7f7e4e362a8a"
   },
   "outputs": [],
   "source": [
    "# %cd ./gdrive/MyDrive/Colab Notebooks/ml2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "featured-stone",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-04-13T07:37:44.002545Z",
     "iopub.status.busy": "2023-04-13T07:37:44.002222Z",
     "iopub.status.idle": "2023-04-13T07:37:45.096711Z",
     "shell.execute_reply": "2023-04-13T07:37:45.094948Z",
     "shell.execute_reply.started": "2023-04-13T07:37:44.002513Z"
    }
   },
   "outputs": [],
   "source": [
    "!mkdir ./chkp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "starting-guide",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "collapsed": true,
    "execution": {
     "iopub.execute_input": "2023-04-13T07:37:45.107352Z",
     "iopub.status.busy": "2023-04-13T07:37:45.104451Z",
     "iopub.status.idle": "2023-04-13T07:38:15.679220Z",
     "shell.execute_reply": "2023-04-13T07:38:15.677863Z",
     "shell.execute_reply.started": "2023-04-13T07:37:45.107302Z"
    },
    "id": "NNYGCW1xdUsn",
    "outputId": "7475a204-3194-4961-d71c-16779fdaf981"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Collecting bitsandbytes\n",
      "  Downloading bitsandbytes-0.38.1-py3-none-any.whl (104.3 MB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m104.3/104.3 MB\u001b[0m \u001b[31m8.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m00:01\u001b[0m00:01\u001b[0m\n",
      "\u001b[?25hInstalling collected packages: bitsandbytes\n",
      "Successfully installed bitsandbytes-0.38.1\n",
      "\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\u001b[33m\n",
      "\u001b[0mRequirement already satisfied: datasets in /opt/conda/lib/python3.7/site-packages (2.1.0)\n",
      "Requirement already satisfied: transformers in /opt/conda/lib/python3.7/site-packages (4.27.4)\n",
      "Requirement already satisfied: pandas in /opt/conda/lib/python3.7/site-packages (from datasets) (1.3.5)\n",
      "Requirement already satisfied: multiprocess in /opt/conda/lib/python3.7/site-packages (from datasets) (0.70.14)\n",
      "Requirement already satisfied: tqdm>=4.62.1 in /opt/conda/lib/python3.7/site-packages (from datasets) (4.64.1)\n",
      "Requirement already satisfied: xxhash in /opt/conda/lib/python3.7/site-packages (from datasets) (3.2.0)\n",
      "Requirement already satisfied: packaging in /opt/conda/lib/python3.7/site-packages (from datasets) (23.0)\n",
      "Requirement already satisfied: pyarrow>=5.0.0 in /opt/conda/lib/python3.7/site-packages (from datasets) (5.0.0)\n",
      "Requirement already satisfied: importlib-metadata in /opt/conda/lib/python3.7/site-packages (from datasets) (4.11.4)\n",
      "Requirement already satisfied: dill in /opt/conda/lib/python3.7/site-packages (from datasets) (0.3.6)\n",
      "Requirement already satisfied: responses<0.19 in /opt/conda/lib/python3.7/site-packages (from datasets) (0.18.0)\n",
      "Requirement already satisfied: fsspec[http]>=2021.05.0 in /opt/conda/lib/python3.7/site-packages (from datasets) (2023.1.0)\n",
      "Requirement already satisfied: requests>=2.19.0 in /opt/conda/lib/python3.7/site-packages (from datasets) (2.28.2)\n",
      "Requirement already satisfied: huggingface-hub<1.0.0,>=0.1.0 in /opt/conda/lib/python3.7/site-packages (from datasets) (0.13.3)\n",
      "Requirement already satisfied: aiohttp in /opt/conda/lib/python3.7/site-packages (from datasets) (3.8.3)\n",
      "Requirement already satisfied: numpy>=1.17 in /opt/conda/lib/python3.7/site-packages (from datasets) (1.21.6)\n",
      "Requirement already satisfied: filelock in /opt/conda/lib/python3.7/site-packages (from transformers) (3.9.0)\n",
      "Requirement already satisfied: pyyaml>=5.1 in /opt/conda/lib/python3.7/site-packages (from transformers) (6.0)\n",
      "Requirement already satisfied: tokenizers!=0.11.3,<0.14,>=0.11.1 in /opt/conda/lib/python3.7/site-packages (from transformers) (0.13.2)\n",
      "Requirement already satisfied: regex!=2019.12.17 in /opt/conda/lib/python3.7/site-packages (from transformers) (2021.11.10)\n",
      "Requirement already satisfied: charset-normalizer<3.0,>=2.0 in /opt/conda/lib/python3.7/site-packages (from aiohttp->datasets) (2.1.1)\n",
      "Requirement already satisfied: multidict<7.0,>=4.5 in /opt/conda/lib/python3.7/site-packages (from aiohttp->datasets) (6.0.4)\n",
      "Requirement already satisfied: yarl<2.0,>=1.0 in /opt/conda/lib/python3.7/site-packages (from aiohttp->datasets) (1.8.2)\n",
      "Requirement already satisfied: attrs>=17.3.0 in /opt/conda/lib/python3.7/site-packages (from aiohttp->datasets) (22.2.0)\n",
      "Requirement already satisfied: async-timeout<5.0,>=4.0.0a3 in /opt/conda/lib/python3.7/site-packages (from aiohttp->datasets) (4.0.2)\n",
      "Requirement already satisfied: asynctest==0.13.0 in /opt/conda/lib/python3.7/site-packages (from aiohttp->datasets) (0.13.0)\n",
      "Requirement already satisfied: aiosignal>=1.1.2 in /opt/conda/lib/python3.7/site-packages (from aiohttp->datasets) (1.3.1)\n",
      "Requirement already satisfied: frozenlist>=1.1.1 in /opt/conda/lib/python3.7/site-packages (from aiohttp->datasets) (1.3.3)\n",
      "Requirement already satisfied: typing-extensions>=3.7.4 in /opt/conda/lib/python3.7/site-packages (from aiohttp->datasets) (4.4.0)\n",
      "Requirement already satisfied: urllib3<1.27,>=1.21.1 in /opt/conda/lib/python3.7/site-packages (from requests>=2.19.0->datasets) (1.26.14)\n",
      "Requirement already satisfied: idna<4,>=2.5 in /opt/conda/lib/python3.7/site-packages (from requests>=2.19.0->datasets) (3.4)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in /opt/conda/lib/python3.7/site-packages (from requests>=2.19.0->datasets) (2022.12.7)\n",
      "Requirement already satisfied: zipp>=0.5 in /opt/conda/lib/python3.7/site-packages (from importlib-metadata->datasets) (3.11.0)\n",
      "Requirement already satisfied: python-dateutil>=2.7.3 in /opt/conda/lib/python3.7/site-packages (from pandas->datasets) (2.8.2)\n",
      "Requirement already satisfied: pytz>=2017.3 in /opt/conda/lib/python3.7/site-packages (from pandas->datasets) (2023.3)\n",
      "Requirement already satisfied: six>=1.5 in /opt/conda/lib/python3.7/site-packages (from python-dateutil>=2.7.3->pandas->datasets) (1.16.0)\n",
      "\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\u001b[33m\n",
      "\u001b[0m"
     ]
    }
   ],
   "source": [
    "!pip install bitsandbytes\n",
    "!pip install datasets transformers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "moral-tribute",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "execution": {
     "iopub.execute_input": "2023-04-13T07:38:15.685046Z",
     "iopub.status.busy": "2023-04-13T07:38:15.684565Z",
     "iopub.status.idle": "2023-04-13T07:38:26.677158Z",
     "shell.execute_reply": "2023-04-13T07:38:26.676012Z",
     "shell.execute_reply.started": "2023-04-13T07:38:15.685009Z"
    },
    "id": "x7B_tbh3dn_j",
    "outputId": "76ceff04-811c-4812-f1c4-f032dd894747"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "===================================BUG REPORT===================================\n",
      "Welcome to bitsandbytes. For bug reports, please submit your error trace to: https://github.com/TimDettmers/bitsandbytes/issues\n",
      "================================================================================\n",
      "CUDA SETUP: CUDA runtime path found: /usr/local/cuda/lib64/libcudart.so\n",
      "CUDA SETUP: Highest compute capability among GPUs detected: 6.1\n",
      "CUDA SETUP: Detected CUDA version 110\n",
      "CUDA SETUP: Loading binary /home/sapetrov/.local/lib/python3.8/site-packages/bitsandbytes/libbitsandbytes_cuda110_nocublaslt.so...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/sapetrov/.local/lib/python3.8/site-packages/bitsandbytes/cuda_setup/main.py:136: UserWarning: WARNING: The following directories listed in your path were found to be non-existent: {PosixPath('/usr/local/cuda/extras/CUPTI/lib64')}\n",
      "  warn(msg)\n",
      "/home/sapetrov/.local/lib/python3.8/site-packages/bitsandbytes/cuda_setup/main.py:136: UserWarning: WARNING: Compute capability < 7.5 detected! Only slow 8-bit matmul is supported for your GPU!\n",
      "  warn(msg)\n"
     ]
    }
   ],
   "source": [
    "# You can also adapt this script on your own causal language modeling task. Pointers for this are left as comments.\n",
    "\n",
    "import bitsandbytes as bnb\n",
    "import pandas as pd\n",
    "\n",
    "\n",
    "import argparse\n",
    "import logging\n",
    "import math\n",
    "import os\n",
    "import random\n",
    "from itertools import chain\n",
    "from pathlib import Path\n",
    "\n",
    "import datasets\n",
    "import torch\n",
    "from datasets import load_dataset\n",
    "from torch.utils.data import DataLoader, Dataset\n",
    "from tqdm.auto import tqdm\n",
    "\n",
    "import transformers\n",
    "from huggingface_hub import Repository\n",
    "from transformers import (\n",
    "    CONFIG_MAPPING,\n",
    "    MODEL_MAPPING,\n",
    "    AdamW,\n",
    "    AutoConfig,\n",
    "    AutoModelForCausalLM,\n",
    "    AutoTokenizer,\n",
    "    SchedulerType,\n",
    "    default_data_collator,\n",
    "    get_scheduler,\n",
    "    set_seed,\n",
    ")\n",
    "from transformers.file_utils import get_full_repo_name\n",
    "from transformers.utils.versions import require_version"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "oriental-assessment",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-04-13T07:38:26.682977Z",
     "iopub.status.busy": "2023-04-13T07:38:26.680823Z",
     "iopub.status.idle": "2023-04-13T07:38:26.863443Z",
     "shell.execute_reply": "2023-04-13T07:38:26.862056Z",
     "shell.execute_reply.started": "2023-04-13T07:38:26.682937Z"
    },
    "id": "danish-blast"
   },
   "outputs": [],
   "source": [
    "logger = logging.getLogger(__name__)\n",
    "\n",
    "require_version(\"datasets>=1.16.1\", \"To fix: pip install -r examples/pytorch/language-modeling/requirements.txt\")\n",
    "\n",
    "MODEL_CONFIG_CLASSES = list(MODEL_MAPPING.keys())\n",
    "MODEL_TYPES = tuple(conf.model_type for conf in MODEL_CONFIG_CLASSES)\n",
    "\n",
    "\n",
    "def parse_args():\n",
    "    parser = argparse.ArgumentParser(description=\"Finetune a transformers model on a causal language modeling task\")\n",
    "    parser.add_argument(\n",
    "        \"--dataset_name\",\n",
    "        type=str,\n",
    "        default=None,\n",
    "        help=\"The name of the dataset to use (via the datasets library).\",\n",
    "    )\n",
    "    parser.add_argument(\n",
    "        \"--dataset_config_name\",\n",
    "        type=str,\n",
    "        default=None,\n",
    "        help=\"The configuration name of the dataset to use (via the datasets library).\",\n",
    "    )\n",
    "    parser.add_argument(\n",
    "        \"--text_column_name\",\n",
    "        type=str,\n",
    "        default=None,\n",
    "        help=\"The name of the column containing the text data.\",\n",
    "    )\n",
    "    parser.add_argument(\n",
    "        \"--dataset_streaming\",\n",
    "        action=\"store_true\",\n",
    "        help=\"If passed, will use dataset streaming (via the datasets library)\",\n",
    "    )\n",
    "    parser.add_argument(\n",
    "        \"--model_name_or_path\",\n",
    "        type=str,\n",
    "        help=\"Path to pretrained model or model identifier from huggingface.co/models.\",\n",
    "        required=False,\n",
    "    )\n",
    "    parser.add_argument(\n",
    "        \"--config_name\",\n",
    "        type=str,\n",
    "        default=None,\n",
    "        help=\"Pretrained config name or path if not the same as model_name\",\n",
    "    )\n",
    "    parser.add_argument(\n",
    "        \"--tokenizer_name\",\n",
    "        type=str,\n",
    "        default=None,\n",
    "        help=\"Pretrained tokenizer name or path if not the same as model_name\",\n",
    "    )\n",
    "    parser.add_argument(\n",
    "        \"--use_slow_tokenizer\",\n",
    "        action=\"store_true\",\n",
    "        help=\"If passed, will use a slow tokenizer (not backed by the 🤗 Tokenizers library).\",\n",
    "    )\n",
    "    parser.add_argument(\n",
    "        \"--per_device_train_batch_size\",\n",
    "        type=int,\n",
    "        default=1,\n",
    "        help=\"Batch size (per device) for the training dataloader.\",\n",
    "    )\n",
    "    parser.add_argument(\n",
    "        \"--learning_rate\",\n",
    "        type=float,\n",
    "        default=5e-5,\n",
    "        help=\"Initial learning rate (after the potential warmup period) to use.\",\n",
    "    )\n",
    "    parser.add_argument(\"--weight_decay\", type=float, default=0.0, help=\"Weight decay to use.\")\n",
    "    parser.add_argument(\"--num_train_epochs\", type=int, default=1, help=\"Total number of training epochs to perform.\")\n",
    "    parser.add_argument(\n",
    "        \"--max_train_steps\",\n",
    "        type=int,\n",
    "        default=None,\n",
    "        help=\"Total number of training steps to perform. If provided, overrides num_train_epochs.\",\n",
    "    )\n",
    "    parser.add_argument(\n",
    "        \"--gradient_accumulation_steps\",\n",
    "        type=int,\n",
    "        default=1,\n",
    "        help=\"Number of updates steps to accumulate before performing a backward/update pass.\",\n",
    "    )\n",
    "    parser.add_argument(\n",
    "        \"--lr_scheduler_type\",\n",
    "        type=SchedulerType,\n",
    "        default=\"linear\",\n",
    "        help=\"The scheduler type to use.\",\n",
    "        choices=[\"linear\", \"cosine\", \"cosine_with_restarts\", \"polynomial\", \"constant\", \"constant_with_warmup\"],\n",
    "    )\n",
    "    parser.add_argument(\n",
    "        \"--num_warmup_steps\", type=int, default=3000, help=\"Number of steps for the warmup in the lr scheduler.\"\n",
    "    )\n",
    "    parser.add_argument(\"--output_dir\", type=str, default=None, help=\"Where to store the final model.\")\n",
    "    parser.add_argument(\"--seed\", type=int, default=None, help=\"A seed for reproducible training.\")\n",
    "    parser.add_argument(\n",
    "        \"--model_type\",\n",
    "        type=str,\n",
    "        default=None,\n",
    "        help=\"Model type to use if training from scratch.\",\n",
    "        choices=MODEL_TYPES,\n",
    "    )\n",
    "    parser.add_argument(\n",
    "        \"--block_size\",\n",
    "        type=int,\n",
    "        default=None,\n",
    "        help=\"Optional input sequence length after tokenization. The training dataset will be truncated in block of this size for training. Default to the model max input length for single sentence inputs (take into account special tokens).\",\n",
    "    )\n",
    "    parser.add_argument(\n",
    "        \"--preprocessing_num_workers\",\n",
    "        type=int,\n",
    "        default=None,\n",
    "        help=\"The number of processes to use for the preprocessing.\",\n",
    "    )\n",
    "    parser.add_argument(\n",
    "        \"--overwrite_cache\", type=bool, default=False, help=\"Overwrite the cached training and evaluation sets\"\n",
    "    )\n",
    "    parser.add_argument(\n",
    "        \"--no_keep_linebreaks\", action=\"store_true\", help=\"Do not keep line breaks when using TXT files.\"\n",
    "    )\n",
    "    parser.add_argument(\"--push_to_hub\", action=\"store_true\", help=\"Whether or not to push the model to the Hub.\")\n",
    "    parser.add_argument(\n",
    "        \"--hub_model_id\", type=str, help=\"The name of the repository to keep in sync with the local `output_dir`.\"\n",
    "    )\n",
    "    parser.add_argument(\"--hub_token\", type=str, help=\"The token to use to push to the Model Hub.\")\n",
    "    args = parser.parse_args(args=[])\n",
    "\n",
    "    if args.push_to_hub:\n",
    "        assert args.output_dir is not None, \"Need an `output_dir` to create a repo when `--push_to_hub` is passed.\"\n",
    "\n",
    "    return args"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "operating-shape",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-04-13T07:38:26.865365Z",
     "iopub.status.busy": "2023-04-13T07:38:26.865002Z",
     "iopub.status.idle": "2023-04-13T07:38:26.877602Z",
     "shell.execute_reply": "2023-04-13T07:38:26.876435Z",
     "shell.execute_reply.started": "2023-04-13T07:38:26.865326Z"
    },
    "id": "colored-degree"
   },
   "outputs": [],
   "source": [
    "args = parse_args() # get default arguments\n",
    "\n",
    "# If passed along, set the training seed now.\n",
    "if args.seed is not None:\n",
    "    set_seed(args.seed)\n",
    "\n",
    "# args.dataset_name = 'c4'\n",
    "# args.dataset_streaming = True\n",
    "# args.dataset_config_name = \"en\"\n",
    "args.text_column_name = \"prompt\"\n",
    "args.model_name_or_path = 'gpt2'\n",
    "args.block_size = 1024\n",
    "args.max_train_steps = 1_000_000\n",
    "args.log_loss_interval = 25\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "governing-token",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-04-13T07:38:26.879970Z",
     "iopub.status.busy": "2023-04-13T07:38:26.879380Z",
     "iopub.status.idle": "2023-04-13T07:38:26.886955Z",
     "shell.execute_reply": "2023-04-13T07:38:26.885542Z",
     "shell.execute_reply.started": "2023-04-13T07:38:26.879928Z"
    },
    "id": "NVd9O0h6eKPt"
   },
   "outputs": [],
   "source": [
    "device = torch.device('cuda:3' if torch.cuda.is_available() else 'cpu')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "general-airline",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-04-13T07:38:26.889949Z",
     "iopub.status.busy": "2023-04-13T07:38:26.888757Z",
     "iopub.status.idle": "2023-04-13T07:38:28.977649Z",
     "shell.execute_reply": "2023-04-13T07:38:28.976624Z",
     "shell.execute_reply.started": "2023-04-13T07:38:26.889905Z"
    },
    "id": "geological-projector"
   },
   "outputs": [],
   "source": [
    "# CONFIG\n",
    "\n",
    "if args.config_name:\n",
    "    config = AutoConfig.from_pretrained(args.config_name)\n",
    "elif args.model_name_or_path:\n",
    "    config = AutoConfig.from_pretrained(args.model_name_or_path)\n",
    "else:\n",
    "    config = CONFIG_MAPPING[args.model_type]()\n",
    "    logger.warning(\"You are instantiating a new config instance from scratch.\")\n",
    "\n",
    "# TOKENIZER\n",
    "if args.tokenizer_name:\n",
    "    tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_name, use_fast=not args.use_slow_tokenizer)\n",
    "elif args.model_name_or_path:\n",
    "    tokenizer = AutoTokenizer.from_pretrained(args.model_name_or_path, use_fast=not args.use_slow_tokenizer)\n",
    "else:\n",
    "    raise ValueError(\n",
    "        \"You are instantiating a new tokenizer from scratch. This is not supported by this script.\"\n",
    "        \"You can do it from another script, save it, and load it from here, using --tokenizer_name.\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "voluntary-discrimination",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "execution": {
     "iopub.execute_input": "2023-04-13T07:38:28.979741Z",
     "iopub.status.busy": "2023-04-13T07:38:28.979348Z",
     "iopub.status.idle": "2023-04-13T07:38:33.942769Z",
     "shell.execute_reply": "2023-04-13T07:38:33.941632Z",
     "shell.execute_reply.started": "2023-04-13T07:38:28.979702Z"
    },
    "id": "bronze-missouri",
    "outputId": "2db06a62-4049-44f3-92c2-a77a84cbf903"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 1 µs, sys: 2 µs, total: 3 µs\n",
      "Wall time: 7.39 µs\n"
     ]
    }
   ],
   "source": [
    "%time\n",
    "# DATASET\n",
    "\n",
    "class Comments(Dataset):  \n",
    "    def __init__(self, \n",
    "                 filename, \n",
    "                 truncate=False, \n",
    "                 gpt2_type=args.model_name_or_path, \n",
    "                 text_column_name=args.text_column_name,\n",
    "                 max_length=1024, \n",
    "                 tokenizer=None, \n",
    "                 block_size=None):\n",
    "        \n",
    "        self.df_full = pd.read_csv(filename).iloc[:, 1:]\n",
    "        self.texts = self.preprocess(self.df_full, block_size, text_column_name, tokenizer)\n",
    "        self.count = self.df_full.shape[0]\n",
    "        self.tokenizer = tokenizer\n",
    "        \n",
    "    def preprocess(self, data, block_size, text_column_name, tokenizer):\n",
    "        texts = data[text_column_name].apply(lambda x: x[:block_size]+\" <EOS>\" if len(x) > block_size else x)\n",
    "#         texts = texts.apply(tokenizer.encode)\n",
    "        return texts\n",
    "\n",
    "    def __len__(self):\n",
    "        return self.count\n",
    "\n",
    "    def __getitem__(self, item):\n",
    "        output = torch.tensor(self.tokenizer.encode(self.texts[item]))\n",
    "        return output \n",
    "\n",
    "\n",
    "if args.block_size is None:\n",
    "    block_size = tokenizer.model_max_length\n",
    "    if block_size > 1024:\n",
    "        logger.warning(\n",
    "            f\"The tokenizer picked seems to have a very large `model_max_length` ({tokenizer.model_max_length}). \"\n",
    "            \"Picking 1024 instead. You can change that default value by passing --block_size xxx.\"\n",
    "        )\n",
    "    block_size = 1024\n",
    "else:\n",
    "    if args.block_size > tokenizer.model_max_length:\n",
    "        logger.warning(\n",
    "            f\"The block_size passed ({args.block_size}) is larger than the maximum length for the model\"\n",
    "            f\"({tokenizer.model_max_length}). Using block_size={tokenizer.model_max_length}.\"\n",
    "        )\n",
    "    block_size = min(args.block_size, tokenizer.model_max_length)\n",
    "    \n",
    "train_dataset = Comments(\"./prompts_train_02.csv\", \n",
    "                           text_column_name=args.text_column_name, \n",
    "                           tokenizer=tokenizer,\n",
    "                           block_size=block_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "macro-burke",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "execution": {
     "iopub.execute_input": "2023-04-13T07:38:33.944575Z",
     "iopub.status.busy": "2023-04-13T07:38:33.944206Z",
     "iopub.status.idle": "2023-04-13T07:38:33.997107Z",
     "shell.execute_reply": "2023-04-13T07:38:33.995932Z",
     "shell.execute_reply.started": "2023-04-13T07:38:33.944538Z"
    },
    "id": "-JgDNXJaqs5s",
    "outputId": "884ff9db-db3e-4cd8-9471-3e5f79b12295"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([   27,    33,  2640,    29, 28662,  2149,    25,  4162,  8314,   262,\n",
       "          4100, 12558,  1475,   396,   287,  2177,    30,  9440, 10979,    25,\n",
       "           880,   340,  2152,   780, 17180,   460,   470,   655,  1011,   606,\n",
       "           736,    13,  6949,  1049,   410,   312,  1279,    36,  2640,    29]),\n",
       " \"<BOS> TOPIC: Why Does the Mac Mini Exist in 2017? COMMENT: well it exist because apple can't just take them back. anyway great vid <EOS>\")"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_dataset[100], train_dataset.df_full.iloc[100].prompt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "healthy-orlando",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "execution": {
     "iopub.execute_input": "2023-04-13T07:38:34.004675Z",
     "iopub.status.busy": "2023-04-13T07:38:34.003678Z",
     "iopub.status.idle": "2023-04-13T07:38:34.369396Z",
     "shell.execute_reply": "2023-04-13T07:38:34.367801Z",
     "shell.execute_reply.started": "2023-04-13T07:38:34.004621Z"
    },
    "id": "RRYgPsTkqwil",
    "outputId": "6c23ad59-4b2d-4378-dd53-0f76945d48fc"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([   27,    33,  2640,    29, 28662,  2149,    25,   770,  6119,   314,\n",
       "         38514,   284,  4294,   303,   262,   642,    87,    20,    87,    20,\n",
       "          6256,  1134,   338, 23315,  9440, 10979,    25,  4222,   466,   257,\n",
       "          2193,  2068,   329,  6155,   319,   534,  2832,  1279,    36,  2640,\n",
       "            29]),\n",
       " \"<BOS> TOPIC: This Week I Learned to Solve the 5x5x5 Rubik's Cube COMMENT: Please do a learn quick for walking on your hands <EOS>\")"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_dataset[400], train_dataset.df_full.iloc[400].prompt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "found-content",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-04-13T07:38:34.371947Z",
     "iopub.status.busy": "2023-04-13T07:38:34.371464Z",
     "iopub.status.idle": "2023-04-13T07:38:34.378418Z",
     "shell.execute_reply": "2023-04-13T07:38:34.377119Z",
     "shell.execute_reply.started": "2023-04-13T07:38:34.371890Z"
    },
    "id": "synthetic-ending"
   },
   "outputs": [],
   "source": [
    "# LOADER\n",
    "train_dataloader = DataLoader(\n",
    "    train_dataset, batch_size=args.per_device_train_batch_size, shuffle=True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "composed-commonwealth",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "execution": {
     "iopub.execute_input": "2023-04-13T07:38:34.381958Z",
     "iopub.status.busy": "2023-04-13T07:38:34.380998Z",
     "iopub.status.idle": "2023-04-13T07:38:46.712862Z",
     "shell.execute_reply": "2023-04-13T07:38:46.711547Z",
     "shell.execute_reply.started": "2023-04-13T07:38:34.381889Z"
    },
    "id": "nervous-dairy",
    "outputId": "7a783f3f-9548-4c59-d1b8-4187bbab4f1c"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "GPT2LMHeadModel(\n",
       "  (transformer): GPT2Model(\n",
       "    (wte): Embedding(50257, 768)\n",
       "    (wpe): Embedding(1024, 768)\n",
       "    (drop): Dropout(p=0.1, inplace=False)\n",
       "    (h): ModuleList(\n",
       "      (0): GPT2Block(\n",
       "        (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): GPT2Attention(\n",
       "          (c_attn): Conv1D()\n",
       "          (c_proj): Conv1D()\n",
       "          (attn_dropout): Dropout(p=0.1, inplace=False)\n",
       "          (resid_dropout): Dropout(p=0.1, inplace=False)\n",
       "        )\n",
       "        (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "        (mlp): GPT2MLP(\n",
       "          (c_fc): Conv1D()\n",
       "          (c_proj): Conv1D()\n",
       "          (act): NewGELUActivation()\n",
       "          (dropout): Dropout(p=0.1, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (1): GPT2Block(\n",
       "        (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): GPT2Attention(\n",
       "          (c_attn): Conv1D()\n",
       "          (c_proj): Conv1D()\n",
       "          (attn_dropout): Dropout(p=0.1, inplace=False)\n",
       "          (resid_dropout): Dropout(p=0.1, inplace=False)\n",
       "        )\n",
       "        (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "        (mlp): GPT2MLP(\n",
       "          (c_fc): Conv1D()\n",
       "          (c_proj): Conv1D()\n",
       "          (act): NewGELUActivation()\n",
       "          (dropout): Dropout(p=0.1, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (2): GPT2Block(\n",
       "        (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): GPT2Attention(\n",
       "          (c_attn): Conv1D()\n",
       "          (c_proj): Conv1D()\n",
       "          (attn_dropout): Dropout(p=0.1, inplace=False)\n",
       "          (resid_dropout): Dropout(p=0.1, inplace=False)\n",
       "        )\n",
       "        (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "        (mlp): GPT2MLP(\n",
       "          (c_fc): Conv1D()\n",
       "          (c_proj): Conv1D()\n",
       "          (act): NewGELUActivation()\n",
       "          (dropout): Dropout(p=0.1, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (3): GPT2Block(\n",
       "        (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): GPT2Attention(\n",
       "          (c_attn): Conv1D()\n",
       "          (c_proj): Conv1D()\n",
       "          (attn_dropout): Dropout(p=0.1, inplace=False)\n",
       "          (resid_dropout): Dropout(p=0.1, inplace=False)\n",
       "        )\n",
       "        (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "        (mlp): GPT2MLP(\n",
       "          (c_fc): Conv1D()\n",
       "          (c_proj): Conv1D()\n",
       "          (act): NewGELUActivation()\n",
       "          (dropout): Dropout(p=0.1, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (4): GPT2Block(\n",
       "        (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): GPT2Attention(\n",
       "          (c_attn): Conv1D()\n",
       "          (c_proj): Conv1D()\n",
       "          (attn_dropout): Dropout(p=0.1, inplace=False)\n",
       "          (resid_dropout): Dropout(p=0.1, inplace=False)\n",
       "        )\n",
       "        (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "        (mlp): GPT2MLP(\n",
       "          (c_fc): Conv1D()\n",
       "          (c_proj): Conv1D()\n",
       "          (act): NewGELUActivation()\n",
       "          (dropout): Dropout(p=0.1, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (5): GPT2Block(\n",
       "        (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): GPT2Attention(\n",
       "          (c_attn): Conv1D()\n",
       "          (c_proj): Conv1D()\n",
       "          (attn_dropout): Dropout(p=0.1, inplace=False)\n",
       "          (resid_dropout): Dropout(p=0.1, inplace=False)\n",
       "        )\n",
       "        (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "        (mlp): GPT2MLP(\n",
       "          (c_fc): Conv1D()\n",
       "          (c_proj): Conv1D()\n",
       "          (act): NewGELUActivation()\n",
       "          (dropout): Dropout(p=0.1, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (6): GPT2Block(\n",
       "        (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): GPT2Attention(\n",
       "          (c_attn): Conv1D()\n",
       "          (c_proj): Conv1D()\n",
       "          (attn_dropout): Dropout(p=0.1, inplace=False)\n",
       "          (resid_dropout): Dropout(p=0.1, inplace=False)\n",
       "        )\n",
       "        (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "        (mlp): GPT2MLP(\n",
       "          (c_fc): Conv1D()\n",
       "          (c_proj): Conv1D()\n",
       "          (act): NewGELUActivation()\n",
       "          (dropout): Dropout(p=0.1, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (7): GPT2Block(\n",
       "        (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): GPT2Attention(\n",
       "          (c_attn): Conv1D()\n",
       "          (c_proj): Conv1D()\n",
       "          (attn_dropout): Dropout(p=0.1, inplace=False)\n",
       "          (resid_dropout): Dropout(p=0.1, inplace=False)\n",
       "        )\n",
       "        (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "        (mlp): GPT2MLP(\n",
       "          (c_fc): Conv1D()\n",
       "          (c_proj): Conv1D()\n",
       "          (act): NewGELUActivation()\n",
       "          (dropout): Dropout(p=0.1, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (8): GPT2Block(\n",
       "        (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): GPT2Attention(\n",
       "          (c_attn): Conv1D()\n",
       "          (c_proj): Conv1D()\n",
       "          (attn_dropout): Dropout(p=0.1, inplace=False)\n",
       "          (resid_dropout): Dropout(p=0.1, inplace=False)\n",
       "        )\n",
       "        (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "        (mlp): GPT2MLP(\n",
       "          (c_fc): Conv1D()\n",
       "          (c_proj): Conv1D()\n",
       "          (act): NewGELUActivation()\n",
       "          (dropout): Dropout(p=0.1, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (9): GPT2Block(\n",
       "        (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): GPT2Attention(\n",
       "          (c_attn): Conv1D()\n",
       "          (c_proj): Conv1D()\n",
       "          (attn_dropout): Dropout(p=0.1, inplace=False)\n",
       "          (resid_dropout): Dropout(p=0.1, inplace=False)\n",
       "        )\n",
       "        (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "        (mlp): GPT2MLP(\n",
       "          (c_fc): Conv1D()\n",
       "          (c_proj): Conv1D()\n",
       "          (act): NewGELUActivation()\n",
       "          (dropout): Dropout(p=0.1, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (10): GPT2Block(\n",
       "        (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): GPT2Attention(\n",
       "          (c_attn): Conv1D()\n",
       "          (c_proj): Conv1D()\n",
       "          (attn_dropout): Dropout(p=0.1, inplace=False)\n",
       "          (resid_dropout): Dropout(p=0.1, inplace=False)\n",
       "        )\n",
       "        (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "        (mlp): GPT2MLP(\n",
       "          (c_fc): Conv1D()\n",
       "          (c_proj): Conv1D()\n",
       "          (act): NewGELUActivation()\n",
       "          (dropout): Dropout(p=0.1, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (11): GPT2Block(\n",
       "        (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): GPT2Attention(\n",
       "          (c_attn): Conv1D()\n",
       "          (c_proj): Conv1D()\n",
       "          (attn_dropout): Dropout(p=0.1, inplace=False)\n",
       "          (resid_dropout): Dropout(p=0.1, inplace=False)\n",
       "        )\n",
       "        (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "        (mlp): GPT2MLP(\n",
       "          (c_fc): Conv1D()\n",
       "          (c_proj): Conv1D()\n",
       "          (act): NewGELUActivation()\n",
       "          (dropout): Dropout(p=0.1, inplace=False)\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "    (ln_f): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "  )\n",
       "  (lm_head): Linear(in_features=768, out_features=50257, bias=False)\n",
       ")"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# MODEL\n",
    "\n",
    "if args.model_name_or_path:\n",
    "    model = AutoModelForCausalLM.from_pretrained(\n",
    "        args.model_name_or_path,\n",
    "        from_tf=bool(\".ckpt\" in args.model_name_or_path),\n",
    "        config=config,\n",
    "    )\n",
    "else:\n",
    "    logger.info(\"Training new model from scratch\")\n",
    "    model = AutoModelForCausalLM.from_config(config)\n",
    "\n",
    "model.resize_token_embeddings(len(tokenizer))\n",
    "\n",
    "model.gradient_checkpointing_enable()\n",
    "model.to(device) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "encouraging-taylor",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-04-13T07:38:46.715450Z",
     "iopub.status.busy": "2023-04-13T07:38:46.714765Z",
     "iopub.status.idle": "2023-04-13T07:38:46.723309Z",
     "shell.execute_reply": "2023-04-13T07:38:46.721965Z",
     "shell.execute_reply.started": "2023-04-13T07:38:46.715408Z"
    },
    "id": "3IMjj23-scMp"
   },
   "outputs": [],
   "source": [
    "def extract_comment_from_prompt(text):\n",
    "    if type(text) == str:\n",
    "        starts = text.rfind(\"COMMENT: \") + len(\"COMMENT: \")\n",
    "        ends = text.find(\" <EOS>\")\n",
    "        result = text[starts:ends]\n",
    "    else:\n",
    "        text = pd.Series(text)\n",
    "        starts = text.str.find(\"COMMENT: \") + len(\"COMMENT: \")\n",
    "        ends = text.str.find(\" <EOS>\")\n",
    "        result = [sentence[start:end] for (sentence, start, end) in zip(text, starts, ends)]\n",
    "    return result\n",
    "\n",
    "def extract_masked_prompt(text):\n",
    "    text = pd.Series(text)\n",
    "    ends = text.str.find(\"COMMENT: \") + len(\"COMMENT: \")\n",
    "    result = [sentence[:end] for (sentence, end) in zip(text, ends)]\n",
    "    return result"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "spare-advice",
   "metadata": {
    "id": "thirty-chorus"
   },
   "source": [
    "# Обучение"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "endless-japanese",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-04-13T07:38:46.725811Z",
     "iopub.status.busy": "2023-04-13T07:38:46.725307Z",
     "iopub.status.idle": "2023-04-13T07:38:46.733001Z",
     "shell.execute_reply": "2023-04-13T07:38:46.731847Z",
     "shell.execute_reply.started": "2023-04-13T07:38:46.725741Z"
    },
    "id": "5K3sbNi2hd4k"
   },
   "outputs": [],
   "source": [
    "args.num_warmup_steps = 0\n",
    "args.learning_rate = 2e-6"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "pediatric-exception",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-04-13T07:38:46.735492Z",
     "iopub.status.busy": "2023-04-13T07:38:46.735044Z",
     "iopub.status.idle": "2023-04-13T07:38:46.762379Z",
     "shell.execute_reply": "2023-04-13T07:38:46.761446Z",
     "shell.execute_reply.started": "2023-04-13T07:38:46.735446Z"
    },
    "id": "vital-bunch"
   },
   "outputs": [],
   "source": [
    "#optimizer = torch.optim.Adam(model.parameters(), lr=args.learning_rate) # this crashes with out-of-memory error\n",
    "optimizer = bnb.optim.Adam8bit(model.parameters(), lr=args.learning_rate)\n",
    "\n",
    "lr_scheduler = get_scheduler(\n",
    "    name=args.lr_scheduler_type,\n",
    "    optimizer=optimizer,\n",
    "    num_warmup_steps=args.num_warmup_steps,\n",
    "    # num_training_steps=args.max_train_steps,\n",
    "    num_training_steps=-1,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "complimentary-explosion",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-04-13T07:38:46.764325Z",
     "iopub.status.busy": "2023-04-13T07:38:46.763773Z",
     "iopub.status.idle": "2023-04-13T07:38:46.768892Z",
     "shell.execute_reply": "2023-04-13T07:38:46.767819Z",
     "shell.execute_reply.started": "2023-04-13T07:38:46.764284Z"
    },
    "id": "verbal-amendment"
   },
   "outputs": [],
   "source": [
    "output_dir = \"./chkp\"\n",
    "output_prefix = \"GPT2_02\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "engaging-dining",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-04-13T07:38:46.771129Z",
     "iopub.status.busy": "2023-04-13T07:38:46.770425Z",
     "iopub.status.idle": "2023-04-13T07:38:46.778278Z",
     "shell.execute_reply": "2023-04-13T07:38:46.777148Z",
     "shell.execute_reply.started": "2023-04-13T07:38:46.771088Z"
    },
    "id": "-7ywhSxJidDn"
   },
   "outputs": [],
   "source": [
    "def pack_tensor(new_tensor, packed_tensor, max_seq_len=1024):\n",
    "    if packed_tensor is None:\n",
    "        return new_tensor, True, None\n",
    "    if new_tensor.size()[1] + packed_tensor.size()[1] > max_seq_len:\n",
    "        return packed_tensor, False, new_tensor\n",
    "    else:\n",
    "        packed_tensor = torch.cat([new_tensor, packed_tensor[:, 1:]], dim=1)\n",
    "        return packed_tensor, True, None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "absent-graduate",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-04-13T07:38:46.780542Z",
     "iopub.status.busy": "2023-04-13T07:38:46.779902Z",
     "iopub.status.idle": "2023-04-13T07:38:46.796166Z",
     "shell.execute_reply": "2023-04-13T07:38:46.795026Z",
     "shell.execute_reply.started": "2023-04-13T07:38:46.780503Z"
    },
    "id": "uAf84EaGgcth"
   },
   "outputs": [],
   "source": [
    "def train(device, train_dataloader, model, tokenizer, \n",
    "          optimizer, scheduler=lr_scheduler,\n",
    "          batch_size=16, epochs=5, \n",
    "          # lr=2e-5, max_seq_len=400, warmup_steps=200,\n",
    "          gpt2_type=\"gpt2\", output_dir=\".\", output_prefix=output_prefix,\n",
    "          test_mode=False,save_model_on_epoch=True, save_every_steps=50000\n",
    "):\n",
    "    acc_steps = 100\n",
    "    model = model.to(device)\n",
    "    model.train()\n",
    "\n",
    "    # optimizer = AdamW(model.parameters(), lr=lr)\n",
    "    # scheduler = get_linear_schedule_with_warmup(\n",
    "    #     optimizer, num_warmup_steps=warmup_steps, num_training_steps=-1\n",
    "    # )\n",
    "\n",
    "    # train_dataloader = DataLoader(dataset, batch_size=1, shuffle=True)\n",
    "    loss=0\n",
    "    accumulating_batch_count = 0\n",
    "    input_tensor = None\n",
    "\n",
    "    for epoch in range(epochs):\n",
    "        losses = []\n",
    "        print(f\"Training epoch {epoch}\")\n",
    "        progress_bar = tqdm(enumerate(train_dataloader), total=len(train_dataloader))\n",
    "        for idx, entry in progress_bar:\n",
    "            progress_bar.set_description(\n",
    "                f\"Epoch[{epoch}/{epochs - 1}]Step[{idx}/{len(train_dataloader) - 1}]\"\n",
    "            )\n",
    "            progress_bar.update(1)\n",
    "\n",
    "            if idx % save_every_steps == 0 and idx != 0:\n",
    "                torch.save(\n",
    "                    model.state_dict(),\n",
    "                    os.path.join(output_dir, f\"{output_prefix}_Ep{epoch}_St{idx}.pt\"),\n",
    "                )\n",
    "            ####################################################################\n",
    "            (input_tensor, carry_on, remainder) = pack_tensor(entry, input_tensor)\n",
    "\n",
    "            if carry_on and idx != len(train_dataloader) - 1:\n",
    "                continue\n",
    "\n",
    "            input_tensor = input_tensor.to(device)\n",
    "            outputs = model(input_tensor, labels=input_tensor)\n",
    "            loss = outputs[0]\n",
    "            losses.append(loss.item())\n",
    "            loss.backward()\n",
    "\n",
    "            if (accumulating_batch_count % batch_size) == 0:\n",
    "                optimizer.step()\n",
    "                scheduler.step()\n",
    "                optimizer.zero_grad()\n",
    "                model.zero_grad()\n",
    "            \n",
    "            accumulating_batch_count += 1\n",
    "            input_tensor = None\n",
    "            ####################################################################\n",
    "            if idx % args.log_loss_interval == 0 and idx > 0:\n",
    "                try:\n",
    "                    perplexity = math.exp(sum(losses)/len(losses))\n",
    "                except OverflowError:\n",
    "                    perplexity = float(\"inf\")\n",
    "                losses = []\n",
    "                print(f\"epoch: {epoch+1}, step: {idx}, perplexity: {perplexity}\")\n",
    "            ####################################################################\n",
    "        if save_model_on_epoch:\n",
    "            torch.save(\n",
    "                model.state_dict(),\n",
    "                os.path.join(output_dir, f\"{output_prefix}_Ep{epoch}.pt\"),\n",
    "            )\n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "lovely-literature",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "collapsed": true,
    "execution": {
     "iopub.execute_input": "2023-04-13T07:38:46.801378Z",
     "iopub.status.busy": "2023-04-13T07:38:46.800355Z"
    },
    "id": "W3Ri0NNCg4Qe",
    "outputId": "02af296b-07ac-401d-ba65-d85415112aa2"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training epoch 0\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a7f2d532e09843febf5684c5e9286649",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/476955 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`...\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 1, step: 50, perplexity: 45.25104748121019\n",
      "epoch: 1, step: 1675, perplexity: 46.89404939629187\n",
      "epoch: 1, step: 1775, perplexity: 45.67016438820054\n",
      "epoch: 1, step: 4425, perplexity: 46.27442498457143\n",
      "epoch: 1, step: 4675, perplexity: 46.33711766792145\n",
      "epoch: 1, step: 4700, perplexity: 47.29383675330516\n",
      "epoch: 1, step: 4725, perplexity: 41.434275857784186\n",
      "epoch: 1, step: 5350, perplexity: 44.818841583009934\n",
      "epoch: 1, step: 5400, perplexity: 44.04448381574429\n",
      "epoch: 1, step: 5450, perplexity: 41.01916698574974\n",
      "epoch: 1, step: 5475, perplexity: 36.76212514732277\n",
      "epoch: 1, step: 5525, perplexity: 42.98661102590692\n",
      "epoch: 1, step: 5800, perplexity: 46.79927847203265\n",
      "epoch: 1, step: 5825, perplexity: 43.36244377499858\n",
      "epoch: 1, step: 5850, perplexity: 44.43151998738242\n",
      "epoch: 1, step: 6175, perplexity: 46.694394016233396\n",
      "epoch: 1, step: 7650, perplexity: 47.71399650477441\n",
      "epoch: 1, step: 7675, perplexity: 42.552602008841475\n",
      "epoch: 1, step: 7700, perplexity: 38.91476936206802\n",
      "epoch: 1, step: 7725, perplexity: 42.29166594321123\n",
      "epoch: 1, step: 7775, perplexity: 44.5169628468203\n",
      "epoch: 1, step: 7800, perplexity: 45.2317937660302\n",
      "epoch: 1, step: 8975, perplexity: 45.83914679702818\n",
      "epoch: 1, step: 9000, perplexity: 54.94590120773058\n",
      "epoch: 1, step: 9100, perplexity: 44.75679894046579\n",
      "epoch: 1, step: 9650, perplexity: 44.605178824687464\n",
      "epoch: 1, step: 9750, perplexity: 46.31359605524846\n",
      "epoch: 1, step: 9775, perplexity: 45.88498867974935\n",
      "epoch: 1, step: 10825, perplexity: 45.837629851187515\n",
      "epoch: 1, step: 10900, perplexity: 51.73074640904678\n",
      "epoch: 1, step: 10925, perplexity: 49.43539737819579\n",
      "epoch: 1, step: 11550, perplexity: 47.84605454193307\n",
      "epoch: 1, step: 14150, perplexity: 46.0390291611813\n",
      "epoch: 1, step: 14175, perplexity: 43.177389334476324\n",
      "epoch: 1, step: 14225, perplexity: 43.57114970995617\n",
      "epoch: 1, step: 23000, perplexity: 46.22528685145157\n",
      "epoch: 1, step: 23375, perplexity: 44.67609716750473\n",
      "epoch: 1, step: 23450, perplexity: 40.41546671233162\n",
      "epoch: 1, step: 25150, perplexity: 46.345488908857064\n",
      "epoch: 1, step: 26550, perplexity: 46.66293801873785\n",
      "epoch: 1, step: 26675, perplexity: 43.30410204412815\n",
      "epoch: 1, step: 26825, perplexity: 47.35319370352264\n",
      "epoch: 1, step: 27950, perplexity: 46.318981874455005\n",
      "epoch: 1, step: 27975, perplexity: 47.35065542994076\n",
      "epoch: 1, step: 28000, perplexity: 66.89805026528835\n",
      "epoch: 1, step: 30475, perplexity: 46.26360236792797\n",
      "epoch: 1, step: 31075, perplexity: 44.91654366976862\n",
      "epoch: 1, step: 31125, perplexity: 50.79035528811557\n"
     ]
    }
   ],
   "source": [
    "model = train(device, train_dataloader, model, \n",
    "              tokenizer, optimizer, lr_scheduler,\n",
    "              epochs=1, gpt2_type=args.model_name_or_path, \n",
    "              output_dir=output_dir,save_every_steps=50000,\n",
    "              output_prefix=output_prefix)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "educated-williams",
   "metadata": {
    "id": "hjrG609ieqHT"
   },
   "source": [
    "# Генерация"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "overhead-webcam",
   "metadata": {
    "id": "8OAiLKEXa67v"
   },
   "outputs": [],
   "source": [
    "def model_loader(model, device, CHECKPOINTS_DIR, WEIGHTS_NAME):\n",
    "    model_file = os.path.join(CHECKPOINTS_DIR, WEIGHTS_NAME)\n",
    "    try:\n",
    "        checkpoint = torch.load(model_file)\n",
    "    except:\n",
    "        checkpoint = torch.load(model_file, map_location=device)\n",
    "\n",
    "    model.load_state_dict(checkpoint)\n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "ordinary-gateway",
   "metadata": {
    "id": "Ts9i32Xme2V3"
   },
   "outputs": [],
   "source": [
    "from tqdm import trange\n",
    "import torch.nn.functional as F\n",
    "from IPython.display import clear_output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "metropolitan-minimum",
   "metadata": {
    "id": "ARCkBDc4eoS2"
   },
   "outputs": [],
   "source": [
    "def generate(model, tokenizer, prompt,\n",
    "             entry_count=10, entry_length=30, top_p=0.8, temperature=1.):\n",
    "    model.eval()\n",
    "    generated_num = 0\n",
    "    generated_list = []\n",
    "    filter_value = -float(\"Inf\")\n",
    "    with torch.no_grad():\n",
    "        for entry_idx in trange(entry_count):\n",
    "            entry_finished = False\n",
    "            generated = torch.tensor(tokenizer.encode(prompt)).unsqueeze(0)\n",
    "            for i in range(entry_length):\n",
    "                outputs = model(generated, labels=generated)\n",
    "                loss, logits = outputs[:2]\n",
    "                logits = logits[:, -1, :] / (temperature if temperature > 0 else 1.0)\n",
    "\n",
    "                sorted_logits, sorted_indices = torch.sort(logits, descending=True)\n",
    "                cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)\n",
    "\n",
    "                sorted_indices_to_remove = cumulative_probs > top_p\n",
    "                sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()\n",
    "                sorted_indices_to_remove[..., 0] = 0\n",
    "\n",
    "                indices_to_remove = sorted_indices[sorted_indices_to_remove]\n",
    "                logits[:, indices_to_remove] = filter_value\n",
    "\n",
    "                next_token = torch.multinomial(F.softmax(logits, dim=-1), num_samples=1)\n",
    "                generated = torch.cat((generated, next_token), dim=1)\n",
    "                if next_token in tokenizer.encode(\"<EOS>\"):\n",
    "                    entry_finished = True\n",
    "                if entry_finished:\n",
    "                    generated_num = generated_num + 1\n",
    "                    output_list = list(generated.squeeze().numpy())\n",
    "                    output_text = tokenizer.decode(output_list)\n",
    "                    generated_list.append(output_text)\n",
    "                    break\n",
    "            if not entry_finished:\n",
    "                output_list = list(generated.squeeze().numpy())\n",
    "                output_text = f\"{tokenizer.decode(output_list)} <EOS>\" \n",
    "                generated_list.append(output_text)\n",
    "    return generated_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "deadly-bridge",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'Sexiest Male Vocalist Riff-Off w/ Usher & Luke Evans'"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ts = train_dataset.df_full.iloc[1000]\n",
    "ts.title"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "cloudy-consequence",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'<BOS> TOPIC: Sexiest Male Vocalist Riff-Off w/ Usher & Luke Evans COMMENT: '"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "prompt = extract_masked_prompt(ts.prompt)[0]\n",
    "prompt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "approximate-canadian",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "collapsed": true,
    "id": "utURdN8mfER4",
    "outputId": "e0a1a561-ab19-4703-ea20-6d1900f791d1"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 0/3 [00:15<?, ?it/s]\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-37-e311e7898236>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      3\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mname\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mnames\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      4\u001b[0m     \u001b[0mmodel\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmodel_loader\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmodel\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdevice\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m\"./chkp\"\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34mf\"GPT2_02_Ep0_St{name}.pt\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 5\u001b[0;31m     \u001b[0mres\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mgenerate\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmodel\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mto\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'cpu'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtokenizer\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mprompt\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mentry_count\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m3\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      6\u001b[0m     \u001b[0;32mfor\u001b[0m \u001b[0mi\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mextract_comment_from_prompt\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mres\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      7\u001b[0m         \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m<ipython-input-15-3bb6b764d281>\u001b[0m in \u001b[0;36mgenerate\u001b[0;34m(model, tokenizer, prompt, entry_count, entry_length, top_p, temperature)\u001b[0m\n\u001b[1;32m     10\u001b[0m             \u001b[0mgenerated\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtensor\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtokenizer\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mencode\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mprompt\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0munsqueeze\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     11\u001b[0m             \u001b[0;32mfor\u001b[0m \u001b[0mi\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mrange\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mentry_length\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 12\u001b[0;31m                 \u001b[0moutputs\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmodel\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mgenerated\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mlabels\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mgenerated\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     13\u001b[0m                 \u001b[0mloss\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mlogits\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0moutputs\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;36m2\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     14\u001b[0m                 \u001b[0mlogits\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mlogits\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m-\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m:\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m/\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0mtemperature\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0mtemperature\u001b[0m \u001b[0;34m>\u001b[0m \u001b[0;36m0\u001b[0m \u001b[0;32melse\u001b[0m \u001b[0;36m1.0\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.8/site-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_call_impl\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m   1188\u001b[0m         if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks\n\u001b[1;32m   1189\u001b[0m                 or _global_forward_hooks or _global_forward_pre_hooks):\n\u001b[0;32m-> 1190\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mforward_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1191\u001b[0m         \u001b[0;31m# Do not call functions when jit is used\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1192\u001b[0m         \u001b[0mfull_backward_hooks\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnon_full_backward_hooks\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.8/dist-packages/transformers/models/gpt2/modeling_gpt2.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, input_ids, past_key_values, attention_mask, token_type_ids, position_ids, head_mask, inputs_embeds, encoder_hidden_states, encoder_attention_mask, labels, use_cache, output_attentions, output_hidden_states, return_dict)\u001b[0m\n\u001b[1;32m   1045\u001b[0m         \u001b[0mreturn_dict\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mreturn_dict\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0mreturn_dict\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0;32mNone\u001b[0m \u001b[0;32melse\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mconfig\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0muse_return_dict\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1046\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1047\u001b[0;31m         transformer_outputs = self.transformer(\n\u001b[0m\u001b[1;32m   1048\u001b[0m             \u001b[0minput_ids\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1049\u001b[0m             \u001b[0mpast_key_values\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mpast_key_values\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.8/site-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_call_impl\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m   1188\u001b[0m         if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks\n\u001b[1;32m   1189\u001b[0m                 or _global_forward_hooks or _global_forward_pre_hooks):\n\u001b[0;32m-> 1190\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mforward_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1191\u001b[0m         \u001b[0;31m# Do not call functions when jit is used\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1192\u001b[0m         \u001b[0mfull_backward_hooks\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnon_full_backward_hooks\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.8/dist-packages/transformers/models/gpt2/modeling_gpt2.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, input_ids, past_key_values, attention_mask, token_type_ids, position_ids, head_mask, inputs_embeds, encoder_hidden_states, encoder_attention_mask, use_cache, output_attentions, output_hidden_states, return_dict)\u001b[0m\n\u001b[1;32m    888\u001b[0m                 )\n\u001b[1;32m    889\u001b[0m             \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 890\u001b[0;31m                 outputs = block(\n\u001b[0m\u001b[1;32m    891\u001b[0m                     \u001b[0mhidden_states\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    892\u001b[0m                     \u001b[0mlayer_past\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mlayer_past\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.8/site-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_call_impl\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m   1188\u001b[0m         if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks\n\u001b[1;32m   1189\u001b[0m                 or _global_forward_hooks or _global_forward_pre_hooks):\n\u001b[0;32m-> 1190\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mforward_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1191\u001b[0m         \u001b[0;31m# Do not call functions when jit is used\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1192\u001b[0m         \u001b[0mfull_backward_hooks\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnon_full_backward_hooks\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.8/dist-packages/transformers/models/gpt2/modeling_gpt2.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, hidden_states, layer_past, attention_mask, head_mask, encoder_hidden_states, encoder_attention_mask, use_cache, output_attentions)\u001b[0m\n\u001b[1;32m    393\u001b[0m         \u001b[0mresidual\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mhidden_states\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    394\u001b[0m         \u001b[0mhidden_states\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mln_1\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mhidden_states\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 395\u001b[0;31m         attn_outputs = self.attn(\n\u001b[0m\u001b[1;32m    396\u001b[0m             \u001b[0mhidden_states\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    397\u001b[0m             \u001b[0mlayer_past\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mlayer_past\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.8/site-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_call_impl\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m   1188\u001b[0m         if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks\n\u001b[1;32m   1189\u001b[0m                 or _global_forward_hooks or _global_forward_pre_hooks):\n\u001b[0;32m-> 1190\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mforward_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1191\u001b[0m         \u001b[0;31m# Do not call functions when jit is used\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1192\u001b[0m         \u001b[0mfull_backward_hooks\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnon_full_backward_hooks\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.8/dist-packages/transformers/models/gpt2/modeling_gpt2.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, hidden_states, layer_past, attention_mask, head_mask, encoder_hidden_states, encoder_attention_mask, use_cache, output_attentions)\u001b[0m\n\u001b[1;32m    334\u001b[0m             \u001b[0mattn_output\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mattn_weights\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_upcast_and_reordered_attn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mquery\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mkey\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mvalue\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mattention_mask\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mhead_mask\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    335\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 336\u001b[0;31m             \u001b[0mattn_output\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mattn_weights\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_attn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mquery\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mkey\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mvalue\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mattention_mask\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mhead_mask\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    337\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    338\u001b[0m         \u001b[0mattn_output\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_merge_heads\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mattn_output\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnum_heads\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mhead_dim\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.8/dist-packages/transformers/models/gpt2/modeling_gpt2.py\u001b[0m in \u001b[0;36m_attn\u001b[0;34m(self, query, key, value, attention_mask, head_mask)\u001b[0m\n\u001b[1;32m    220\u001b[0m             \u001b[0mattn_weights\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mattn_weights\u001b[0m \u001b[0;34m*\u001b[0m \u001b[0mhead_mask\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    221\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 222\u001b[0;31m         \u001b[0mattn_output\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmatmul\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mattn_weights\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mvalue\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    223\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    224\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0mattn_output\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mattn_weights\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "names = [\"50000\", \"100000\", \"200000\", \"300000\", \"400000\"]\n",
    "\n",
    "for name in names:\n",
    "    model = model_loader(model, device, \"./chkp\", f\"GPT2_02_Ep0_St{name}.pt\")    \n",
    "    res = generate(model.to('cpu'), tokenizer, prompt, entry_count=3)\n",
    "    for i in extract_comment_from_prompt(res):\n",
    "        print()\n",
    "        print(i)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "chronic-clinic",
   "metadata": {
    "id": "Wx8AUo92jhnT"
   },
   "source": [
    "# Evaluation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "lucky-trash",
   "metadata": {},
   "outputs": [],
   "source": [
    "def text_generation(test_data):\n",
    "    generated_lyrics = []\n",
    "    for i in trange(len(test_data)):\n",
    "        x = generate(model.to('cpu'), tokenizer, test_data[i], entry_count=1)\n",
    "        generated_lyrics.append(x)\n",
    "        clear_output(wait=False)\n",
    "    return generated_lyrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "contrary-contact",
   "metadata": {
    "id": "5NSIYz2hkA19"
   },
   "outputs": [],
   "source": [
    "test_dataset = Comments(\"prompts_test.csv\", \n",
    "                      text_column_name=args.text_column_name, \n",
    "                      tokenizer=tokenizer,\n",
    "                      block_size=block_size)\n",
    "# test_dataloader = DataLoader(\n",
    "#     test_dataset, batch_size=args.per_device_train_batch_size, shuffle=False\n",
    "# )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "formal-biology",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "qkRlPEpYjg-v",
    "outputId": "eca98bc9-3d58-49b6-c21c-9442335e054e"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "100%|██████████| 10/10 [01:29<00:00,  8.98s/it]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "0"
      ]
     },
     "execution_count": 138,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import statistics\n",
    "from nltk.translate.bleu_score import sentence_bleu\n",
    "\n",
    "df_test = test_dataset.df_full.iloc[:10]\n",
    "\n",
    "scores=[]\n",
    "true_coms = extract_comment_from_prompt(df_test.prompt)\n",
    "test_prompts = extract_masked_prompt(df_test.prompt)\n",
    "fake_coms = text_generation(test_prompts)\n",
    "fake_coms = [com[0] for com in fake_coms]\n",
    "fake_coms = extract_comment_from_prompt(pd.Series(fake_coms))\n",
    "\n",
    "for i in range(len(true_coms)):\n",
    "  reference = true_coms[i]\n",
    "  candidate = fake_coms[i]\n",
    "  scores.append(sentence_bleu(reference, candidate))\n",
    "\n",
    "\n",
    "statistics.mean(scores)"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "provenance": [],
   "toc_visible": true
  },
  "gpuClass": "standard",
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}