fine_tuning.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import math\n",
    "\n",
    "from pathlib import Path\n",
    "from typing import Dict, List\n",
    "from datasets import load_dataset\n",
    "from huggingface_hub import notebook_login\n",
    "from transformers import AutoTokenizer, AutoModelWithLMHead, AutoModelForCausalLM\n",
    "from transformers import TrainingArguments, Trainer\n",
    "from transformers import DataCollatorForLanguageModeling\n",
    "from transformers import pipeline\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# https://huggingface.co/docs/transformers/tasks/language_modeling\n",
    "# https://huggingface.co/tinkoff-ai/ruDialoGPT-medium\n",
    "# https://github.com/tinkoff-ai/pycon-chit-chat/blob/main/notebooks/lm_training.ipynb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "notebook_login()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Data Preparation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python prepare_messages.py --tg-history-path \"data/result.json\" --output-path \"data/data.csv\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = load_dataset(\"csv\", data_files=\"data/data.csv\", split=\"train\")\n",
    "data = data.filter(lambda example: example[\"context_1\"] != None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "FIRST_SPEAKER_TOKEN = \"@@ПЕРВЫЙ@@\"\n",
    "SECOND_SPEAKER_TOKEN = \"@@ВТОРОЙ@@\"\n",
    "CONTEXT_COLS = [\"context_3\", \"context_2\", \"context_1\"]\n",
    "RESPONSE_COL = [\"response\"]\n",
    "SEP = \" \"\n",
    "\n",
    "\n",
    "def convert_to_dialog(sample: Dict[str, str]) -> Dict[str, str]:\n",
    "    result_dict = dict()\n",
    "    dialog = \"\"\n",
    "    for i in range(len(CONTEXT_COLS)):\n",
    "        key = CONTEXT_COLS[i]\n",
    "        if key in sample and sample[key] is not None:\n",
    "            speaker_token = FIRST_SPEAKER_TOKEN if i % 2 == 0 else SECOND_SPEAKER_TOKEN\n",
    "            dialog += speaker_token + SEP + sample[key] + SEP\n",
    "    \n",
    "    response_key = RESPONSE_COL[0]\n",
    "    if response_key in sample and sample[response_key] is not None:\n",
    "        dialog += SECOND_SPEAKER_TOKEN + SEP + sample[response_key]\n",
    "    result_dict['text'] = dialog\n",
    "    return result_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = data.train_test_split(test_size=0.2, shuffle=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dialog_data = data.map(convert_to_dialog)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer = AutoTokenizer.from_pretrained(\"tinkoff-ai/ruDialoGPT-medium\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def preprocess_function(examples):\n",
    "    return tokenizer(examples[\"text\"], truncation=True, padding=\"max_length\", max_length=128)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenized_data = dialog_data.map(\n",
    "    preprocess_function,\n",
    "    batched=True,\n",
    "    num_proc=4,\n",
    "    remove_columns=data[\"train\"].column_names,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer.pad_token = tokenizer.eos_token\n",
    "data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "training_args = TrainingArguments(\n",
    "    output_dir=\"model\",\n",
    "    evaluation_strategy=\"epoch\",\n",
    "    learning_rate=2e-5,\n",
    "    weight_decay=0.01,\n",
    "    push_to_hub=True,\n",
    "    per_device_train_batch_size=32,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "model = AutoModelForCausalLM.from_pretrained(\"tinkoff-ai/ruDialoGPT-medium\").to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer = Trainer(\n",
    "    model=model,\n",
    "    args=training_args,\n",
    "    train_dataset=tokenized_data[\"train\"],\n",
    "    eval_dataset=tokenized_data[\"test\"],\n",
    "    data_collator=data_collator,\n",
    ")\n",
    "\n",
    "trainer.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "eval_results = trainer.evaluate()\n",
    "print(f\"Perplexity: {math.exp(eval_results['eval_loss']):.2f}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.push_to_hub()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "checkpoint_path = \"PATH TO CHECKPOINT\"\n",
    "\n",
    "model = AutoModelWithLMHead.from_pretrained(checkpoint_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def process_response(response):\n",
    "    response = response.capitalize()\n",
    "    response = response.replace(\"\\n\", \" \")\n",
    "    return response\n",
    "\n",
    "\n",
    "def print_dialogue(prompt, response):\n",
    "    print(\"— \" + prompt)\n",
    "    print(\"— \" + response)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt = \"YOUR PROMPT\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "full_prompt = f'@@ПЕРВЫЙ@@ привет @@ВТОРОЙ@@ привет @@ПЕРВЫЙ@@ {prompt} @@ВТОРОЙ@@'\n",
    "inputs = tokenizer(full_prompt, return_tensors='pt')\n",
    "generated_token_ids = model.generate(\n",
    "    **inputs,\n",
    "    top_k=10,\n",
    "    top_p=0.95,\n",
    "    num_beams=3,\n",
    "    num_return_sequences=3,\n",
    "    do_sample=True,\n",
    "    no_repeat_ngram_size=2,\n",
    "    temperature=1.2,\n",
    "    repetition_penalty=1.2,\n",
    "    length_penalty=1.0,\n",
    "    eos_token_id=50257,\n",
    "    max_new_tokens=40\n",
    ")\n",
    "context_with_response = [tokenizer.decode(sample_token_ids) for sample_token_ids in generated_token_ids]\n",
    "\n",
    "response = context_with_response[2][len(full_prompt)+1:]\n",
    "response = response[:response.find(\"<pad>\")]\n",
    "response = process_response(response)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print_dialogue(prompt, response)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.9.12 ('base')",
   "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.9.13"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "40d3a090f54c6569ab1632332b64b2c03c39dcf918b08424e98f38b5ae0af88f"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}