{ "nbformat": 4, "nbformat_minor": 0, "metadata": { "colab": { "name": "Fine tune RuDialoGPT3 on telegram chat", "provenance": [], "collapsed_sections": [ "uPZXtklAd0Cd", "ESogNuUOEmj_", "psXZnJk0Eo3J" ], "toc_visible": true, "include_colab_link": true }, "kernelspec": { "display_name": "Python 3", "name": "python3" }, "language_info": { "name": "python" } }, "cells": [ { "cell_type": "markdown", "metadata": { "id": "view-in-github", "colab_type": "text" }, "source": [ "\"Open" ] }, { "cell_type": "markdown", "metadata": { "id": "ocoQoLlek3cb" }, "source": [ "# Fine-Tuning DialoGPT3 on your telegram chat" ] }, { "cell_type": "markdown", "metadata": { "id": "_ptkarFllCDr" }, "source": [ "Here is a ready-to-run code for fine-tuning a RuDialoGPT3 model using HuggingFace and PyTorch on **your telegram chat**.\n", "\n", "I used RuDialoGPT-3 trained on forums to fine tune. It was trained by [@Grossmend](https://github.com/Grossmend) on Russian forums. The training process took 12 days using 4x RTX 2080 Ti (2 epochs on 32GB text corpus). The training procedure of GPT-3 for dialogue is described in Grossmend's [blogpost](https://habr.com/ru/company/icl_services/blog/548244/) (in Russian).\n", "\n", "I have created a simple pipeline and fine tuned that model on my own exported telegram chat (~30mb json). It is in fact very easy to get the data from telegram and fine tune a model. Therefore, I made this notebook!" ] }, { "cell_type": "markdown", "metadata": { "id": "GAB9ev-Gd8lH" }, "source": [ "If you want just to try / to talk to my fine-tuned model than go **straight to the Inference section**." ] }, { "cell_type": "markdown", "metadata": { "id": "uPZXtklAd0Cd" }, "source": [ "## Uploading your data for fine-tuning" ] }, { "cell_type": "code", "metadata": { "id": "VL5BXKmva2-Q" }, "source": [ "# installing huggingface datasets and accelerate \n", "! pip install datasets transformers[sentencepiece]\n", "! pip install accelerate\n", "\n", "# [optional] Login to google drive to save models\n", "from google.colab import drive\n", "drive.mount('/content/drive')\n", "\n", "# [optional] Login to wandb to track model's behaviour\n", "'''! pip install wandb\n", "! wandb login\n", "wandb.init(project=\"fine tune RuDialoGPT2 on KirArChat\")'''" ], "execution_count": null, "outputs": [] }, { "cell_type": "code", "metadata": { "cellView": "form", "id": "Iq78W4qhrYmN" }, "source": [ "#@title Imports\n", "import sys\n", "import re\n", "import json\n", "\n", "from sklearn.model_selection import train_test_split\n", "from tqdm import tqdm\n", "\n", "import torch\n", "from transformers import TextDataset, DataCollatorForLanguageModeling\n", "from torch.utils.data import DataLoader\n", "\n", "from accelerate import Accelerator\n", "from transformers import AdamW, AutoModelForSequenceClassification, get_scheduler" ], "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "metadata": { "id": "7fRNBMkYnAUV" }, "source": [ "Next cell downloads model and tokenizer using HuggingFace.\n", "\n", "You can start with my version or @Grossmend's: \"Grossmend/rudialogpt3_medium_based_on_gpt2\". Moreover, you can even start with any different DialoGPT trained on your language (with the notation of |x|y|text)." ] }, { "cell_type": "code", "metadata": { "id": "fn9KxEnfaxwo" }, "source": [ "from transformers import AutoModelForCausalLM, AutoTokenizer\n", "\n", "checkpoint = \"Kirili4ik/ruDialoGpt3-medium-finetuned-telegram\" \n", "tokenizer = AutoTokenizer.from_pretrained(checkpoint)\n", "model = AutoModelForCausalLM.from_pretrained(checkpoint)" ], "execution_count": null, "outputs": [] }, { "cell_type": "code", "metadata": { "id": "SulpoPQxpJrK", "cellView": "form" }, "source": [ "#@title Utility functions\n", "def get_length_param(text: str, tokenizer) -> str:\n", " \"\"\"Maps text to 1 of 4 buckets based on length after encoding.\n", "\n", " Parameters\n", " ----------\n", " text: str\n", " The text to be given 1 of 4 length parameters.\n", "\n", " tokenizer: HuggingFace tokenizer \n", " Tokenizer that used to compute the length of the text after encoding.\n", " For more info ee https://huggingface.co/transformers/main_classes/tokenizer.html\n", "\n", " Returns\n", " -------\n", " len_param: str\n", " One of four buckets: \n", " '1' for short, '2' for medium, '3' for long texts and '-' for all others. \n", " \"\"\"\n", " tokens_count = len(tokenizer.encode(text))\n", " if tokens_count <= 15:\n", " len_param = '1'\n", " elif tokens_count <= 50:\n", " len_param = '2'\n", " elif tokens_count <= 256:\n", " len_param = '3'\n", " else:\n", " len_param = '-'\n", " return len_param\n", "\n", "\n", "def get_user_param(text: dict, machine_name_in_chat: str) -> str:\n", " \"\"\"Maps text by 1/0 for it to be the person or the machine in the dialog\n", "\n", " Parameters\n", " ----------\n", " text: Dict[..., 'from', ...]\n", " Dict containing field 'from' with the name of the user who sent the message\n", "\n", " machine_name_in_chat: str\n", " Str with the name of the machine - it will be predicted\n", " \"\"\"\n", " if text['from'] == machine_name_in_chat:\n", " return '1' # machine\n", " else:\n", " return '0' # human\n", "\n", "\n", "def build_text_file(data_json: dict, dest_path: str, \n", " tokenizer, machine_name_in_chat='Кирилл Гельван'):\n", " \"\"\"Create a text file for training in special format for ruDialoGPT-3.\n", "\n", " Parameters\n", " ----------\n", " data_json: dict\n", " Dict containing 'text' (message) and 'from' (user who sent the message)\n", " \n", " dest_path: str\n", " String containing path to write data there\n", "\n", " tokenizer: HuggingFace tokenizer \n", " Tokenizer that used to compute the length of the text after encoding.\n", " For more info ee https://huggingface.co/transformers/main_classes/tokenizer.html\n", " \"\"\"\n", " f = open(dest_path, 'w')\n", " new_data = ''\n", " for i in range(len(data_json) - 1):\n", " message, next_message = data_json[i], data_json[i+1]\n", " if message['text'] == '' or type(message['text']) != str:\n", " continue\n", " if next_message['text'] == '' or type(next_message['text']) != str:\n", " continue\n", "\n", " user = get_user_param(message, machine_name_in_chat=machine_name_in_chat)\n", " length = get_length_param(data_json[i+1]['text'], tokenizer)\n", " message_text = re.sub(r\"\\n\", \". \", message['text'])\n", " new_data += f\"|{user}|{length}|{message_text}{tokenizer.eos_token}\" + \"\\n\"\n", "\n", " f.write(new_data)\n", "\n", "\n", "def load_dataset(train_path, test_path, tokenizer):\n", " \"\"\"Creates train and test PyTorch datasets and collate_fn using HuggingFace.\n", "\n", " Parameters\n", " ----------\n", " train_path: str\n", " String containing path to train data\n", " \n", " test_path: str\n", " String containing path to test data\n", "\n", " tokenizer: HuggingFace tokenizer \n", " Tokenizer that used to compute the length of the text after encoding.\n", " For more info ee https://huggingface.co/transformers/main_classes/tokenizer.html\n", " \"\"\"\n", " train_dataset = TextDataset(\n", " tokenizer = tokenizer,\n", " file_path = train_path,\n", " block_size = 256)\n", " \n", " test_dataset = TextDataset(\n", " tokenizer = tokenizer,\n", " file_path = test_path,\n", " block_size = 256) \n", " \n", " data_collator = DataCollatorForLanguageModeling(\n", " tokenizer=tokenizer, mlm=False\n", " )\n", " return train_dataset, test_dataset, data_collator" ], "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "metadata": { "id": "wS5aTe48GF_N" }, "source": [ "1) Export your telegram chat\n", "\n", "![](https://raw.githubusercontent.com/Kirili4ik/ruDialoGpt3-finetune-colab/main/how-to-export-chat.jpg)\n", "\n", "2) Upload it to colab\n", "\n", "![](https://raw.githubusercontent.com/Kirili4ik/ruDialoGpt3-finetune-colab/main/how-to-upload-json.jpg)\n", "\n", "3) Next cell creates train and test set from it\n", "\n", "4) :tada:" ] }, { "cell_type": "code", "metadata": { "id": "19JKNqTS2Nu7", "cellView": "form" }, "source": [ "#@markdown Your telegram chat json path 'ChatExport.../YourChatName.json':\n", "path_to_telegram_chat_json = 'example: /content/drive/MyDrive/char27.json' #@param {type : \"string\"}\n", "#@markdown Name of the user to predict by GPT-3:\n", "machine_name_in_chat = 'example: Kirill Gelvan' #@param {type : \"string\"}\n", "\n", "\n", "with open(path_to_telegram_chat_json) as f: data = json.load(f)['messages']\n", "\n", "# test data is first 10% of chat, train - last 90%\n", "train, test = data[int(len(data)*0.1):], data[:int(len(data)*0.1)]\n", "\n", "build_text_file(train, 'train_dataset.txt', tokenizer)\n", "build_text_file(test, 'test_dataset.txt', tokenizer)\n", "\n", "print(\"Train dataset length: \" + str(len(train)) + \"samples\")\n", "print(\"Test dataset length: \" + str(len(test)) + \"samples\")" ], "execution_count": null, "outputs": [] }, { "cell_type": "code", "metadata": { "id": "qO1-aAHF6TxB" }, "source": [ "# let's look at our data\n", "! head -n 10 train_dataset.txt" ], "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "metadata": { "id": "J6dMhVaeIO8x" }, "source": [ "Here the first number is the spearker number - '1' for GPT and '0' for the person. \n", "\n", "The second number is the lengths of the expected answer: '1' for short, '2' for medium, '3' for long texts and '-' for all others. \n" ] }, { "cell_type": "code", "metadata": { "id": "-ty6A-qTzhya" }, "source": [ "# Create PyTorch Datasets\n", "train_dataset, test_dataset, data_collator = load_dataset('train_dataset.txt', 'test_dataset.txt', tokenizer)\n", "\n", "# Create PyTorch Dataloaders\n", "train_loader = DataLoader(train_dataset, shuffle=True, batch_size=2, collate_fn=data_collator)\n", "test_loader = DataLoader(test_dataset, batch_size=2, collate_fn=data_collator)" ], "execution_count": null, "outputs": [] }, { "cell_type": "code", "metadata": { "id": "NWhfc7ElAbkY" }, "source": [ "# this cell checks 1 forward pass\n", "try:\n", " for batch in train_loader:\n", " break\n", " {k: v.shape for k, v in batch.items()}\n", "\n", " outputs = model(**batch)\n", "except:\n", " print(\"Unexpected error:\", sys.exc_info()[0])\n", " raise" ], "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "metadata": { "id": "ESogNuUOEmj_" }, "source": [ "## Fine-tuning" ] }, { "cell_type": "code", "metadata": { "id": "mZBWIviea2-Y", "cellView": "form" }, "source": [ "#@title Fine-tuning params\n", "num_epochs = 3 #@param {type:\"integer\"}\n", "optimizer = AdamW(model.parameters(), lr=3e-5) #@param\n", "save_checkpoint_path = 'exmaple: drive/MyDrive/GPT2_checkpoint-more-data-2ep.pt' #@param {type:\"string\"}\n", "\n", "\n", "num_training_steps = num_epochs * len(train_dataset)\n", "lr_scheduler = get_scheduler(\n", " \"linear\",\n", " optimizer=optimizer,\n", " num_warmup_steps=100,\n", " num_training_steps=num_training_steps\n", ")\n", "\n", "accelerator = Accelerator()\n", "train_dl, test_dl, model, optimizer = accelerator.prepare(\n", " train_loader, test_loader, model, optimizer\n", ")\n", "# wandb.watch(model, log=\"all\")" ], "execution_count": null, "outputs": [] }, { "cell_type": "code", "metadata": { "id": "rEV3EcZOCOhw" }, "source": [ "progress_bar = tqdm(range(num_training_steps))\n", "\n", "for epoch in range(num_epochs):\n", " \n", " ### TRAIN EPOCH\n", " model.train()\n", " for batch in train_dl:\n", " optimizer.zero_grad()\n", " outputs = model(**batch)\n", " loss = outputs.loss\n", " accelerator.backward(loss)\n", " \n", " # wandb.log({'train_loss':loss.item()})\n", " optimizer.step()\n", " lr_scheduler.step()\n", " progress_bar.update(1)\n", "\n", " ### SAVE\n", " torch.save({\n", " 'model_state_dict': model.state_dict(),\n", " }, save_checkpoint_path)\n", " \n", " ### VALIDATE ONCE\n", " cum_loss = 0\n", " model.eval()\n", " with torch.inference_mode():\n", " for batch in test_dl:\n", " outputs = model(**batch)\n", " cum_loss += float(outputs.loss.item())\n", " \n", " print(cum_loss/len(test_loader))\n", " # wandb.log({'val_mean_loss':cum_loss/len(test_loader)})" ], "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "metadata": { "id": "psXZnJk0Eo3J" }, "source": [ "## Inference" ] }, { "cell_type": "code", "metadata": { "cellView": "form", "id": "3N13Nwd1axA4" }, "source": [ "#@title Installs and Utility functions\n", "\n", "%%capture\n", "# installing huggingface datasets and accelerate \n", "! pip install datasets transformers[sentencepiece]\n", "! pip install accelerate\n", "\n", "def get_length_param(text: str, tokenizer) -> str:\n", " \"\"\"Maps text to 1 of 4 buckets based on length after encoding.\n", "\n", " Parameters\n", " ----------\n", " text: str\n", " The text to be given 1 of 4 length parameters.\n", "\n", " tokenizer: HuggingFace tokenizer \n", " Tokenizer that used to compute the length of the text after encoding.\n", " For more info ee https://huggingface.co/transformers/main_classes/tokenizer.html\n", "\n", " Returns\n", " -------\n", " len_param: str\n", " One of four buckets: \n", " '1' for short, '2' for medium, '3' for long texts and '-' for all others. \n", " \"\"\"\n", " tokens_count = len(tokenizer.encode(text))\n", " if tokens_count <= 15:\n", " len_param = '1'\n", " elif tokens_count <= 50:\n", " len_param = '2'\n", " elif tokens_count <= 256:\n", " len_param = '3'\n", " else:\n", " len_param = '-'\n", " return len_param\n", "\n", "\n", "def get_user_param(text: dict, machine_name_in_chat: str) -> str:\n", " \"\"\"Maps text by 1/0 for it to be the person or the machine in the dialogue\n", "\n", " Parameters\n", " ----------\n", " text: Dict[..., 'from', ...]\n", " Dict containing field 'from' with the name of the user who sent the message\n", "\n", " machine_name_in_chat: str\n", " Str with the name of the machine - it will be predicted\n", " \"\"\"\n", " if text['from'] == machine_name_in_chat:\n", " return '1' # machine\n", " else:\n", " return '0' # human\n", "\n", "\n", "def build_text_file(data_json: dict, dest_path: str, \n", " tokenizer, machine_name_in_chat='Кирилл Гельван'):\n", " \"\"\"Create a text file for training in special format for ruDialoGPT-3.\n", "\n", " Parameters\n", " ----------\n", " data_json: dict\n", " Dict containing 'text' (message) and 'from' (user who sent the message)\n", " \n", " dest_path: str\n", " String containing path to write data there\n", "\n", " tokenizer: HuggingFace tokenizer \n", " Tokenizer that used to compute the length of the text after encoding.\n", " For more info ee https://huggingface.co/transformers/main_classes/tokenizer.html\n", " \"\"\"\n", " f = open(dest_path, 'w')\n", " new_data = ''\n", " for i in range(len(data_json) - 1):\n", " message, next_message = data_json[i], data_json[i+1]\n", " if message['text'] == '' or type(message['text']) != str:\n", " continue\n", " if next_message['text'] == '' or type(next_message['text']) != str:\n", " continue\n", "\n", " user = get_user_param(message, machine_name_in_chat=machine_name_in_chat)\n", " length = get_length_param(data_json[i+1]['text'], tokenizer)\n", " message_text = re.sub(r\"\\n\", \". \", message['text'])\n", " new_data += f\"|{user}|{length}|{message_text}{tokenizer.eos_token}\" + \"\\n\"\n", "\n", " f.write(new_data)\n", "\n", "\n", "def load_dataset(train_path, test_path, tokenizer):\n", " \"\"\"Creates train and test PyTorch datasets and collate_fn using HuggingFace.\n", "\n", " Parameters\n", " ----------\n", " train_path: str\n", " String containing path to train data\n", " \n", " test_path: str\n", " String containing path to test data\n", "\n", " tokenizer: HuggingFace tokenizer \n", " Tokenizer that used to compute the length of the text after encoding.\n", " For more info ee https://huggingface.co/transformers/main_classes/tokenizer.html\n", " \"\"\"\n", " train_dataset = TextDataset(\n", " tokenizer = tokenizer,\n", " file_path = train_path,\n", " block_size = 256)\n", " \n", " test_dataset = TextDataset(\n", " tokenizer = tokenizer,\n", " file_path = test_path,\n", " block_size = 256) \n", " \n", " data_collator = DataCollatorForLanguageModeling(\n", " tokenizer=tokenizer, mlm=False\n", " )\n", " return train_dataset, test_dataset, data_collator" ], "execution_count": null, "outputs": [] }, { "cell_type": "code", "metadata": { "id": "vvsSRglEA0kt" }, "source": [ "import torch\n", "from transformers import AutoModelForCausalLM, AutoTokenizer\n", "\n", "# Download checkpoint:\n", "checkpoint = \"Kirili4ik/ruDialoGpt3-medium-finetuned-telegram\" \n", "tokenizer = AutoTokenizer.from_pretrained(checkpoint)\n", "model = AutoModelForCausalLM.from_pretrained(checkpoint)\n", "\n", "# [optional] Insert your checkpoint if needed:\n", "'''from google.colab import drive\n", "drive.mount('/content/drive')\n", "checkpoint = torch.load('drive/MyDrive/GPT2_checkpoint.pt', map_location='cpu')\n", "model.load_state_dict(checkpoint['model_state_dict'])'''\n", "\n", "model = model.to('cpu')\n", "model.eval()\n", "print()" ], "execution_count": null, "outputs": [] }, { "cell_type": "code", "metadata": { "id": "MGdCxVnOhK_K" }, "source": [ "### INFERENCE\n", "\n", "chat_history_ids = torch.zeros((1, 0), dtype=torch.int)\n", "\n", "while True:\n", " \n", " next_who = input(\"Who's phrase?\\t\") #input(\"H / G?\") # Human or GPT\n", "\n", " # In case Human\n", " if next_who == \"H\":\n", " input_user = input(\"===> Human: \")\n", " \n", " # encode the new user input, add parameters and return a tensor in Pytorch\n", " new_user_input_ids = tokenizer.encode(f\"|0|{get_length_param(input_user, tokenizer)}|\" \\\n", " + input_user + tokenizer.eos_token, return_tensors=\"pt\")\n", " # append the new user input tokens to the chat history\n", " chat_history_ids = torch.cat([chat_history_ids, new_user_input_ids], dim=-1)\n", "\n", " if next_who == \"G\":\n", "\n", " next_len = input(\"Phrase len? 1/2/3/-\\t\") #input(\"Exp. len?(-/1/2/3): \")\n", " # encode the new user input, add parameters and return a tensor in Pytorch\n", " new_user_input_ids = tokenizer.encode(f\"|1|{next_len}|\", return_tensors=\"pt\")\n", " # append the new user input tokens to the chat history\n", " chat_history_ids = torch.cat([chat_history_ids, new_user_input_ids], dim=-1)\n", " \n", " # print(tokenizer.decode(chat_history_ids[-1])) # uncomment to see full gpt input\n", " \n", " # save previous len\n", " input_len = chat_history_ids.shape[-1]\n", " # generated a response; PS you can read about the parameters at hf.co/blog/how-to-generate\n", " chat_history_ids = model.generate(\n", " chat_history_ids,\n", " num_return_sequences=1, # use for more variants, but have to print [i]\n", " max_length=512,\n", " no_repeat_ngram_size=3,\n", " do_sample=True,\n", " top_k=50,\n", " top_p=0.9,\n", " temperature = 0.6, # 0 for greedy\n", " mask_token_id=tokenizer.mask_token_id,\n", " eos_token_id=tokenizer.eos_token_id,\n", " unk_token_id=tokenizer.unk_token_id,\n", " pad_token_id=tokenizer.pad_token_id,\n", " device='cpu'\n", " )\n", " \n", " # pretty print last ouput tokens from bot\n", " print(f\"===> GPT-3: {tokenizer.decode(chat_history_ids[:, input_len:][0], skip_special_tokens=True)}\")" ], "execution_count": null, "outputs": [] }, { "cell_type": "code", "metadata": { "id": "mjEQiv5TMjZW" }, "source": [ "" ], "execution_count": null, "outputs": [] } ] }