![\"Open](\"https://colab.research.google.com/assets/colab-badge.svg\"){kind=icon}
"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "_6q2San5Xh5N",
"colab_type": "text"
},
"source": [
"# T5 on TPU 💥🚀"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "knpacarPX2AN",
"colab_type": "text"
},
"source": [
"In this notebook we will see how to train T5 model on TPU with Huggingface's awesome new [trainer](https://github.com/huggingface/transformers/blob/master/src/transformers/trainer.py). We will train T5 base model on SQUAD dataset for QA task. We will use the recently released amazing [nlp](https://github.com/huggingface/nlp) package to load and process the dataset in just few lines.\n",
"\n",
"First make sure you are connected to the high RAM instance. This will not work on 12 GB colab instance."
]
},
{
"cell_type": "code",
"metadata": {
"id": "QLGiFCDqvuil",
"colab_type": "code",
"colab": {}
},
"source": [
"# Crash on purpose to get more ram :\n",
"import torch\n",
"torch.tensor([10.]*10000000000)"
],
"execution_count": 0,
"outputs": []
},
{
"cell_type": "markdown",
"metadata": {
"id": "40CJSrN9ZiIP",
"colab_type": "text"
},
"source": [
"Let's install [PyTorch/XLA](https://github.com/pytorch/xla) which enables PyTorch on TPU. Make sure you install the nightly version, as the trainer breaks on other versions."
]
},
{
"cell_type": "code",
"metadata": {
"id": "coOmS2s_xDBy",
"colab_type": "code",
"outputId": "6ed0947e-4061-4ba7-9eca-72e7ed935bd6",
"colab": {
"base_uri": "https://localhost:8080/",
"height": 907
}
},
"source": [
"VERSION = \"nightly\" #@param [\"1.5\" , \"20200325\", \"nightly\"]\n",
"!curl https://raw.githubusercontent.com/pytorch/xla/master/contrib/scripts/env-setup.py -o pytorch-xla-env-setup.py\n",
"!python pytorch-xla-env-setup.py --version $VERSION"
],
"execution_count": 5,
"outputs": [
{
"output_type": "stream",
"text": [
" % Total % Received % Xferd Average Speed Time Time Time Current\n",
" Dload Upload Total Spent Left Speed\n",
"\r 0 0 0 0 0 0 0 0 --:--:-- --:--:-- --:--:-- 0\r100 4264 100 4264 0 0 60914 0 --:--:-- --:--:-- --:--:-- 60914\n",
"Updating TPU and VM. This may take around 2 minutes.\n",
"Updating TPU runtime to pytorch-nightly ...\n",
"Uninstalling torch-1.5.0+cu101:\n",
"Done updating TPU runtime: "
]
},
{
"cell_type": "code",
"metadata": {
"id": "PJX4vkjj6wYz",
"colab_type": "code",
"outputId": "83a8a420-48cd-4d49-bc60-2693268481c6",
"colab": {
"base_uri": "https://localhost:8080/",
"height": 121
}
},
"source": [
"from google.colab import drive\n",
"drive.mount('/content/drive')"
],
"execution_count": 0,
"outputs": [
{
"output_type": "stream",
"text": [
"Go to this URL in a browser: https://accounts.google.com/o/oauth2/auth?client_id=947318989803-6bn6qk8qdgf4n4g3pfee6491hc0brc4i.apps.googleusercontent.com&redirect_uri=urn%3aietf%3awg%3aoauth%3a2.0%3aoob&response_type=code&scope=email%20https%3a%2f%2fwww.googleapis.com%2fauth%2fdocs.test%20https%3a%2f%2fwww.googleapis.com%2fauth%2fdrive%20https%3a%2f%2fwww.googleapis.com%2fauth%2fdrive.photos.readonly%20https%3a%2f%2fwww.googleapis.com%2fauth%2fpeopleapi.readonly\n",
"\n",
"Enter your authorization code:\n",
"··········\n",
"Mounted at /content/drive\n"
],
"name": "stdout"
}
]
},
{
"cell_type": "code",
"metadata": {
"id": "1V5cInhu42Wk",
"colab_type": "code",
"outputId": "5501a5f1-fc49-4df7-f7a0-31cc37647337",
"colab": {
"base_uri": "https://localhost:8080/",
"height": 302
}
},
"source": [
"!nvidia-smi"
],
"execution_count": 0,
"outputs": [
{
"output_type": "stream",
"text": [
"Sat May 9 15:05:25 2020 \n",
"+-----------------------------------------------------------------------------+\n",
"| NVIDIA-SMI 440.82 Driver Version: 418.67 CUDA Version: 10.1 |\n",
"|-------------------------------+----------------------+----------------------+\n",
"| GPU Name Persistence-M| Bus-Id Disp.A | Volatile Uncorr. ECC |\n",
"| Fan Temp Perf Pwr:Usage/Cap| Memory-Usage | GPU-Util Compute M. |\n",
"|===============================+======================+======================|\n",
"| 0 Tesla P100-PCIE... Off | 00000000:00:04.0 Off | 0 |\n",
"| N/A 34C P0 25W / 250W | 0MiB / 16280MiB | 0% Default |\n",
"+-------------------------------+----------------------+----------------------+\n",
" \n",
"+-----------------------------------------------------------------------------+\n",
"| Processes: GPU Memory |\n",
"| GPU PID Type Process name Usage |\n",
"|=============================================================================|\n",
"| No running processes found |\n",
"+-----------------------------------------------------------------------------+\n"
],
"name": "stdout"
}
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "epWcPHhJ7v7j",
"colab_type": "text"
},
"source": [
"Instal apex if you want to do 16 bit training. You'll probably need to restart the notebook after installing apex"
]
},
{
"cell_type": "code",
"metadata": {
"id": "k1Xy7ZG-7gHt",
"colab_type": "code",
"colab": {}
},
"source": [
"# !export CUDA_HOME=/usr/local/cuda-10.1\n",
"# !git clone https://github.com/NVIDIA/apex\n",
"# !pip install -v --no-cache-dir --global-option=\"--cpp_ext\" --global-option=\"--cuda_ext\" ./apex"
],
"execution_count": 0,
"outputs": []
},
{
"cell_type": "code",
"metadata": {
"id": "SDVQ04fGRb1v",
"colab_type": "code",
"outputId": "11689986-ca27-4ab0-f14d-5ee4f0eba40d",
"colab": {
"base_uri": "https://localhost:8080/",
"height": 1000
}
},
"source": [
"!pip install transformers\n",
"!pip install pytorch_lightning"
],
"execution_count": 0,
"outputs": [
{
"output_type": "stream",
"text": [
"Collecting transformers\n",
"\u001b[?25l Downloading https://files.pythonhosted.org/packages/cd/38/c9527aa055241c66c4d785381eaf6f80a28c224cae97daa1f8b183b5fabb/transformers-2.9.0-py3-none-any.whl (635kB)\n",
"\r\u001b[K |▌ | 10kB 20.5MB/s eta 0:00:01\r\u001b[K |█ | 20kB 1.7MB/s eta 0:00:01\r\u001b[K |█▌ | 30kB 2.3MB/s eta 0:00:01\r\u001b[K |██ | 40kB 2.6MB/s eta 0:00:01\r\u001b[K |██▋ | 51kB 2.0MB/s eta 0:00:01\r\u001b[K |███ | 61kB 2.3MB/s eta 0:00:01\r\u001b[K |███▋ | 71kB 2.5MB/s eta 0:00:01\r\u001b[K |████▏ | 81kB 2.7MB/s eta 0:00:01\r\u001b[K |████▋ | 92kB 3.0MB/s eta 0:00:01\r\u001b[K |█████▏ | 102kB 2.8MB/s eta 0:00:01\r\u001b[K |█████▊ | 112kB 2.8MB/s eta 0:00:01\r\u001b[K |██████▏ | 122kB 2.8MB/s eta 0:00:01\r\u001b[K |██████▊ | 133kB 2.8MB/s eta 0:00:01\r\u001b[K |███████▏ | 143kB 2.8MB/s eta 0:00:01\r\u001b[K |███████▊ | 153kB 2.8MB/s eta 0:00:01\r\u001b[K |████████▎ | 163kB 2.8MB/s eta 0:00:01\r\u001b[K |████████▊ | 174kB 2.8MB/s eta 0:00:01\r\u001b[K |█████████▎ | 184kB 2.8MB/s eta 0:00:01\r\u001b[K |█████████▉ | 194kB 2.8MB/s eta 0:00:01\r\u001b[K |██████████▎ | 204kB 2.8MB/s eta 0:00:01\r\u001b[K |██████████▉ | 215kB 2.8MB/s eta 0:00:01\r\u001b[K |███████████▍ | 225kB 2.8MB/s eta 0:00:01\r\u001b[K |███████████▉ | 235kB 2.8MB/s eta 0:00:01\r\u001b[K |████████████▍ | 245kB 2.8MB/s eta 0:00:01\r\u001b[K |████████████▉ | 256kB 2.8MB/s eta 0:00:01\r\u001b[K |█████████████▍ | 266kB 2.8MB/s eta 0:00:01\r\u001b[K |██████████████ | 276kB 2.8MB/s eta 0:00:01\r\u001b[K |██████████████▍ | 286kB 2.8MB/s eta 0:00:01\r\u001b[K |███████████████ | 296kB 2.8MB/s eta 0:00:01\r\u001b[K |███████████████▌ | 307kB 2.8MB/s eta 0:00:01\r\u001b[K |████████████████ | 317kB 2.8MB/s eta 0:00:01\r\u001b[K |████████████████▌ | 327kB 2.8MB/s eta 0:00:01\r\u001b[K |█████████████████ | 337kB 2.8MB/s eta 0:00:01\r\u001b[K |█████████████████▌ | 348kB 2.8MB/s eta 0:00:01\r\u001b[K |██████████████████ | 358kB 2.8MB/s eta 0:00:01\r\u001b[K |██████████████████▋ | 368kB 2.8MB/s eta 0:00:01\r\u001b[K |███████████████████ | 378kB 2.8MB/s eta 0:00:01\r\u001b[K |███████████████████▋ | 389kB 2.8MB/s eta 0:00:01\r\u001b[K |████████████████████ | 399kB 2.8MB/s eta 0:00:01\r\u001b[K |████████████████████▋ | 409kB 2.8MB/s eta 0:00:01\r\u001b[K |█████████████████████▏ | 419kB 2.8MB/s eta 0:00:01\r\u001b[K |█████████████████████▋ | 430kB 2.8MB/s eta 0:00:01\r\u001b[K |██████████████████████▏ | 440kB 2.8MB/s eta 0:00:01\r\u001b[K |██████████████████████▊ | 450kB 2.8MB/s eta 0:00:01\r\u001b[K |███████████████████████▏ | 460kB 2.8MB/s eta 0:00:01\r\u001b[K |███████████████████████▊ | 471kB 2.8MB/s eta 0:00:01\r\u001b[K |████████████████████████▎ | 481kB 2.8MB/s eta 0:00:01\r\u001b[K |████████████████████████▊ | 491kB 2.8MB/s eta 0:00:01\r\u001b[K |█████████████████████████▎ | 501kB 2.8MB/s eta 0:00:01\r\u001b[K |█████████████████████████▊ | 512kB 2.8MB/s eta 0:00:01\r\u001b[K |██████████████████████████▎ | 522kB 2.8MB/s eta 0:00:01\r\u001b[K |██████████████████████████▉ | 532kB 2.8MB/s eta 0:00:01\r\u001b[K |███████████████████████████▎ | 542kB 2.8MB/s eta 0:00:01\r\u001b[K |███████████████████████████▉ | 552kB 2.8MB/s eta 0:00:01\r\u001b[K |████████████████████████████▍ | 563kB 2.8MB/s eta 0:00:01\r\u001b[K |████████████████████████████▉ | 573kB 2.8MB/s eta 0:00:01\r\u001b[K |█████████████████████████████▍ | 583kB 2.8MB/s eta 0:00:01\r\u001b[K |██████████████████████████████ | 593kB 2.8MB/s eta 0:00:01\r\u001b[K |██████████████████████████████▍ | 604kB 2.8MB/s eta 0:00:01\r\u001b[K |███████████████████████████████ | 614kB 2.8MB/s eta 0:00:01\r\u001b[K |███████████████████████████████▍| 624kB 2.8MB/s eta 0:00:01\r\u001b[K |████████████████████████████████| 634kB 2.8MB/s eta 0:00:01\r\u001b[K |████████████████████████████████| 645kB 2.8MB/s \n",
"\u001b[?25hRequirement already satisfied: requests in /usr/local/lib/python3.6/dist-packages (from transformers) (2.23.0)\n",
"Requirement already satisfied: dataclasses; python_version < \"3.7\" in /usr/local/lib/python3.6/dist-packages (from transformers) (0.7)\n",
"Requirement already satisfied: tqdm>=4.27 in /usr/local/lib/python3.6/dist-packages (from transformers) (4.41.1)\n",
"Collecting tokenizers==0.7.0\n",
"\u001b[?25l Downloading https://files.pythonhosted.org/packages/14/e5/a26eb4716523808bb0a799fcfdceb6ebf77a18169d9591b2f46a9adb87d9/tokenizers-0.7.0-cp36-cp36m-manylinux1_x86_64.whl (3.8MB)\n",
"\u001b[K |████████████████████████████████| 3.8MB 12.8MB/s \n",
"\u001b[?25hRequirement already satisfied: regex!=2019.12.17 in /usr/local/lib/python3.6/dist-packages (from transformers) (2019.12.20)\n",
"Collecting sacremoses\n",
"\u001b[?25l Downloading https://files.pythonhosted.org/packages/7d/34/09d19aff26edcc8eb2a01bed8e98f13a1537005d31e95233fd48216eed10/sacremoses-0.0.43.tar.gz (883kB)\n",
"\u001b[K |████████████████████████████████| 890kB 28.2MB/s \n",
"\u001b[?25hCollecting sentencepiece\n",
"\u001b[?25l Downloading https://files.pythonhosted.org/packages/98/2c/8df20f3ac6c22ac224fff307ebc102818206c53fc454ecd37d8ac2060df5/sentencepiece-0.1.86-cp36-cp36m-manylinux1_x86_64.whl (1.0MB)\n",
"\u001b[K |████████████████████████████████| 1.0MB 42.6MB/s \n",
"\u001b[?25hRequirement already satisfied: numpy in /usr/local/lib/python3.6/dist-packages (from transformers) (1.18.4)\n",
"Requirement already satisfied: filelock in /usr/local/lib/python3.6/dist-packages (from transformers) (3.0.12)\n",
"Requirement already satisfied: chardet<4,>=3.0.2 in /usr/local/lib/python3.6/dist-packages (from requests->transformers) (3.0.4)\n",
"Requirement already satisfied: urllib3!=1.25.0,!=1.25.1,<1.26,>=1.21.1 in /usr/local/lib/python3.6/dist-packages (from requests->transformers) (1.24.3)\n",
"Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.6/dist-packages (from requests->transformers) (2020.4.5.1)\n",
"Requirement already satisfied: idna<3,>=2.5 in /usr/local/lib/python3.6/dist-packages (from requests->transformers) (2.9)\n",
"Requirement already satisfied: six in /usr/local/lib/python3.6/dist-packages (from sacremoses->transformers) (1.12.0)\n",
"Requirement already satisfied: click in /usr/local/lib/python3.6/dist-packages (from sacremoses->transformers) (7.1.2)\n",
"Requirement already satisfied: joblib in /usr/local/lib/python3.6/dist-packages (from sacremoses->transformers) (0.14.1)\n",
"Building wheels for collected packages: sacremoses\n",
" Building wheel for sacremoses (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
" Created wheel for sacremoses: filename=sacremoses-0.0.43-cp36-none-any.whl size=893260 sha256=1d6422ddbf7526c5762d09193b36548e7b07bf9cc526057f833254f31a68c87c\n",
" Stored in directory: /root/.cache/pip/wheels/29/3c/fd/7ce5c3f0666dab31a50123635e6fb5e19ceb42ce38d4e58f45\n",
"Successfully built sacremoses\n",
"Installing collected packages: tokenizers, sacremoses, sentencepiece, transformers\n",
"Successfully installed sacremoses-0.0.43 sentencepiece-0.1.86 tokenizers-0.7.0 transformers-2.9.0\n",
"Collecting pytorch_lightning\n",
"\u001b[?25l Downloading https://files.pythonhosted.org/packages/75/ac/ac03f1f3fa950d96ca52f07d33fdbf5add05f164c1ac4eae179231dfa93d/pytorch_lightning-0.7.5-py3-none-any.whl (233kB)\n",
"\u001b[K |████████████████████████████████| 235kB 2.8MB/s \n",
"\u001b[?25hRequirement already satisfied: numpy>=1.16.4 in /usr/local/lib/python3.6/dist-packages (from pytorch_lightning) (1.18.4)\n",
"Requirement already satisfied: tqdm>=4.41.0 in /usr/local/lib/python3.6/dist-packages (from pytorch_lightning) (4.41.1)\n",
"Collecting future>=0.17.1\n",
"\u001b[?25l Downloading https://files.pythonhosted.org/packages/45/0b/38b06fd9b92dc2b68d58b75f900e97884c45bedd2ff83203d933cf5851c9/future-0.18.2.tar.gz (829kB)\n",
"\u001b[K |████████████████████████████████| 829kB 8.4MB/s \n",
"\u001b[?25hRequirement already satisfied: tensorboard>=1.14 in /usr/local/lib/python3.6/dist-packages (from pytorch_lightning) (2.2.1)\n",
"Requirement already satisfied: torch>=1.1 in /usr/local/lib/python3.6/dist-packages (from pytorch_lightning) (1.5.0+cu101)\n",
"Requirement already satisfied: six>=1.10.0 in /usr/local/lib/python3.6/dist-packages (from tensorboard>=1.14->pytorch_lightning) (1.12.0)\n",
"Requirement already satisfied: tensorboard-plugin-wit>=1.6.0 in /usr/local/lib/python3.6/dist-packages (from tensorboard>=1.14->pytorch_lightning) (1.6.0.post3)\n",
"Requirement already satisfied: google-auth<2,>=1.6.3 in /usr/local/lib/python3.6/dist-packages (from tensorboard>=1.14->pytorch_lightning) (1.7.2)\n",
"Requirement already satisfied: markdown>=2.6.8 in /usr/local/lib/python3.6/dist-packages (from tensorboard>=1.14->pytorch_lightning) (3.2.1)\n",
"Requirement already satisfied: wheel>=0.26; python_version >= \"3\" in /usr/local/lib/python3.6/dist-packages (from tensorboard>=1.14->pytorch_lightning) (0.34.2)\n",
"Requirement already satisfied: google-auth-oauthlib<0.5,>=0.4.1 in /usr/local/lib/python3.6/dist-packages (from tensorboard>=1.14->pytorch_lightning) (0.4.1)\n",
"Requirement already satisfied: grpcio>=1.24.3 in /usr/local/lib/python3.6/dist-packages (from tensorboard>=1.14->pytorch_lightning) (1.28.1)\n",
"Requirement already satisfied: protobuf>=3.6.0 in /usr/local/lib/python3.6/dist-packages (from tensorboard>=1.14->pytorch_lightning) (3.10.0)\n",
"Requirement already satisfied: absl-py>=0.4 in /usr/local/lib/python3.6/dist-packages (from tensorboard>=1.14->pytorch_lightning) (0.9.0)\n",
"Requirement already satisfied: requests<3,>=2.21.0 in /usr/local/lib/python3.6/dist-packages (from tensorboard>=1.14->pytorch_lightning) (2.23.0)\n",
"Requirement already satisfied: setuptools>=41.0.0 in /usr/local/lib/python3.6/dist-packages (from tensorboard>=1.14->pytorch_lightning) (46.1.3)\n",
"Requirement already satisfied: werkzeug>=0.11.15 in /usr/local/lib/python3.6/dist-packages (from tensorboard>=1.14->pytorch_lightning) (1.0.1)\n",
"Requirement already satisfied: pyasn1-modules>=0.2.1 in /usr/local/lib/python3.6/dist-packages (from google-auth<2,>=1.6.3->tensorboard>=1.14->pytorch_lightning) (0.2.8)\n",
"Requirement already satisfied: cachetools<3.2,>=2.0.0 in /usr/local/lib/python3.6/dist-packages (from google-auth<2,>=1.6.3->tensorboard>=1.14->pytorch_lightning) (3.1.1)\n",
"Requirement already satisfied: rsa<4.1,>=3.1.4 in /usr/local/lib/python3.6/dist-packages (from google-auth<2,>=1.6.3->tensorboard>=1.14->pytorch_lightning) (4.0)\n",
"Requirement already satisfied: requests-oauthlib>=0.7.0 in /usr/local/lib/python3.6/dist-packages (from google-auth-oauthlib<0.5,>=0.4.1->tensorboard>=1.14->pytorch_lightning) (1.3.0)\n",
"Requirement already satisfied: idna<3,>=2.5 in /usr/local/lib/python3.6/dist-packages (from requests<3,>=2.21.0->tensorboard>=1.14->pytorch_lightning) (2.9)\n",
"Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.6/dist-packages (from requests<3,>=2.21.0->tensorboard>=1.14->pytorch_lightning) (2020.4.5.1)\n",
"Requirement already satisfied: urllib3!=1.25.0,!=1.25.1,<1.26,>=1.21.1 in /usr/local/lib/python3.6/dist-packages (from requests<3,>=2.21.0->tensorboard>=1.14->pytorch_lightning) (1.24.3)\n",
"Requirement already satisfied: chardet<4,>=3.0.2 in /usr/local/lib/python3.6/dist-packages (from requests<3,>=2.21.0->tensorboard>=1.14->pytorch_lightning) (3.0.4)\n",
"Requirement already satisfied: pyasn1<0.5.0,>=0.4.6 in /usr/local/lib/python3.6/dist-packages (from pyasn1-modules>=0.2.1->google-auth<2,>=1.6.3->tensorboard>=1.14->pytorch_lightning) (0.4.8)\n",
"Requirement already satisfied: oauthlib>=3.0.0 in /usr/local/lib/python3.6/dist-packages (from requests-oauthlib>=0.7.0->google-auth-oauthlib<0.5,>=0.4.1->tensorboard>=1.14->pytorch_lightning) (3.1.0)\n",
"Building wheels for collected packages: future\n",
" Building wheel for future (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
" Created wheel for future: filename=future-0.18.2-cp36-none-any.whl size=491057 sha256=2748e4b7f9acd3e1e87b8118cdcb4cb5a4bf5ed682d99f3866e26265ab336042\n",
" Stored in directory: /root/.cache/pip/wheels/8b/99/a0/81daf51dcd359a9377b110a8a886b3895921802d2fc1b2397e\n",
"Successfully built future\n",
"Installing collected packages: future, pytorch-lightning\n",
" Found existing installation: future 0.16.0\n",
" Uninstalling future-0.16.0:\n",
" Successfully uninstalled future-0.16.0\n",
"Successfully installed future-0.18.2 pytorch-lightning-0.7.5\n"
],
"name": "stdout"
}
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "HVxGfmEMCKs_",
"colab_type": "text"
},
"source": [
"## T5 fine-tuning\n",
"\n",
"This notebook is to showcase how to fine-tune [T5 model](https://arxiv.org/abs/1910.10683) with Huggigface's [Transformers](https://github.com/huggingface/transformers/) to solve different NLP tasks using text-2-text approach proposed in the T5 paper. For demo I chose 3 non text-2-text problems just to reiterate the fact from the paper that how widely applicable this text-2-text framework is and how it can be used for different tasks without changing the model at all.\n",
"\n",
"This is a rough draft so if you find any issues with this notebook or have any questions reach out to me via [Twitter](https://twitter.com/psuraj28).\n",
"\n"
]
},
{
"cell_type": "code",
"metadata": {
"id": "HS8mNXq6bdxq",
"colab_type": "code",
"outputId": "b0a32f10-f2ef-4d49-b433-266e8206040b",
"colab": {
"base_uri": "https://localhost:8080/",
"height": 84
}
},
"source": [
"import argparse\n",
"import glob\n",
"import os\n",
"import json\n",
"import time\n",
"import logging\n",
"import random\n",
"import re\n",
"from itertools import chain\n",
"from string import punctuation\n",
"\n",
"import nltk\n",
"nltk.download('punkt')\n",
"from nltk.tokenize import sent_tokenize\n",
"\n",
"import pandas as pd\n",
"import numpy as np\n",
"import torch\n",
"from torch.utils.data import Dataset, DataLoader\n",
"import pytorch_lightning as pl\n",
"\n",
"\n",
"from transformers import (\n",
" AdamW,\n",
" T5ForConditionalGeneration,\n",
" T5Tokenizer,\n",
" get_linear_schedule_with_warmup\n",
")"
],
"execution_count": 0,
"outputs": [
{
"output_type": "stream",
"text": [
"[nltk_data] Downloading package punkt to /root/nltk_data...\n",
"[nltk_data] Unzipping tokenizers/punkt.zip.\n"
],
"name": "stdout"
},
{
"output_type": "stream",
"text": [
"INFO:transformers.file_utils:PyTorch version 1.5.0+cu101 available.\n",
"INFO:transformers.file_utils:TensorFlow version 2.2.0-rc4 available.\n"
],
"name": "stderr"
}
]
},
{
"cell_type": "code",
"metadata": {
"id": "IswYuhWaz7QJ",
"colab_type": "code",
"colab": {}
},
"source": [
"def set_seed(seed):\n",
" random.seed(seed)\n",
" np.random.seed(seed)\n",
" torch.manual_seed(seed)\n",
" if torch.cuda.is_available():\n",
" torch.cuda.manual_seed_all(seed)\n",
"\n",
"set_seed(42)"
],
"execution_count": 0,
"outputs": []
},
{
"cell_type": "markdown",
"metadata": {
"id": "RKNr7fgzcKpZ",
"colab_type": "text"
},
"source": [
"## Model\n",
"\n",
"We'll be using the awesome [pytorch-lightning](https://github.com/PytorchLightning/pytorch-lightning) library for training. Most of the below code is adapted from here https://github.com/huggingface/transformers/blob/master/examples/lightning_base.py\n",
"\n",
"The trainer is generic and can be used for any text-2-text task. You'll just need to change the dataset. Rest of the code will stay unchanged for all the tasks.\n",
"\n",
"This is the most intresting and powrfull thing about the text-2-text format. You can fine-tune the model on variety of NLP tasks by just formulating the problem in text-2-text setting. No need to change hyperparameters, learning rate, optimizer or loss function. Just plug in your dataset and you are ready to go!"
]
},
{
"cell_type": "code",
"metadata": {
"id": "B7uVNBtXST5X",
"colab_type": "code",
"colab": {}
},
"source": [
"class T5FineTuner(pl.LightningModule):\n",
" def __init__(self, hparams):\n",
" super(T5FineTuner, self).__init__()\n",
" self.hparams = hparams\n",
" \n",
" self.model = T5ForConditionalGeneration.from_pretrained(hparams.model_name_or_path)\n",
" self.tokenizer = T5Tokenizer.from_pretrained(hparams.tokenizer_name_or_path)\n",
" \n",
" def is_logger(self):\n",
" return self.trainer.proc_rank <= 0\n",
" \n",
" def forward(\n",
" self, input_ids, attention_mask=None, decoder_input_ids=None, decoder_attention_mask=None, lm_labels=None\n",
" ):\n",
" return self.model(\n",
" input_ids,\n",
" attention_mask=attention_mask,\n",
" decoder_input_ids=decoder_input_ids,\n",
" decoder_attention_mask=decoder_attention_mask,\n",
" lm_labels=lm_labels,\n",
" )\n",
"\n",
" def _step(self, batch):\n",
" lm_labels = batch[\"target_ids\"]\n",
" lm_labels[lm_labels[:, :] == self.tokenizer.pad_token_id] = -100\n",
"\n",
" outputs = self(\n",
" input_ids=batch[\"source_ids\"],\n",
" attention_mask=batch[\"source_mask\"],\n",
" lm_labels=lm_labels,\n",
" decoder_attention_mask=batch['target_mask']\n",
" )\n",
"\n",
" loss = outputs[0]\n",
"\n",
" return loss\n",
"\n",
" def training_step(self, batch, batch_idx):\n",
" loss = self._step(batch)\n",
"\n",
" tensorboard_logs = {\"train_loss\": loss}\n",
" return {\"loss\": loss, \"log\": tensorboard_logs}\n",
" \n",
" def training_epoch_end(self, outputs):\n",
" avg_train_loss = torch.stack([x[\"loss\"] for x in outputs]).mean()\n",
" tensorboard_logs = {\"avg_train_loss\": avg_train_loss}\n",
" return {\"avg_train_loss\": avg_train_loss, \"log\": tensorboard_logs, 'progress_bar': tensorboard_logs}\n",
"\n",
" def validation_step(self, batch, batch_idx):\n",
" loss = self._step(batch)\n",
" return {\"val_loss\": loss}\n",
" \n",
" def validation_epoch_end(self, outputs):\n",
" avg_loss = torch.stack([x[\"val_loss\"] for x in outputs]).mean()\n",
" tensorboard_logs = {\"val_loss\": avg_loss}\n",
" return {\"avg_val_loss\": avg_loss, \"log\": tensorboard_logs, 'progress_bar': tensorboard_logs}\n",
"\n",
" def configure_optimizers(self):\n",
" \"Prepare optimizer and schedule (linear warmup and decay)\"\n",
"\n",
" model = self.model\n",
" no_decay = [\"bias\", \"LayerNorm.weight\"]\n",
" optimizer_grouped_parameters = [\n",
" {\n",
" \"params\": [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)],\n",
" \"weight_decay\": self.hparams.weight_decay,\n",
" },\n",
" {\n",
" \"params\": [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)],\n",
" \"weight_decay\": 0.0,\n",
" },\n",
" ]\n",
" optimizer = AdamW(optimizer_grouped_parameters, lr=self.hparams.learning_rate, eps=self.hparams.adam_epsilon)\n",
" self.opt = optimizer\n",
" return [optimizer]\n",
" \n",
" def optimizer_step(self, epoch, batch_idx, optimizer, optimizer_idx, second_order_closure=None):\n",
" if self.trainer.use_tpu:\n",
" xm.optimizer_step(optimizer)\n",
" else:\n",
" optimizer.step()\n",
" optimizer.zero_grad()\n",
" self.lr_scheduler.step()\n",
" \n",
" def get_tqdm_dict(self):\n",
" tqdm_dict = {\"loss\": \"{:.3f}\".format(self.trainer.avg_loss), \"lr\": self.lr_scheduler.get_last_lr()[-1]}\n",
"\n",
" return tqdm_dict\n",
"\n",
" def train_dataloader(self):\n",
" train_dataset = get_dataset(tokenizer=self.tokenizer, type_path=\"train\", args=self.hparams)\n",
" dataloader = DataLoader(train_dataset, batch_size=self.hparams.train_batch_size, drop_last=True, shuffle=True, num_workers=4)\n",
" t_total = (\n",
" (len(dataloader.dataset) // (self.hparams.train_batch_size * max(1, self.hparams.n_gpu)))\n",
" // self.hparams.gradient_accumulation_steps\n",
" * float(self.hparams.num_train_epochs)\n",
" )\n",
" scheduler = get_linear_schedule_with_warmup(\n",
" self.opt, num_warmup_steps=self.hparams.warmup_steps, num_training_steps=t_total\n",
" )\n",
" self.lr_scheduler = scheduler\n",
" return dataloader\n",
"\n",
" def val_dataloader(self):\n",
" val_dataset = get_dataset(tokenizer=self.tokenizer, type_path=\"val\", args=self.hparams)\n",
" return DataLoader(val_dataset, batch_size=self.hparams.eval_batch_size, num_workers=4)"
],
"execution_count": 0,
"outputs": []
},
{
"cell_type": "code",
"metadata": {
"id": "oh1R5C-GwMqx",
"colab_type": "code",
"colab": {}
},
"source": [
"logger = logging.getLogger(__name__)\n",
"\n",
"class LoggingCallback(pl.Callback):\n",
" def on_validation_end(self, trainer, pl_module):\n",
" logger.info(\"***** Validation results *****\")\n",
" if pl_module.is_logger():\n",
" metrics = trainer.callback_metrics\n",
" # Log results\n",
" for key in sorted(metrics):\n",
" if key not in [\"log\", \"progress_bar\"]:\n",
" logger.info(\"{} = {}\\n\".format(key, str(metrics[key])))\n",
"\n",
" def on_test_end(self, trainer, pl_module):\n",
" logger.info(\"***** Test results *****\")\n",
"\n",
" if pl_module.is_logger():\n",
" metrics = trainer.callback_metrics\n",
"\n",
" # Log and save results to file\n",
" output_test_results_file = os.path.join(pl_module.hparams.output_dir, \"test_results.txt\")\n",
" with open(output_test_results_file, \"w\") as writer:\n",
" for key in sorted(metrics):\n",
" if key not in [\"log\", \"progress_bar\"]:\n",
" logger.info(\"{} = {}\\n\".format(key, str(metrics[key])))\n",
" writer.write(\"{} = {}\\n\".format(key, str(metrics[key])))"
],
"execution_count": 0,
"outputs": []
},
{
"cell_type": "markdown",
"metadata": {
"id": "a4hjvsBJ5Zk5",
"colab_type": "text"
},
"source": [
"Let's define the hyperparameters and other arguments. You can overide this `dict` for specific task as needed. While in most of cases you'll only need to change the `data_dir`and `output_dir`.\n",
"\n",
"Here the batch size is 8 and gradient_accumulation_steps are 16 so the effective batch size is 128"
]
},
{
"cell_type": "code",
"metadata": {
"id": "urduopvizqTq",
"colab_type": "code",
"colab": {}
},
"source": [
"args_dict = dict(\n",
" data_dir=\"\", # path for data files\n",
" output_dir=\"\", # path to save the checkpoints\n",
" model_name_or_path='t5-base',\n",
" tokenizer_name_or_path='t5-base',\n",
" max_seq_length=512,\n",
" learning_rate=3e-4,\n",
" weight_decay=0.0,\n",
" adam_epsilon=1e-8,\n",
" warmup_steps=0,\n",
" train_batch_size=8,\n",
" eval_batch_size=8,\n",
" num_train_epochs=2,\n",
" gradient_accumulation_steps=16,\n",
" n_gpu=1,\n",
" early_stop_callback=False,\n",
" fp_16=False, # if you want to enable 16-bit training then install apex and set this to true\n",
" opt_level='O1', # you can find out more on optimisation levels here https://nvidia.github.io/apex/amp.html#opt-levels-and-properties\n",
" max_grad_norm=1.0, # if you enable 16-bit training then set this to a sensible value, 0.5 is a good default\n",
" seed=42,\n",
")"
],
"execution_count": 0,
"outputs": []
},
{
"cell_type": "markdown",
"metadata": {
"id": "vfhlYUUV2NIh",
"colab_type": "text"
},
"source": [
"## IMDB review classification"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "b3C13iabZvwK",
"colab_type": "text"
},
"source": [
"### Download IMDB Data"
]
},
{
"cell_type": "code",
"metadata": {
"id": "7R0QdcgXuIWW",
"colab_type": "code",
"colab": {}
},
"source": [
"!wget https://ai.stanford.edu/~amaas/data/sentiment/aclImdb_v1.tar.gz\n",
"!tar -xvf aclImdb_v1.tar.gz"
],
"execution_count": 0,
"outputs": []
},
{
"cell_type": "code",
"metadata": {
"id": "ni1cAK7EvXSB",
"colab_type": "code",
"colab": {}
},
"source": [
"train_pos_files = glob.glob('aclImdb/train/pos/*.txt')\n",
"train_neg_files = glob.glob('aclImdb/train/neg/*.txt')"
],
"execution_count": 0,
"outputs": []
},
{
"cell_type": "code",
"metadata": {
"id": "jEsRn5pa0v8d",
"colab_type": "code",
"outputId": "6977ce56-d0b4-4d9f-8548-22003bb07eaf",
"colab": {
"base_uri": "https://localhost:8080/",
"height": 34
}
},
"source": [
"len(train_pos_files), len(train_neg_files)"
],
"execution_count": 0,
"outputs": [
{
"output_type": "execute_result",
"data": {
"text/plain": [
"(12500, 12500)"
]
},
"metadata": {
"tags": []
},
"execution_count": 10
}
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "5zgS8KhlaPiA",
"colab_type": "text"
},
"source": [
"We will use 2000 samples from the train set for validation. Let's choose 1000 postive reviews and 1000 negative reviews for validation and save them in the val directory"
]
},
{
"cell_type": "code",
"metadata": {
"id": "hLvBHcXwzXrk",
"colab_type": "code",
"colab": {}
},
"source": [
"!mkdir aclImdb/val aclImdb/val/pos aclImdb/val/neg"
],
"execution_count": 0,
"outputs": []
},
{
"cell_type": "code",
"metadata": {
"id": "IXZmLZ1pzjiY",
"colab_type": "code",
"colab": {}
},
"source": [
"random.shuffle(train_pos_files)\n",
"random.shuffle(train_neg_files)\n",
"\n",
"val_pos_files = train_pos_files[:1000]\n",
"val_neg_files = train_neg_files[:1000]"
],
"execution_count": 0,
"outputs": []
},
{
"cell_type": "code",
"metadata": {
"id": "5yTS2Jx40UNu",
"colab_type": "code",
"colab": {}
},
"source": [
"import shutil"
],
"execution_count": 0,
"outputs": []
},
{
"cell_type": "code",
"metadata": {
"id": "hJnJpkdb0ZKY",
"colab_type": "code",
"colab": {}
},
"source": [
"for f in val_pos_files:\n",
" shutil.move(f, 'aclImdb/val/pos')\n",
"for f in val_neg_files:\n",
" shutil.move(f, 'aclImdb/val/neg')"
],
"execution_count": 0,
"outputs": []
},
{
"cell_type": "markdown",
"metadata": {
"id": "qdEgCwL7cIyi",
"colab_type": "text"
},
"source": [
"### Prepare Dataset"
]
},
{
"cell_type": "code",
"metadata": {
"id": "McQC1FotigqA",
"colab_type": "code",
"outputId": "f60dbf68-32cf-44e1-9a2f-f9dba38cbbac",
"colab": {
"base_uri": "https://localhost:8080/",
"height": 186,
"referenced_widgets": [
"7d8f60bfc0a248e58028b6e8a477a5f7",
"72dc1e39b931429883e68c0603797896",
"cde60c5e18f04ba792fff8c2ac33f470",
"c0c0df12695b4a1eacf8fa4ccc0ac62c",
"72ea881ce3f445a9983d858b76dd257b",
"d0f0c28a14b242f8990a547ed7f87c04",
"f97741534b554be3b5cdccd45c73b317",
"1e70a3dc7090487fa883e932bff395cb"
]
}
},
"source": [
"tokenizer = T5Tokenizer.from_pretrained('t5-base')"
],
"execution_count": 0,
"outputs": [
{
"output_type": "stream",
"text": [
"INFO:filelock:Lock 139780871368544 acquired on /root/.cache/torch/transformers/68f1b8dbca4350743bb54b8c4169fd38cbabaad564f85a9239337a8d0342af9f.9995af32582a1a7062cb3173c118cb7b4636fa03feb967340f20fc37406f021f.lock\n",
"INFO:transformers.file_utils:https://s3.amazonaws.com/models.huggingface.co/bert/t5-spiece.model not found in cache or force_download set to True, downloading to /root/.cache/torch/transformers/tmpgy9lk1eo\n"
],
"name": "stderr"
},
{
"output_type": "display_data",
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "7d8f60bfc0a248e58028b6e8a477a5f7",
"version_minor": 0,
"version_major": 2
},
"text/plain": [
"HBox(children=(FloatProgress(value=0.0, description='Downloading', max=791656.0, style=ProgressStyle(descripti…"
]
},
"metadata": {
"tags": []
}
},
{
"output_type": "stream",
"text": [
"INFO:transformers.file_utils:storing https://s3.amazonaws.com/models.huggingface.co/bert/t5-spiece.model in cache at /root/.cache/torch/transformers/68f1b8dbca4350743bb54b8c4169fd38cbabaad564f85a9239337a8d0342af9f.9995af32582a1a7062cb3173c118cb7b4636fa03feb967340f20fc37406f021f\n",
"INFO:transformers.file_utils:creating metadata file for /root/.cache/torch/transformers/68f1b8dbca4350743bb54b8c4169fd38cbabaad564f85a9239337a8d0342af9f.9995af32582a1a7062cb3173c118cb7b4636fa03feb967340f20fc37406f021f\n",
"INFO:filelock:Lock 139780871368544 released on /root/.cache/torch/transformers/68f1b8dbca4350743bb54b8c4169fd38cbabaad564f85a9239337a8d0342af9f.9995af32582a1a7062cb3173c118cb7b4636fa03feb967340f20fc37406f021f.lock\n",
"INFO:transformers.tokenization_utils:loading file https://s3.amazonaws.com/models.huggingface.co/bert/t5-spiece.model from cache at /root/.cache/torch/transformers/68f1b8dbca4350743bb54b8c4169fd38cbabaad564f85a9239337a8d0342af9f.9995af32582a1a7062cb3173c118cb7b4636fa03feb967340f20fc37406f021f\n"
],
"name": "stderr"
},
{
"output_type": "stream",
"text": [
"\n"
],
"name": "stdout"
}
]
},
{
"cell_type": "code",
"metadata": {
"id": "wthd9SM74RG8",
"colab_type": "code",
"outputId": "52deb6bd-19c4-4071-8bcb-254925d8e4cc",
"colab": {
"base_uri": "https://localhost:8080/",
"height": 34
}
},
"source": [
"ids_neg = tokenizer.encode('negative ')\n",
"ids_pos = tokenizer.encode('positive ')\n",
"len(ids_neg), len(ids_pos)"
],
"execution_count": 0,
"outputs": [
{
"output_type": "execute_result",
"data": {
"text/plain": [
"(2, 2)"
]
},
"metadata": {
"tags": []
},
"execution_count": 21
}
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "k5sJkyI3a723",
"colab_type": "text"
},
"source": [
"All the examples are converted in the text-2-text format as shown in the paper. However I didn't use any task prefix here. The examples are encoded as follows,\n",
"if the review is positive then the target is 'positive' else 'negative'\n",
"\n",
"**input**: I went to see this\n",
"movie with my husband, and we both\n",
"thought the acting was terrible!\"\n",
"\n",
"**target**: negative\n",
"\n",
"**input**: Despite what others say,\n",
"I thought this movie was funny.\n",
"\n",
"**target**: positive"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "VEYmYHKGcxEq",
"colab_type": "text"
},
"source": [
"The dataset below takes care of reading the review files and processing the examples in text-2-text format.\n",
"\n",
"It cleans the review text by removing the html tags. It also appends the eos token `` at the end of input and target as required by the T5 model \n",
"\n",
"For T5 max input length is 512 and we can choose the max length for target sequence depending upon our dataset. The `T5Tokenizer` encodes both 'postive' and 'negative' as a single ids so I chose the max target length 2, extra 1 for the `` token"
]
},
{
"cell_type": "code",
"metadata": {
"id": "IIY0GenSb72m",
"colab_type": "code",
"colab": {}
},
"source": [
"class ImdbDataset(Dataset):\n",
" def __init__(self, tokenizer, data_dir, type_path, max_len=512):\n",
" self.pos_file_path = os.path.join(data_dir, type_path, 'pos')\n",
" self.neg_file_path = os.path.join(data_dir, type_path, 'neg')\n",
" \n",
" self.pos_files = glob.glob(\"%s/*.txt\" % self.pos_file_path)\n",
" self.neg_files = glob.glob(\"%s/*.txt\" % self.neg_file_path)\n",
" \n",
" self.max_len = max_len\n",
" self.tokenizer = tokenizer\n",
" self.inputs = []\n",
" self.targets = []\n",
"\n",
" self._build()\n",
" \n",
" def __len__(self):\n",
" return len(self.inputs)\n",
" \n",
" def __getitem__(self, index):\n",
" source_ids = self.inputs[index][\"input_ids\"].squeeze()\n",
" target_ids = self.targets[index][\"input_ids\"].squeeze()\n",
"\n",
" src_mask = self.inputs[index][\"attention_mask\"].squeeze() # might need to squeeze\n",
" target_mask = self.targets[index][\"attention_mask\"].squeeze() # might need to squeeze\n",
"\n",
" return {\"source_ids\": source_ids, \"source_mask\": src_mask, \"target_ids\": target_ids, \"target_mask\": target_mask}\n",
" \n",
" def _build(self):\n",
" self._buil_examples_from_files(self.pos_files, 'positive')\n",
" self._buil_examples_from_files(self.neg_files, 'negative')\n",
" \n",
" def _buil_examples_from_files(self, files, sentiment):\n",
" REPLACE_NO_SPACE = re.compile(\"[.;:!\\'?,\\\"()\\[\\]]\")\n",
" REPLACE_WITH_SPACE = re.compile(\"(| \n", " | text | \n", "emotion | \n", "
|---|---|---|
| 0 | \n", "i didnt feel humiliated | \n", "sadness | \n", "
| 1 | \n", "i can go from feeling so hopeless to so damned... | \n", "sadness | \n", "
| 2 | \n", "im grabbing a minute to post i feel greedy wrong | \n", "anger | \n", "
| 3 | \n", "i am ever feeling nostalgic about the fireplac... | \n", "love | \n", "
| 4 | \n", "i am feeling grouchy | \n", "anger | \n", "