base_from_text.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Generic-ByT5-larger.ipynb",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "TPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MjWEQN6Z--FR"
      },
      "source": [
        "# ByT5 Generic\n",
        "\n",
        "Use TPU instance, create a GCS bucket with your train and validation sets as line-by-line text files"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_frSR4et-8Ua"
      },
      "source": [
        "## Dependencies"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "X9L9cVEN6QR5"
      },
      "source": [
        "%%capture\n",
        "! pip install t5 sentencepiece apache_beam --upgrade"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "1E3JfHHp6aK2",
        "outputId": "2736c7c8-4cc2-4a84-f47e-c2ba698485f3"
      },
      "source": [
        "! git clone https://github.com/google-research/byt5 byt5-repo\n",
        "! git clone https://github.com/google-research/multilingual-t5\n",
        "! git clone https://github.com/google-research/text-to-text-transfer-transformer"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Cloning into 'byt5-repo'...\n",
            "remote: Enumerating objects: 58, done.\u001b[K\n",
            "remote: Counting objects: 100% (58/58), done.\u001b[K\n",
            "remote: Compressing objects: 100% (35/35), done.\u001b[K\n",
            "remote: Total 58 (delta 26), reused 54 (delta 22), pack-reused 0\u001b[K\n",
            "Unpacking objects: 100% (58/58), done.\n",
            "Cloning into 'multilingual-t5'...\n",
            "remote: Enumerating objects: 246, done.\u001b[K\n",
            "remote: Counting objects: 100% (69/69), done.\u001b[K\n",
            "remote: Compressing objects: 100% (46/46), done.\u001b[K\n",
            "remote: Total 246 (delta 32), reused 32 (delta 23), pack-reused 177\u001b[K\n",
            "Receiving objects: 100% (246/246), 63.44 KiB | 2.64 MiB/s, done.\n",
            "Resolving deltas: 100% (144/144), done.\n",
            "Cloning into 'text-to-text-transfer-transformer'...\n",
            "remote: Enumerating objects: 3314, done.\u001b[K\n",
            "remote: Counting objects: 100% (259/259), done.\u001b[K\n",
            "remote: Compressing objects: 100% (113/113), done.\u001b[K\n",
            "remote: Total 3314 (delta 154), reused 200 (delta 145), pack-reused 3055\u001b[K\n",
            "Receiving objects: 100% (3314/3314), 5.18 MiB | 13.52 MiB/s, done.\n",
            "Resolving deltas: 100% (2413/2413), done.\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "P0Ehz6Fv_D3i"
      },
      "source": [
        "## Reorder"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "rqTvB7sY-tfo"
      },
      "source": [
        "! mv multilingual-t5/multilingual_t5 ./\n",
        "! mv byt5-repo/byt5 ./"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "nXj_BJjcAcvs"
      },
      "source": [
        "! mkdir models\n",
        "! cp byt5/gin/models/* models/\n",
        "! cp text-to-text-transfer-transformer/t5/models/gin/models/* models/"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "g-dx-vuXctFX"
      },
      "source": [
        "## Custom Text Dataset"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "x3saRKaWct17"
      },
      "source": [
        "# set up train and validation line by line files"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "FFbvtj1OgEw7"
      },
      "source": [
        "## already in byt5/tasks.py\n",
        "import t5\n",
        "import functools\n",
        "\n",
        "DEFAULT_BYTE_OUTPUT_FEATURES = {\n",
        "    \"inputs\": t5.data.Feature(vocabulary=t5.data.ByteVocabulary()),\n",
        "    \"targets\": t5.data.Feature(vocabulary=t5.data.ByteVocabulary())\n",
        "}\n",
        "MEAN_NOISE_SPAN_LENGTH = 20"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2jRpiY2RdPzx"
      },
      "source": [
        "### MANUAL EDITS ###\n",
        "## upload files to google storage\n",
        "## update paths\n",
        "## add this to byt5/tasks.py\n",
        "####################\n",
        "t5.data.TaskRegistry.add(\n",
        "      \"byt5_ex\",\n",
        "      t5.data.TextLineTask,\n",
        "      split_to_filepattern={\n",
        "            \"train\": \"gs://BUCKET/train_lines.txt\",\n",
        "            \"validation\": \"gs://BUCKET/validation_lines.txt\",\n",
        "        },\n",
        "      text_preprocessor=[\n",
        "        functools.partial(\n",
        "          t5.data.preprocessors.parse_tsv,\n",
        "          field_names=['text'],\n",
        "          field_delim='~', # check ASCII char doesn't appear in files, default is tab (\\t)\n",
        "        ),\n",
        "        functools.partial(\n",
        "              t5.data.preprocessors.rekey,\n",
        "              key_map={\n",
        "                  \"inputs\": None,\n",
        "                  \"targets\": \"text\"\n",
        "              }),\n",
        "      ],\n",
        "      token_preprocessor=functools.partial(\n",
        "          t5.data.preprocessors.span_corruption,\n",
        "          mean_noise_span_length=MEAN_NOISE_SPAN_LENGTH),\n",
        "      output_features=DEFAULT_BYTE_OUTPUT_FEATURES,\n",
        "      metric_fns=[])"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "al47b1-y_IwG"
      },
      "source": [
        "## TPU Training"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "3ZsOrQ-ARKJ-"
      },
      "source": [
        "# get tpu ip address and port\n",
        "import os\n",
        "os.environ['COLAB_TPU_ADDR']"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "TPNvQQ6F-BWA"
      },
      "source": [
        "! python -m t5.models.mesh_transformer_main \\\n",
        "  --gin_file=\"./models/byt5.base.gin\" \\\n",
        "  --gin_param=\"MIXTURE_NAME = 'byt5_ex'\" \\\n",
        "  --gin_param=\"utils.run.sequence_length = {'inputs': 256, 'targets': 189}\" \\\n",
        "  --gin_param=\"utils.run.batch_size = ('tokens_per_batch', 131072)\" \\\n",
        "  --gin_param=\"run.train_steps = 100000\" \\\n",
        "  --gin_param=\"utils.tpu_mesh_shape.tpu_topology = 'v3-8'\" \\\n",
        "  --gin_param=\"run.train_dataset_fn = @t5.models.mesh_transformer.mesh_train_dataset_fn\" \\\n",
        "  --gin_param=\"mesh_train_dataset_fn.mixture_or_task_name = 'byt5_ex'\" \\\n",
        "  --module_import=\"byt5.tasks\" \\\n",
        "  --tpu=\"grpc://TPU_LOCATION\" \\\n",
        "  --gin_param=\"utils.tpu_mesh_shape.model_parallelism = 1\" \\\n",
        "  --model_dir=\"gs://BUCKET/byt5_base_ex\" \\\n",
        "  --gcp_project=\"mapmeld-hrd\"\n",
        "\n",
        "  # will save every 5000 steps\n",
        "\n",
        "  #--t5_tfds_data_dir=\"${BUCKET}/t5-tfds\" \\\n",
        "    # --eval_mode=\"perplexity_eval\" \\\n",
        "#  --eval_gin_param=\"mesh_eval_dataset_fn.num_eval_examples = 10000\" \\\n",
        "\n",
        "  #--tpu_zone=\"${ZONE}\" \\\n",
        "#   --gin_param=\"utils.run.learning_rate_schedule=@learning_rate_schedules.rsqrt_no_ramp_down\" \\\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8GisvHROq55s"
      },
      "source": [
        "## Convert TF checkpoint"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_IIyyz9vq8Xk"
      },
      "source": [
        "%%capture\n",
        "! pip install transformers"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "7qP0bm34q73j"
      },
      "source": [
        "! cp ./drive/MyDrive/mlin/dvcorpus/dv-t5/checkpoint ./drive/MyDrive/mlin/dvcorpus/dv-t5/model.ckpt"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "jynr0ad3q_DO"
      },
      "source": [
        "! transformers-cli convert --model_type t5 \\\n",
        "  --tf_checkpoint ./drive/MyDrive/mlin/dvcorpus/dv-t5/ \\\n",
        "  --config ./drive/MyDrive/mlin/dvcorpus/dv-t5/config.json \\\n",
        "  --pytorch_dump_output ./drive/MyDrive/mlin/dvcorpus/dv-t5/"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "v3VfYzXrrGJV"
      },
      "source": [
        "from transformers import TFT5ForConditionalGeneration\n",
        "t_model = TFT5ForConditionalGeneration.from_pretrained('./drive/MyDrive/mlin/dvcorpus/dv-t5', from_pt=True)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "c6btcqfUrGb2"
      },
      "source": [
        "t_model.save_pretrained('./drive/MyDrive/mlin/dvcorpus/dv-t5/')"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}