Fine_tune_RuDialoGPT3_on_telegram_chat.ipynb

{
  "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": [
        "<a href=\"https://colab.research.google.com/github/Kirili4ik/ruDialoGpt3-finetune-colab/blob/main/Fine_tune_RuDialoGPT3_on_telegram_chat.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "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": []
    }
  ]
}