training.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "source": [
    "# Дообучение модели на собственных данных для задачи генерации текста на\n",
    "русском языке"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "Ноутбук подходит для запуска на бесплатном тарифе в Google Colab. Перед\n",
    "запуском ячеек вы должны располагать набором данных (`corpus.db` или\n",
    "`corpus_small.db`. Данные можно получить, запустив `grabber.py`."
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "Доустанавливаем нужные библиотеки"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "!pip install transformers datasets evaluate"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "Необходимые импорты"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "import math\n",
    "\n",
    "from datasets import Dataset\n",
    "\n",
    "from transformers import AutoTokenizer, AutoModelForCausalLM, Trainer, TrainingArguments\n",
    "from huggingface_hub import notebook_login"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "Необходимо указать токен 🤗, чтобы загрузить обученную модель."
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "notebook_login()"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "Пожалуй, самую большую известность получили модели для русского языка от\n",
    "Сбера. Её малый вариант и будем дотренировывать - модели большего размера\n",
    "уже не помещаются в Colab."
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "model_checkpoint = \"sberbank-ai/rugpt3small_based_on_gpt2\""
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "Готовим данные"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "data = Dataset.from_sql(\"SELECT text FROM quotes;\",\n",
    "                        con=\"sqlite:///corpus_small.db\")\n",
    "data = data.train_test_split(test_size=0.25)"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "Перед каждой цитатой добавим затравку: \"Преподаватель говорит: \". Это\n",
    "позволит нам генерировать случайные цитаты без привязки к конкретному\n",
    "инпуту, а также разграничивать сэмплы."
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "data = data.map(lambda x: {\"text\":\n",
    "                               f\"Преподаватель  говорит: \"\n",
    "                               f\"{x['text']}<|endoftext|>\"})"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "Токенизируем и разбиваем на блоки"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "def preprocess_function(examples):\n",
    "    return tokenizer(examples['text'])"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "data = data.map(preprocess_function, batched=True, remove_columns=[\"text\"])"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "В идеале, текст надо делить на блоки, длина которых соответствует длине\n",
    "блоков, на которых была обучена исходная модель. Однако слишком длинные\n",
    "блоки переполняют оперативную память. В бесплатном Colab не помещаются блоки\n",
    " длиннее 512 токенов."
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "block_size = 512\n",
    "\n",
    "def group_texts(examples):\n",
    "    # Concatenate all texts.\n",
    "    concatenated_examples = {k: sum(examples[k], []) for k in examples.keys()}\n",
    "    total_length = len(concatenated_examples[list(examples.keys())[0]])\n",
    "    total_length = (total_length // block_size) * block_size\n",
    "    result = {\n",
    "        k: [t[i : i + block_size] for i in range(0, total_length, block_size)]\n",
    "        for k, t in concatenated_examples.items()\n",
    "    }\n",
    "    result[\"labels\"] = result[\"input_ids\"].copy()\n",
    "    return result"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "data = data.map(group_texts, batched=True)\n",
    "print(data)"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "Загружаем исходную модель"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "model = AutoModelForCausalLM.from_pretrained(model_checkpoint)"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "Задаём имя нашей новой модели и гиперпараметры обучения"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "model_name = f\"{model_checkpoint.split('/')[-1]}-finetuned_teachers_quotes\"\n",
    "\n",
    "training_args = TrainingArguments(\n",
    "    model_name,\n",
    "    evaluation_strategy = \"epoch\",\n",
    "    learning_rate=2e-5,\n",
    "    weight_decay=0.01,\n",
    "    push_to_hub=True\n",
    ")"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "new_model = Trainer(\n",
    "    model=model,\n",
    "    args=training_args,\n",
    "    train_dataset=data[\"train\"],\n",
    "    eval_dataset=data[\"test\"]\n",
    ")"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "Ура, тренируем!"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "new_model.train()"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "Посмотрим на циферки"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "eval_results = new_model.evaluate()\n",
    "print(f\"Perplexity: {math.exp(eval_results['eval_loss']):.2f}\")"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "Perplexity в районе 18 - очень неплохой результат, учитывая неаккуратность\n",
    "наших данных и сравнительно небольшой размер."
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "Выгружаем обученную модель и токенизатор на 🤗. Теперь её можно будет\n",
    "использовать через API, не загружая локально (там больше 500 МБ)."
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "new_model.push_to_hub()\n",
    "tokenizer.push_to_hub(f\"Futyn-Maker/{model_name}\")\n"
   ],
   "metadata": {
    "collapsed": false
   }
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}