Correcao_T5_NoTuning.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Correcao_T5_NoTuning.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU",
    "widgets": {
      "application/vnd.jupyter.widget-state+json": {
        "1a1cda60f820424987bf483644073c5e": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "state": {
            "_view_name": "HBoxView",
            "_dom_classes": [],
            "_model_name": "HBoxModel",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "box_style": "",
            "layout": "IPY_MODEL_9d9562be39fb485da6151955f74b22d4",
            "_model_module": "@jupyter-widgets/controls",
            "children": [
              "IPY_MODEL_4b06efc977314efcb15415974ce0b129",
              "IPY_MODEL_01ad8b879b3c49afab5ee8ac14d0362f"
            ]
          }
        },
        "9d9562be39fb485da6151955f74b22d4": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": "row wrap",
            "width": "100%",
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": "inline-flex",
            "left": null
          }
        },
        "4b06efc977314efcb15415974ce0b129": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "state": {
            "_view_name": "ProgressView",
            "style": "IPY_MODEL_065ed763dd0d477fa837e03362a835cc",
            "_dom_classes": [],
            "description": "Testing:   1%",
            "_model_name": "FloatProgressModel",
            "bar_style": "info",
            "max": 1,
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": 1,
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "orientation": "horizontal",
            "min": 0,
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_5d2df278d2734ee2a33a5b1e83672dd6"
          }
        },
        "01ad8b879b3c49afab5ee8ac14d0362f": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "state": {
            "_view_name": "HTMLView",
            "style": "IPY_MODEL_9b5436fbb17e4130a25e9140569e6442",
            "_dom_classes": [],
            "description": "",
            "_model_name": "HTMLModel",
            "placeholder": "​",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": " 240/20000 [2:32:40&lt;211:08:09, 38.47s/it]",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_f59f22994d244d12808e4544c4a0a38c"
          }
        },
        "065ed763dd0d477fa837e03362a835cc": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "ProgressStyleModel",
            "description_width": "initial",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "bar_color": null,
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "5d2df278d2734ee2a33a5b1e83672dd6": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": "2",
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "9b5436fbb17e4130a25e9140569e6442": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "DescriptionStyleModel",
            "description_width": "",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "f59f22994d244d12808e4544c4a0a38c": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        }
      }
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/fberanizo/spelling-correction/blob/master/Correcao_T5_NoTuning.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0Up4u8rMkQSG",
        "colab_type": "text"
      },
      "source": [
        "# Detecção e correção de erro ortográfico com modelo T5\n",
        "\n",
        "**Nome: Fabio Beranizo Lopes**<br>\n",
        "**Nome: Luiz Pita Almeida**\n",
        "\n",
        "Usaremos o modelo T5 pré-treinado e o dataset Paracrawl Inglês-Português. <br>\n",
        "Truncamos para strings de tamanho 100 para deixar os testes mais rápidos.\n",
        "\n",
        "Métrica de avaliação: F0.5-score <br>\n",
        "https://www.cl.cam.ac.uk/research/nl/bea2019st/#eval\n",
        "\n",
        "O método de correção aplicado foi sugerido pelos docentes:<br>\n",
        "> dado uma frase e um palavra nesta frase a ser corrigida ou não, iremos\n",
        "> mascarar a palavra, rodar o BERT ou T5, e prever as top-10 palavras \n",
        "> alternativas usando mask language modeling. Se a palavra original estiver \n",
        "> entre as top previstas, não sugerir correção. Caso contrário, usar edit \n",
        "> distance para ver qual é a palavra mais próxima, e sugerí-la ao usuário.\n",
        "\n",
        "Passos:\n",
        "\n",
        "1. Geram-se tuplas: `(original, corrected)`\n",
        "2. Aplica-se modelo T5 para prever top-10 palavras.<br>\n",
        "   Caso a palavra original esteja no Top 10 do modelo, é classificada como correta.<br>\n",
        "   Senão, a palavra é classificada como incorreta.\n",
        "\n",
        "**Obs: os notebooks contém excertos de códigos dos colegas de turma.**<br>\n",
        "**Obrigado Diedre, Gabriela, Leard, Lucas e Israel.**\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "CpELBvNmku5a",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        },
        "outputId": "c095c09d-a2f9-4723-e977-8c16e12ffb85"
      },
      "source": [
        "import torch\n",
        "\n",
        "print(f\"Current GPU: {torch.cuda.get_device_name(0)}\")\n",
        "\n",
        "# don't even start if it's not a P100 GPU\n",
        "# if torch.cuda.get_device_name(0) != \"Tesla P100-PCIE-16GB\":\n",
        "#     import os\n",
        "#     os.kill(os.getpid(), 9)"
      ],
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Current GPU: Tesla P100-PCIE-16GB\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "FgW-boJLU0wU",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "#@title Configurações gerais\n",
        "experiment_name = \"no-tuning\"  #@param {type:\"string\"}\n",
        "model_name = \"t5-small\"  #@param [\"t5-small\", \"t5-base\", \"t5-large\", \"t5-3B\", \"t5-11B\"] {type:\"string\"}\n",
        "batch_size = 10  #@param {type:\"integer\"}\n",
        "accumulate_grad_batches = 1  #@param {type:\"integer\"}\n",
        "sequence_length = 100  #@param {type:\"integer\"}\n",
        "learning_rate = 5e-3  #@param {type:\"number\"}\n",
        "decode_mode = \"topk\"  #@param [\"greedy\", \"nucleus\", \"topk\", \"beam\"] {type:\"string\"}\n",
        "k = 10  #@param {type:\"integer\"}"
      ],
      "execution_count": 2,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1U0dTSB-mnGN",
        "colab_type": "text"
      },
      "source": [
        "## Instala dependências\n",
        "\n",
        "- PyTorch Lightning\n",
        "- Hugginface Transformers\n",
        "- ERRANT (ERRor ANnotation Toolkit)\n",
        "- pyxDamerauLevenshtein"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "OOXZxjWRkLMM",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        },
        "outputId": "d33d837c-43e8-40a5-bc9a-49ad8e5556b8"
      },
      "source": [
        "!git clone --quiet https://github.com/fberanizo/Adversarial-Misspellings.git\n",
        "\n",
        "try:\n",
        "    import pytorch_lightning\n",
        "    import transformers\n",
        "except ImportError as e:\n",
        "    # can't import modules, then install\n",
        "    !pip install --quiet pytorch-lightning\n",
        "    !pip install --quiet transformers\n",
        "    !pip install --quiet errant==2.0.0\n",
        "    !pip install pyxDamerauLevenshtein\n",
        "    !python -m spacy download en\n",
        "    # kill kernel (necessary for tqdm)\n",
        "    import os\n",
        "    os.kill(os.getpid(), 9)"
      ],
      "execution_count": 3,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "fatal: destination path 'Adversarial-Misspellings' already exists and is not an empty directory.\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ob7qL6kUVjbu",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 51
        },
        "outputId": "c19590a6-62f7-468a-86e8-753aed74ffdb"
      },
      "source": [
        "# Importar todos os pacotes de uma só vez para evitar duplicados ao longo do notebook.\n",
        "import datetime\n",
        "import errant\n",
        "import gzip\n",
        "import json\n",
        "import numpy as np\n",
        "import nvidia_smi\n",
        "import os\n",
        "import pandas as pd\n",
        "import joblib\n",
        "import psutil\n",
        "import pytorch_lightning as pl\n",
        "import random\n",
        "import spacy\n",
        "import sys\n",
        "import tarfile\n",
        "import tempfile\n",
        "import torch\n",
        "import torch.nn.functional as F\n",
        "\n",
        "from argparse import Namespace\n",
        "from collections import deque\n",
        "from google.colab import drive\n",
        "from itertools import cycle\n",
        "\n",
        "from pyxdameraulevenshtein import damerau_levenshtein_distance, \\\n",
        "    normalized_damerau_levenshtein_distance\n",
        "from pyxdameraulevenshtein import damerau_levenshtein_distance_ndarray, \\\n",
        "    normalized_damerau_levenshtein_distance_ndarray\n",
        "\n",
        "from pytorch_lightning.callbacks import ModelCheckpoint\n",
        "from pytorch_lightning.loggers import TensorBoardLogger\n",
        "from pytorch_lightning import Trainer\n",
        "\n",
        "\n",
        "from tqdm import tqdm\n",
        "from transformers import T5ForConditionalGeneration\n",
        "from transformers import T5Tokenizer\n",
        "from torch.optim import Adam\n",
        "from torch.utils.data import DataLoader\n",
        "from torch.utils.data import Dataset\n",
        "\n",
        "from typing import Dict\n",
        "from typing import List\n",
        "from typing import Tuple\n",
        "\n",
        "# Leard decoding solution\n",
        "import html\n",
        "import unicodedata\n",
        "\n",
        "nlp = spacy.load(\"en\")\n",
        "annotator = errant.load(\"en\", nlp)\n",
        "\n",
        "import nltk\n",
        "nltk.download(\"stopwords\")\n",
        "\n",
        "sys.path.insert(0, \"/content/Adversarial-Misspellings\")\n",
        "import attacks\n",
        "\n",
        "nvidia_smi.nvmlInit()\n",
        "handle = nvidia_smi.nvmlDeviceGetHandleByIndex(0)\n",
        "\n",
        "def hardware_stats():\n",
        "    \"\"\"\n",
        "    Returns a dict containing some hardware related stats\n",
        "    \"\"\"\n",
        "    res = nvidia_smi.nvmlDeviceGetUtilizationRates(handle)\n",
        "    return {\"cpu\": f\"{str(psutil.cpu_percent())}%\",\n",
        "            \"mem\": f\"{str(psutil.virtual_memory().percent)}%\",\n",
        "            \"gpu\": f\"{str(res.gpu)}%\",\n",
        "            \"gpu_mem\": f\"{str(res.memory)}%\"}"
      ],
      "execution_count": 4,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "[nltk_data] Downloading package stopwords to /root/nltk_data...\n",
            "[nltk_data]   Package stopwords is already up-to-date!\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Tm8H4Rfwm2gE",
        "colab_type": "text"
      },
      "source": [
        "## Define random seeds\n",
        "\n",
        "Importante: Fix seeds so we can replicate results"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "bJlZDb1VY29r",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import random\n",
        "\n",
        "seed = 0\n",
        "random.seed(seed)\n",
        "torch.random.manual_seed(seed)\n",
        "torch.cuda.manual_seed(seed)"
      ],
      "execution_count": 5,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ETfkvMGl4JA1",
        "colab_type": "text"
      },
      "source": [
        "## Mapeia Google Drive\n",
        "\n",
        "Iremos salvar os checkpoints (pesos do modelo) no google drive, para que possamos continuar o treino de onde paramos."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "hvGjweSKfbA2",
        "colab": {}
      },
      "source": [
        "# drive.mount(\"/content/drive\")\n",
        "# base_path = \"/content/drive/My Drive/PF-Correcao/t5-no-tuning\"\n",
        "base_path = \"/content/t5-no-tuning\"\n",
        "os.environ[\"BASE_PATH\"] = base_path"
      ],
      "execution_count": 6,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BMhp7hvN9mS6",
        "colab_type": "text"
      },
      "source": [
        "## ERRANT Scorer\n",
        "\n",
        "Comando para avaliação que compara um arquivo M2 \"hipótese\" contra um arquivo M2 \"referência\".<br>\n",
        "\n",
        "### **Example**\n",
        "**Original**: This are gramamtical sentence .<br>\n",
        "**Corrected**: This is a grammatical sentence .<br>\n",
        "**Output M2**:<br>\n",
        "S This are gramamtical sentence .<br>\n",
        "A 1 2|||R:VERB:SVA|||is|||REQUIRED|||-NONE-|||0<br>\n",
        "A 2 2|||M:DET|||a|||REQUIRED|||-NONE-|||0<br>\n",
        "A 2 3|||R:SPELL|||grammatical|||REQUIRED|||-NONE-|||0<br>\n",
        "A -1 -1|||noop|||-NONE-|||REQUIRED|||-NONE-|||1<br>\n",
        "\n",
        "In M2 format, a line preceded by S denotes an original sentence while a line preceded by A indicates an edit annotation. Each edit line consists of the start and end token offset of the edit, the error type, and the tokenized correction string. The next two fields are included for historical reasons (see the CoNLL-2014 shared task) while the last field is the annotator id."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "aWPWngE89ss_",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        },
        "outputId": "95c6e0d8-069d-4c20-ed73-8a88ad6d990b"
      },
      "source": [
        "%%writefile orig.txt\n",
        "Eu não cei pra ondi vou .\n",
        "Podi até num dá em nada .\n",
        "Minha vida segui o sol .\n",
        "No horizonti dessa istrada ."
      ],
      "execution_count": 7,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Overwriting orig.txt\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Dkdn9fTBS4lk",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        },
        "outputId": "38122cd9-c539-42e4-eac4-9cb701e5b033"
      },
      "source": [
        "%%writefile ref.txt\n",
        "Eu não sei pra onde vou .\n",
        "Pode até não dar em nada .\n",
        "Minha vida segue o sol .\n",
        "No horizonte dessa estrada ."
      ],
      "execution_count": 8,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Overwriting ref.txt\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "5eQ142HVSZu3",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        },
        "outputId": "0c0f7968-da7a-43ec-a82c-b42a156c77a9"
      },
      "source": [
        "%%writefile hyp.txt\n",
        "Eu não sei pra ondi vou .\n",
        "Podi até não dar em nada .\n",
        "Minha vida segui u sol .\n",
        "Num horizonte dessa estrada ."
      ],
      "execution_count": 9,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Overwriting hyp.txt\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "a6KxhquG8U44",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!errant_parallel -orig orig.txt -cor ref.txt -out ref.m2 > /dev/null\n",
        "!errant_parallel -orig orig.txt -cor hyp.txt -out hyp.m2 > /dev/null"
      ],
      "execution_count": 10,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Qj0qyFWK9pI3",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 119
        },
        "outputId": "2d0ad5f1-5458-44bb-b1d3-d0652fe6fa79"
      },
      "source": [
        "import pandas as pd\n",
        "!errant_compare -hyp hyp.m2 -ref ref.m2\n",
        "# x = !errant_compare -hyp hyp.m2 -ref ref.m2\n",
        "# df = pd.DataFrame(data=x[2:4])[0].str.split('\\t', expand=True)\n",
        "# new_header = df.iloc[0] #grab the first row for the header\n",
        "# df = df[1:] #take the data less the header row\n",
        "# df.columns = new_header\n",
        "# df\n",
        "# df[\"F0.5\"][1]\n",
        "# d = df.apply(pd.to_numeric).to_dict('r')\n"
      ],
      "execution_count": 11,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "\n",
            "=========== Span-Based Correction ============\n",
            "TP\tFP\tFN\tPrec\tRec\tF0.5\n",
            "5\t2\t3\t0.7143\t0.625\t0.6944\n",
            "==============================================\n",
            "\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0W1V2fdg-RqU",
        "colab_type": "text"
      },
      "source": [
        "## Gerador \n",
        "\n",
        "From:\n",
        "https://github.com/huggingface/transformers/issues/3985"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "YmxoZCZoIq6N",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# tokenizer = T5Tokenizer.from_pretrained(model_name)\n",
        "# model = T5ForConditionalGeneration.from_pretrained(model_name)\n",
        "# # Input text\n",
        "# original = 'This are gramamtical sentence .'\n",
        "# correct = 'This is a grammatical sentence .'\n",
        "\n",
        "# text = 'This <extra_id_0> sentence. </s>'\n",
        "\n",
        "# encoded = tokenizer.encode_plus(text, add_special_tokens=True, return_tensors='pt')\n",
        "# input_ids = encoded['input_ids']\n",
        "\n",
        "# # Generating 20 sequences with maximum length set to 5\n",
        "# outputs = model.generate(input_ids=input_ids, \n",
        "#                           num_beams=200, num_return_sequences=20,\n",
        "#                           max_length=5)\n",
        "\n",
        "# _0_index = text.index('<extra_id_0>')\n",
        "# _result_prefix = text[:_0_index]\n",
        "# _result_suffix = text[_0_index+12:]  # 12 is the length of <extra_id_0>\n",
        "\n",
        "# def _filter(output, end_token='<extra_id_1>'):\n",
        "#     # The first token is <unk> (inidex at 0) and the second token is <extra_id_0> (indexed at 32099)\n",
        "#     _txt = tokenizer.decode(output[2:], skip_special_tokens=False, clean_up_tokenization_spaces=False)\n",
        "#     if end_token in _txt:\n",
        "#         _end_token_index = _txt.index(end_token)\n",
        "#         return _result_prefix + _txt[:_end_token_index] + _result_suffix\n",
        "#     else:\n",
        "#         return _result_prefix + _txt + _result_suffix\n",
        "\n",
        "# results = list(map(_filter, outputs))\n",
        "# results\n",
        "# del tokenizer\n",
        "# del model"
      ],
      "execution_count": 12,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "TgUkD55Dkam2"
      },
      "source": [
        "## Classe Dataset\n",
        "Gerenciamento dos dados, e um pequeno teste."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "R2GaMtgRkaze",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 697
        },
        "outputId": "3b3d6085-174b-466e-87e9-50d555d63671"
      },
      "source": [
        "hparams = {\"model_name\": model_name, \"seq_len\": sequence_length, \"batch_size\": batch_size}\n",
        "class ParaCrawl(Dataset):\n",
        "    \"\"\"\n",
        "    Loads data from preprocessed file and manages them.\n",
        "    \"\"\"\n",
        "    VALID_MODES = [\"train\", \"validation\", \"test\"]\n",
        "    TOKENIZER = T5Tokenizer.from_pretrained(hparams[\"model_name\"],\n",
        "                                            cache_dir=base_path)\n",
        "    def __init__(self, mode: str, seq_len: int):\n",
        "        \"\"\"\n",
        "        mode: One of train, validation or test \n",
        "        seq_len: limit to returned encoded tokens\n",
        "        \"\"\"\n",
        "        super().__init__()\n",
        "        assert mode in ParaCrawl.VALID_MODES\n",
        "\n",
        "        self.mode = mode\n",
        "        self.seq_len = seq_len\n",
        "\n",
        "        file_name = os.path.join(base_path, f\"{mode}.pkl\")\n",
        "        if not os.path.isfile(file_name):\n",
        "            print(\"Pre-processed files not found, preparing data.\")\n",
        "            self.prepare_data()\n",
        "        \n",
        "        with open(file_name, \"rb\") as preprocessed_file:\n",
        "            self.data = joblib.load(preprocessed_file)\n",
        "\n",
        "    def __len__(self):\n",
        "        return len(self.data)\n",
        "\n",
        "    def __getitem__(self, i: int):\n",
        "        \"\"\"\n",
        "        Unpacks line from data.\n",
        "\n",
        "        returns: input (corrputed), target (corrected)\n",
        "        \"\"\"\n",
        "        input, target = self.data[i]\n",
        "\n",
        "        return input, target\n",
        "\n",
        "    def get_dataloader(self, batch_size: int, shuffle: bool):\n",
        "        return DataLoader(self, batch_size=batch_size, shuffle=shuffle, \n",
        "                          num_workers=4)\n",
        "\n",
        "    @staticmethod\n",
        "    def load_text_pairs(path):\n",
        "        \"\"\"\n",
        "        Load pairs from original files, selects pt, then corrupts the samples.\n",
        "        \"\"\"\n",
        "        text_pairs = []\n",
        "        for line in tqdm(gzip.open(path, mode=\"rt\")):\n",
        "            text_pt = line.strip().split(\"\\t\")[1]\n",
        "            text_pt = text_pt[:hparams[\"seq_len\"]].rsplit(\" \", 1)[0]\n",
        "            try:\n",
        "                attack_list = deque(attacks.all_one_attack(text_pt, include_ends=True))\n",
        "                text_corrupt = list(map(lambda a: a[1], random.sample(attack_list, k=10)))\n",
        "                text_pairs.extend(list(zip(text_corrupt, cycle([text_pt]))))\n",
        "            except ValueError:\n",
        "                pass\n",
        "\n",
        "        return text_pairs\n",
        "\n",
        "    @staticmethod\n",
        "    def prepare_data(train_size=9997000, val_size=3000):\n",
        "        \"\"\"\n",
        "        Performs everything needed to get the data ready.\n",
        "        Addition of Eos token and encoding is performed in runtime.\n",
        "        \"\"\"\n",
        "        if not os.path.isfile(\"paracrawl_enpt_train.tsv.gz\"):    \n",
        "            !wget -nc https://storage.googleapis.com/neuralresearcher_data/unicamp/ia376e_2020s1/paracrawl_enpt_train.tsv.gz -P \"$BASE_PATH\"\n",
        "            !wget -nc https://storage.googleapis.com/neuralresearcher_data/unicamp/ia376e_2020s1/paracrawl_enpt_test.tsv.gz -P \"$BASE_PATH\"\n",
        "\n",
        "        data = {}\n",
        "        test_data = ParaCrawl.load_text_pairs(os.path.join(base_path, \"paracrawl_enpt_test.tsv.gz\"))\n",
        "        train_val_data = ParaCrawl.load_text_pairs(os.path.join(base_path, \"paracrawl_enpt_train.tsv.gz\"))\n",
        "\n",
        "        random.shuffle(train_val_data)\n",
        "\n",
        "        train_data = train_val_data[:train_size]\n",
        "        val_data = train_val_data[train_size:train_size + val_size]\n",
        "\n",
        "        for mode, data in zip(ParaCrawl.VALID_MODES, [train_data, val_data, test_data]):\n",
        "            file_name = os.path.join(base_path, f\"{mode}.pkl\")\n",
        "            with open(file_name, \"wb\") as pkl_file:\n",
        "                joblib.dump(data, pkl_file)\n",
        "            print(f\"Pre-processed data saved as {file_name}.\")\n",
        "\n",
        "\n",
        "datasets = {m: ParaCrawl(mode=m, seq_len=hparams[\"seq_len\"]) for m in ParaCrawl.VALID_MODES}\n",
        "\n",
        "# Testing datasets\n",
        "for mode, dataset in datasets.items():\n",
        "    print(f\"\\n{mode} dataset length: {len(dataset)}\\n\")\n",
        "    print(\"Random sample\")\n",
        "    input, target = random.choice(dataset)\n",
        "    print(\"input\\n\", input, end=\"\\n\\n\")\n",
        "    print(\"target\\n\", target, end=\"\\n\\n\")"
      ],
      "execution_count": 13,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Pre-processed files not found, preparing data.\n",
            "File ‘/content/t5-no-tuning/paracrawl_enpt_train.tsv.gz’ already there; not retrieving.\n",
            "\n",
            "File ‘/content/t5-no-tuning/paracrawl_enpt_test.tsv.gz’ already there; not retrieving.\n",
            "\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "20000it [00:36, 552.64it/s]\n",
            "1000000it [30:16, 550.47it/s]\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "Pre-processed data saved as /content/t5-no-tuning/train.pkl.\n",
            "Pre-processed data saved as /content/t5-no-tuning/validation.pkl.\n",
            "Pre-processed data saved as /content/t5-no-tuning/test.pkl.\n",
            "\n",
            "train dataset length: 9997000\n",
            "\n",
            "Random sample\n",
            "input\n",
            " Será justas uma sociedade que não aceita as diferenças?, será o amor mais forte que as\n",
            "\n",
            "target\n",
            " Será justa uma sociedade que não aceita as diferenças?, será o amor mais forte que as\n",
            "\n",
            "\n",
            "validation dataset length: 2670\n",
            "\n",
            "Random sample\n",
            "input\n",
            " História da arte (1) Epistemologia e rmétodos\n",
            "\n",
            "target\n",
            " História da arte (1) Epistemologia e métodos\n",
            "\n",
            "\n",
            "test dataset length: 200000\n",
            "\n",
            "Random sample\n",
            "input\n",
            " Uma vaga que npão faz\n",
            "\n",
            "target\n",
            " Uma vaga que não faz\n",
            "\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cloyt0tIwIiD",
        "colab_type": "text"
      },
      "source": [
        "## Dataloaders\n",
        "\n",
        "Verificação se dataloaders estão funcionando corretamente."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ZoKiQXCvwGrP",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 68
        },
        "outputId": "bcfbed5d-ea8e-4593-a423-279e52103d6a"
      },
      "source": [
        "shuffle = {\"train\": True, \"validation\": False, \"test\": False}\n",
        "debug_dataloaders = {mode: datasets[mode].get_dataloader(batch_size=hparams[\"batch_size\"], \n",
        "                                                         shuffle=shuffle[mode])\n",
        "                     for mode in ParaCrawl.VALID_MODES}\n",
        "\n",
        "# Testing dataloaders\n",
        "for mode, dataloader in debug_dataloaders.items():\n",
        "    print(\"{} number of batches: {}\".format(mode, len(dataloader)))\n",
        "    batch = next(iter(dataloader))"
      ],
      "execution_count": 14,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "train number of batches: 999700\n",
            "validation number of batches: 267\n",
            "test number of batches: 20000\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YGH2rc3lthSB",
        "colab_type": "text"
      },
      "source": [
        "## Lightning Module\n",
        "\n",
        "Aqui a classe principal do PyTorch Lightning é definida.\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "jg-hZEktbvnr",
        "colab": {}
      },
      "source": [
        "class T5Corrector(pl.LightningModule):\n",
        "    def __init__(self, hparams):\n",
        "        super().__init__()\n",
        "\n",
        "        self.hparams = hparams\n",
        "        self.t5 = T5ForConditionalGeneration.from_pretrained(hparams.model_name,\n",
        "                                                             cache_dir=hparams.base_path)\n",
        "        self.tokenizer = ParaCrawl.TOKENIZER\n",
        "        self.start_token = ParaCrawl.TOKENIZER.convert_tokens_to_ids('<extra_id_0>')\n",
        "        self.end_token = ParaCrawl.TOKENIZER.convert_tokens_to_ids('<extra_id_1>')\n",
        "\n",
        "    def select_correction(self, word, hypotheses):\n",
        "        \"\"\"\n",
        "        Selects the most probable correction for a given word and given hypotheses.\n",
        "        \"\"\"\n",
        "        if word in hypotheses:\n",
        "            # print('Only copy:', word)\n",
        "            return word\n",
        "        else:\n",
        "            distances = damerau_levenshtein_distance_ndarray(word, np.array(hypotheses))\n",
        "            # print(distances.shape)\n",
        "            # print('distances', distances)\n",
        "            idx_min_distance = np.argmin(distances)\n",
        "            # print('idx_min_distance', idx_min_distance)\n",
        "            # print('idx_min_distance distance', distances[idx_min_distance])\n",
        "            # print('idx_min_distance', hypotheses[idx_min_distance])\n",
        "            if distances[idx_min_distance] < 10:\n",
        "                # print('Replaced:', word, 'by', hypotheses[idx_min_distance])\n",
        "                return hypotheses[idx_min_distance]\n",
        "        # print('No suggestion, word copied:', word)\n",
        "        return word\n",
        "\n",
        "    def generate(self, original, end_token=\"<extra_id_1>\"):\n",
        "        \"\"\"\n",
        "        Generates a correction hypothesis for a given sentence.\n",
        "        \"\"\"\n",
        "        hypothesis = \"\"\n",
        "        words = original.split()\n",
        "        # we could not fit n masked_sentences in memory:\n",
        "        # masked_ids = []\n",
        "        # input_ids = torch.stack(masked_ids).squeeze().to(\"cuda\")\n",
        "        # masked_ids.append(input_ids)\n",
        "\n",
        "        for idx, word in enumerate(words):\n",
        "            # masks the i-th word\n",
        "            input = f\"{words[:idx]} <extra_id_0> {words[idx + 1:]} {self.tokenizer.eos_token}\"\n",
        "            input_ids = self.tokenizer.encode(input,\n",
        "                                              max_length=self.hparams.seq_len,\n",
        "                                              pad_to_max_length=True,\n",
        "                                              add_special_tokens=True,\n",
        "                                              return_tensors=\"pt\").to(\"cuda\")\n",
        "\n",
        "            # generates k hypothesis for the sentence\n",
        "            hypotheses = self.t5.generate(input_ids=input_ids,\n",
        "                                        top_k=self.hparams.k,\n",
        "                                        do_sample=True,\n",
        "                                        num_return_sequences=self.hparams.k)\n",
        "            hypotheses = [self.tokenizer.decode(output_ids[2:],\n",
        "                                                skip_special_tokens=False,\n",
        "                                                clean_up_tokenization_spaces=False) for output_ids in hypotheses]\n",
        "\n",
        "            hypothesis_word = self.select_correction(word, hypotheses)\n",
        "            hypothesis = f\"{hypothesis} {hypothesis_word}\"\n",
        "\n",
        "        return hypothesis\n",
        "\n",
        "    def forward(self, x):\n",
        "        inputs, targets = x\n",
        "\n",
        "        if self.training:\n",
        "            input_ids = []\n",
        "            attention_mask = []\n",
        "            lm_labels = []\n",
        "            # for each sample in batch\n",
        "            for input, target in zip(inputs, targets):\n",
        "                input = f\"{input} {self.tokenizer.eos_token}\"\n",
        "                d = self.tokenizer.encode_plus(\n",
        "                    input,\n",
        "                    max_length=self.hparams.seq_len,\n",
        "                    pad_to_max_length=True,\n",
        "                    add_special_tokens=True,\n",
        "                    return_tensors=\"pt\").to(\"cuda\")\n",
        "                input_ids.append(d[\"input_ids\"])\n",
        "                attention_mask.append(d[\"attention_mask\"])\n",
        "                lm_labels.append(self.tokenizer.encode(\n",
        "                    target,\n",
        "                    max_length=self.hparams.seq_len,\n",
        "                    pad_to_max_length=True,\n",
        "                    add_special_tokens=True,\n",
        "                    return_tensors=\"pt\"))\n",
        "\n",
        "            input_ids = torch.stack(input_ids).squeeze(1).to(\"cuda\")\n",
        "            attention_mask = torch.stack(attention_mask).squeeze(1).to(\"cuda\")\n",
        "            lm_labels = torch.stack(lm_labels).squeeze(1).to(\"cuda\")\n",
        "\n",
        "            outputs = self.t5(input_ids=input_ids, \n",
        "                              attention_mask=attention_mask,\n",
        "                              lm_labels=lm_labels)\n",
        "            loss, predicted_scores = outputs[:2]\n",
        "            return loss, predicted_scores, inputs, targets\n",
        "        else:\n",
        "            predicts = []\n",
        "            # for each sample in batch\n",
        "            for input, target in zip(inputs, targets):\n",
        "                predicts.append(self.generate(input))\n",
        "                # print(\"original\", orig)\n",
        "                # print(\"corrected\", corr)\n",
        "                # print(\"outputs\", outputs[-1])\n",
        "            return predicts, inputs, targets\n",
        "\n",
        "    def training_step(self, batch, batch_idx):\n",
        "        loss, predicted_scores, inputs, targets = self(batch)\n",
        "\n",
        "        return {\"loss\": loss, \"log\": {\"loss\": loss}, \"progress_bar\": hardware_stats()}\n",
        "\n",
        "    def validation_step(self, batch, batch_idx):\n",
        "        predicts, inputs, targets = self(batch)\n",
        "\n",
        "        with open(\"orig.txt\", \"w\") as f:\n",
        "            for input in inputs:\n",
        "                input = input.replace(\"\\n\", \"\")\n",
        "                f.write(f\"{input}\\n\")\n",
        "\n",
        "        with open(\"ref.txt\", \"w\") as f:\n",
        "            for target in targets:\n",
        "                target = target.replace(\"\\n\", \"\")\n",
        "                f.write(f\"{target}\\n\")\n",
        "\n",
        "        with open(\"hyp.txt\", \"w\") as f:\n",
        "            for pred in predicts:\n",
        "                pred = pred.replace(\"\\n\", \"\")\n",
        "                f.write(f\"{pred}\\n\")\n",
        "\n",
        "        !errant_parallel -orig orig.txt -cor ref.txt -out ref.m2 > /dev/null\n",
        "        !errant_parallel -orig orig.txt -cor hyp.txt -out hyp.m2 > /dev/null\n",
        "        x = !errant_compare -hyp hyp.m2 -ref ref.m2\n",
        "        df = pd.DataFrame(data=x[2:4])[0].str.split(\"\\t\", expand=True)\n",
        "        new_header = df.iloc[0] #grab the first row for the header\n",
        "        df = df[1:].apply(pd.to_numeric) #take the data less the header row\n",
        "        df.columns = new_header\n",
        "\n",
        "        true_positive = df[\"TP\"][1]\n",
        "        false_positive = df[\"FP\"][1]\n",
        "        false_negative = df[\"FN\"][1]\n",
        "        precision = df[\"Prec\"][1]\n",
        "        recall = df[\"Rec\"][1]\n",
        "        f_score = df[\"F0.5\"][1]\n",
        "\n",
        "        progress_bar = hardware_stats()\n",
        "        progress_bar.update({\"precision\": precision, \"recall\": recall, \"f_score\": f_score})\n",
        "        # print(\"progress_bar\", progress_bar)\n",
        "\n",
        "        return {\"true_positive\": true_positive, \"false_positive\": false_positive, \"false_negative\": false_negative,\n",
        "                \"precision\": precision, \"recall\": recall, \"f_score\": f_score,\n",
        "                \"predicts\": predicts, \"inputs\": inputs, \"targets\": targets, \"progress_bar\": progress_bar}\n",
        "\n",
        "    def test_step(self, batch, batch_idx):\n",
        "        predicts, inputs, targets = self(batch)\n",
        "\n",
        "        with open(\"orig.txt\", \"w\") as f:\n",
        "            for input in inputs:\n",
        "                input = input.replace(\"\\n\", \"\")\n",
        "                f.write(f\"{input}\\n\")\n",
        "\n",
        "        with open(\"ref.txt\", \"w\") as f:\n",
        "            for target in targets:\n",
        "                target = target.replace(\"\\n\", \"\")\n",
        "                f.write(f\"{target}\\n\")\n",
        "\n",
        "        with open(\"hyp.txt\", \"w\") as f:\n",
        "            for pred in predicts:\n",
        "                pred = pred.replace(\"\\n\", \"\")\n",
        "                f.write(f\"{pred}\\n\")\n",
        "\n",
        "        !errant_parallel -orig orig.txt -cor ref.txt -out ref.m2 > /dev/null\n",
        "        !errant_parallel -orig orig.txt -cor hyp.txt -out hyp.m2 > /dev/null\n",
        "        x = !errant_compare -hyp hyp.m2 -ref ref.m2\n",
        "        df = pd.DataFrame(data=x[2:4])[0].str.split(\"\\t\", expand=True)\n",
        "        new_header = df.iloc[0] #grab the first row for the header\n",
        "        df = df[1:].apply(pd.to_numeric) #take the data less the header row\n",
        "        df.columns = new_header\n",
        "\n",
        "        true_positive = df[\"TP\"][1]\n",
        "        false_positive = df[\"FP\"][1]\n",
        "        false_negative = df[\"FN\"][1]\n",
        "        precision = df[\"Prec\"][1]\n",
        "        recall = df[\"Rec\"][1]\n",
        "        f_score = df[\"F0.5\"][1]\n",
        "\n",
        "        progress_bar = hardware_stats()\n",
        "        progress_bar.update({\"precision\": precision, \"recall\": recall, \"f_score\": f_score})\n",
        "        # print(\"progress_bar\", progress_bar)\n",
        "\n",
        "        return {\"true_positive\": true_positive, \"false_positive\": false_positive, \"false_negative\": false_negative,\n",
        "                \"precision\": precision, \"recall\": recall, \"f_score\": f_score,\n",
        "                \"predicts\": predicts, \"inputs\": inputs, \"targets\": targets, \"progress_bar\": progress_bar}\n",
        "\n",
        "    def training_epoch_end(self, outputs):\n",
        "        avg_loss = torch.stack([x[\"loss\"] for x in outputs]).mean()\n",
        "\n",
        "        return {\"log\": {\"train_loss\": avg_loss}} \n",
        "\n",
        "    def validation_epoch_end(self, outputs):\n",
        "        avg_precision = sum([x[\"precision\"] for x in outputs]) / len(outputs)\n",
        "        avg_recall = sum([x[\"recall\"] for x in outputs]) / len(outputs)\n",
        "        avg_f_score = sum([x[\"f_score\"] for x in outputs]) / len(outputs)\n",
        "\n",
        "        tensorboard_logs = {\"avg_precision\": avg_precision,\n",
        "                            \"avg_recall\": avg_recall,\n",
        "                            \"avg_f_score\": avg_f_score}\n",
        "\n",
        "        origs = sum([list(x[\"inputs\"]) for x in outputs], [])\n",
        "        trues = sum([list(x[\"targets\"]) for x in outputs], [])\n",
        "        preds = sum([list(x[\"predicts\"]) for x in outputs], [])\n",
        "\n",
        "        n = random.choice(range(len(trues)))\n",
        "        print(f\"\\Input: {origs[n]}\\nTarget: {trues[n]}\\nPrediction: {preds[n]}\\n\")\n",
        "\n",
        "        return {\"avg_precision\": avg_precision, \"avg_recall\": avg_recall, \"avg_f_score\": avg_f_score,\n",
        "                \"log\": tensorboard_logs, \"progress_bar\": tensorboard_logs}\n",
        "\n",
        "    def test_epoch_end(self, outputs):\n",
        "        avg_precision = sum([x[\"precision\"] for x in outputs]) / len(outputs)\n",
        "        avg_recall = sum([x[\"recall\"] for x in outputs]) / len(outputs)\n",
        "        avg_f_score = sum([x[\"f_score\"] for x in outputs]) / len(outputs)\n",
        "\n",
        "        tensorboard_logs = {\"avg_precision\": avg_precision,\n",
        "                            \"avg_recall\": avg_recall,\n",
        "                            \"avg_f_score\": avg_f_score}\n",
        "\n",
        "        origs = sum([list(x[\"inputs\"]) for x in outputs], [])\n",
        "        trues = sum([list(x[\"targets\"]) for x in outputs], [])\n",
        "        preds = sum([list(x[\"predicts\"]) for x in outputs], [])\n",
        "\n",
        "        n = random.choice(range(len(trues)))\n",
        "        print(f\"\\Input: {origs[n]}\\nTarget: {trues[n]}\\nPrediction: {preds[n]}\\n\")\n",
        "        \n",
        "        return {\"avg_precision\": avg_precision, \"avg_recall\": avg_recall, \"avg_f_score\": avg_f_score,\n",
        "                \"log\": tensorboard_logs, \"progress_bar\": tensorboard_logs}\n",
        "\n",
        "    def configure_optimizers(self):\n",
        "        return Adam(self.parameters(), lr=self.hparams.lr)    \n",
        "\n",
        "    def train_dataloader(self):\n",
        "        if self.hparams.overfit_pct > 0:\n",
        "            logging.info(\"Disabling train shuffle due to overfit_pct.\")\n",
        "            shuffle = False\n",
        "        else:\n",
        "            shuffle = True\n",
        "        dataset = ParaCrawl(\"train\", seq_len=self.hparams.seq_len)\n",
        "        return dataset.get_dataloader(batch_size=self.hparams.batch_size, shuffle=shuffle)\n",
        "\n",
        "    def val_dataloader(self):\n",
        "        dataset = ParaCrawl(\"validation\", seq_len=self.hparams.seq_len)\n",
        "        return dataset.get_dataloader(batch_size=self.hparams.batch_size, shuffle=False)\n",
        "\n",
        "    def test_dataloader(self):\n",
        "        dataset = ParaCrawl(\"test\", seq_len=self.hparams.seq_len)\n",
        "        return dataset.get_dataloader(batch_size=self.hparams.batch_size, shuffle=False)"
      ],
      "execution_count": 15,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "YbA8e334UM-N"
      },
      "source": [
        "## Preparação"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "XDAYv5_zURPc",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 221
        },
        "outputId": "f0f1ba47-a34c-42f8-f522-9d92d933fbfc"
      },
      "source": [
        "hparams = {\"name\": experiment_name, \"base_path\": base_path,\n",
        "           \"model_name\": model_name, \"seq_len\": sequence_length,\n",
        "           \"decode_mode\": decode_mode, \"k\": k,\n",
        "           \"lr\": learning_rate, \"batch_size\": batch_size, \"batch_accum\": accumulate_grad_batches,\n",
        "           \"max_epochs\": 3,\n",
        "           \"overfit_pct\": 0, \"debug\": 0,\n",
        "           \"decode_mode\": decode_mode}\n",
        "\n",
        "\n",
        "for key, parameter in hparams.items():\n",
        "    print(\"{}: {}\".format(key, parameter))"
      ],
      "execution_count": 16,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "name: no-tuning\n",
            "base_path: /content/t5-no-tuning\n",
            "model_name: t5-small\n",
            "seq_len: 100\n",
            "decode_mode: topk\n",
            "k: 10\n",
            "lr: 0.005\n",
            "batch_size: 10\n",
            "batch_accum: 1\n",
            "max_epochs: 3\n",
            "overfit_pct: 0\n",
            "debug: 0\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "fsxBZdkDzhNQ",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 68
        },
        "outputId": "dfdbe74c-1948-41aa-b072-3227a55ee2d0"
      },
      "source": [
        "# Instantiate model\n",
        "model = T5Corrector(Namespace(**hparams))\n",
        "\n",
        "# Folder/path management, for logs and checkpoints\n",
        "tensorboard_path = os.path.join(base_path, \"logs\")\n",
        "experiment_name = hparams[\"name\"]\n",
        "model_folder = os.path.join(tensorboard_path, experiment_name)\n",
        "os.makedirs(model_folder, exist_ok=True)\n",
        "ckpt_path = os.path.join(model_folder, \"-{epoch}\")\n",
        "\n",
        "# Callback initialization\n",
        "checkpoint_callback = ModelCheckpoint(prefix=experiment_name, \n",
        "                                      filepath=ckpt_path, \n",
        "                                      mode=\"max\")\n",
        "logger = TensorBoardLogger(tensorboard_path, experiment_name)\n",
        "\n",
        "# PL Trainer initialization\n",
        "trainer = Trainer(gpus=1,\n",
        "                  checkpoint_callback=checkpoint_callback, \n",
        "                  early_stop_callback=False,\n",
        "                  logger=logger,\n",
        "                  accumulate_grad_batches=hparams[\"batch_accum\"],\n",
        "                  max_epochs=hparams[\"max_epochs\"], \n",
        "                  fast_dev_run=bool(hparams[\"debug\"]), \n",
        "                  overfit_pct=hparams[\"overfit_pct\"],\n",
        "                  progress_bar_refresh_rate=1)"
      ],
      "execution_count": 17,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "GPU available: True, used: True\n",
            "TPU available: False, using: 0 TPU cores\n",
            "CUDA_VISIBLE_DEVICES: [0]\n"
          ],
          "name": "stderr"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gcCLkmIvmrx6",
        "colab_type": "text"
      },
      "source": [
        "## Tensorboard"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "T56srWEzmttv",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 838
        },
        "outputId": "1f116855-f2b8-4644-c13d-c0b4526e3c66"
      },
      "source": [
        "%load_ext tensorboard\n",
        "# %tensorboard --logdir \"/content/drive/My Drive/PF-Correcao/t5-no-tuning\"\n",
        "%tensorboard --logdir \"/content/t5-no-tuning\""
      ],
      "execution_count": 18,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "Reusing TensorBoard on port 6006 (pid 1514), started 0:52:43 ago. (Use '!kill 1514' to kill it.)"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "display_data",
          "data": {
            "application/javascript": [
              "\n",
              "        (async () => {\n",
              "            const url = await google.colab.kernel.proxyPort(6006, {\"cache\": true});\n",
              "            const iframe = document.createElement('iframe');\n",
              "            iframe.src = url;\n",
              "            iframe.setAttribute('width', '100%');\n",
              "            iframe.setAttribute('height', '800');\n",
              "            iframe.setAttribute('frameborder', 0);\n",
              "            document.body.appendChild(iframe);\n",
              "        })();\n",
              "    "
            ],
            "text/plain": [
              "<IPython.core.display.Javascript object>"
            ]
          },
          "metadata": {
            "tags": []
          }
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aslaho6VpMSL",
        "colab_type": "text"
      },
      "source": [
        "## Teste"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "efVSbZD1-chp",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 78,
          "referenced_widgets": [
            "1a1cda60f820424987bf483644073c5e",
            "9d9562be39fb485da6151955f74b22d4",
            "4b06efc977314efcb15415974ce0b129",
            "01ad8b879b3c49afab5ee8ac14d0362f",
            "065ed763dd0d477fa837e03362a835cc",
            "5d2df278d2734ee2a33a5b1e83672dd6",
            "9b5436fbb17e4130a25e9140569e6442",
            "f59f22994d244d12808e4544c4a0a38c"
          ]
        },
        "outputId": "d112a8f4-a7db-49bd-cab9-927149ed4348"
      },
      "source": [
        "trainer.test(model)"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "1a1cda60f820424987bf483644073c5e",
              "version_minor": 0,
              "version_major": 2
            },
            "text/plain": [
              "HBox(children=(FloatProgress(value=1.0, bar_style='info', description='Testing', layout=Layout(flex='2'), max=…"
            ]
          },
          "metadata": {
            "tags": []
          }
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6TYZSh-UvOJS",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        ""
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}