Punctuation_and_Capitalization_ipynb.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3gQAKaUssmdO"
      },
      "outputs": [],
      "source": [
        "BRANCH = 'r1.11.0'"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "o_0K1lsW1dj9"
      },
      "outputs": [],
      "source": [
        "\"\"\"\n",
        "このノートブックはローカル（すべての依存ファイルとGPUがある場合）でも、Google Colab上でも実行可能です。\n",
        "\n",
        "Colabのセットアップ手順は以下の通りです。\n",
        "1. 新しいPython 3ノートブックを開く。\n",
        "2. 2. GitHubからこのノートブックをインポートする（File -> Upload Notebook -> \"GITHUB\" tab -> copy/paste GitHub URL）。\n",
        "3. GPUを搭載したインスタンスに接続します（Runtime -> Change runtime type -> Select \"GPU\" for hardware accelerator）。\n",
        "4. このセルを実行して依存関係を設定する。\n",
        "\"\"\"\n",
        "# Google Colabを使用していて、ローカルで実行していない場合は、次のセルを実行してください。\n",
        "\n",
        "# install NeMo\n",
        "!python -m pip install git+https://github.com/NVIDIA/NeMo.git@$BRANCH#egg=nemo_toolkit[nlp]\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "uglDB-pVh__t",
        "pycharm": {
          "name": "#%%\n"
        }
      },
      "outputs": [],
      "source": [
        "# Colabを使っていない場合、以下のエラーを回避するためにjupyter notebookをアップグレードする必要があるかもしれません。\n",
        "# 'ImportError: IProgress not found. Please update jupyter and ipywidgets.'\n",
        "\n",
        "! pip install ipywidgets\n",
        "! jupyter nbextension enable --py widgetsnbextension\n",
        "\n",
        "# このセルを実行した後、カーネルを再起動してください。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dzqD2WDFOIN-"
      },
      "outputs": [],
      "source": [
        "from nemo.utils.exp_manager import exp_manager\n",
        "from nemo.collections import nlp as nemo_nlp\n",
        "\n",
        "import os\n",
        "import wget \n",
        "import torch\n",
        "import pytorch_lightning as pl\n",
        "from omegaconf import OmegaConf"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "daYw_Xll2ZR9"
      },
      "source": [
        "# タスクの説明\n",
        "自動音声認識（ASR）システムは通常、句読点や単語の大文字小文字がないテキストを生成します。\n",
        "このチュートリアルでは、ASR出力をより読みやすくし、名前付きエンティティ認識、機械翻訳、音声合成モデルのパフォーマンスを高めるために、文中の各単語の句読点と大文字を予測するモデルをNeMoに実装する方法について説明します。\n",
        "ここでは、事前学習済みのBERTモデルを使用して、このタスクのモデルを学習する方法を説明します。\n",
        "学習データセットの各単語について、次のことを予測する。\n",
        "\n",
        "- 単語の後に続く句読点と、その単語が大文字かどうかを予測する。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZnuziSwJ1yEB"
      },
      "source": [
        "# データセット\n",
        "このモデルは、以下のフォーマットに従っていれば、どのようなデータセットでも動作させることができます。\n",
        "学習・評価データは、text.txtとlabels.txt*の2つのファイルに分割される。\n",
        "**text.txt**ファイルの各行は、単語をスペースで区切ったテキスト列である。[WORD] [SPACE] [WORD] [SPACE] [WORD]といった具合である。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TXFORGBv2Jqu"
      },
      "source": [
        "\n",
        "\n",
        "```\n",
        "when is the next flight to new york\n",
        "the next flight is ...\n",
        "...\n",
        "```\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Y7ci55rM2QH8"
      },
      "source": [
        "**labels.txt**ファイルには、text.txtの各単語に対応するラベルが、スペースで区切られて格納されています。labels.txtファイルの各ラベルは、2つの記号で構成されています。\n",
        "\n",
        "- ラベルの最初の記号は、単語の後にどのような句読点を付けるかを示します（Oは句読点が不要であることを意味します）。\n",
        "- 2番目の記号は、単語を大文字にするかどうかを決めます（ここで、Uは単語を大文字にすること、Oは大文字にする必要がないことを示します）。\n",
        "\n",
        "このチュートリアルでは、カンマ、ピリオド、疑問符のみを取り上げ、その他の句読点は削除しています。もっと多くの句読点を使いたい場合は、データセットを更新して必要なラベルを追加してください。\n",
        "\n",
        "**labels.txt**の各行は次のような形式であるべきである。\n",
        "[LABEL] [SPACE] [LABEL] [SPACE] [LABEL] (labels.txtの場合). \n",
        "例えば、上記のtext.txtファイルのラベルは次のようになります。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-94C1-864EW1"
      },
      "source": [
        "\n",
        "\n",
        "```\n",
        "OU OO OO OO OO OO OU ?U\n",
        "OU OO OO OO ...\n",
        "...\n",
        "```\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VsEmwIPO4L4V"
      },
      "source": [
        "このチュートリアルで使用したこのタスクのすべての可能なラベルの完全なリストは次のとおりです。 `OO, ,O, .O, ?O, OU, ,U, .U, ?U.`"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SL58EWkd2ZVb"
      },
      "source": [
        "## データをダウンロードし、前処理を行う¶。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "THi6s1Qx2G1k"
      },
      "source": [
        "このノートでは、[Tatoeba collection of sentences](https://tatoeba.org/eng) から英語の例文のサブセットを使用します。このスクリプトはTatoebaデータをダウンロードして前処理をします[NeMo/examples/nlp/token_classification/get_tatoeba_data.py](https://github.com/NVIDIA/NeMo/blob/stable/examples/nlp/token_classification/data/get_tatoeba_data.py)。なお、このモデルでさらに実験を行う場合は、NUM_SAMPLES=-1に設定し、モデルの性能を向上させるために他のデータセットも含めることを検討してください。\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "n8HZrDmr12_-"
      },
      "outputs": [],
      "source": [
        "DATA_DIR = 'PATH_TO_A_DIRECTORY_WHERE_DATA_FROM_THIS_TUTORIAL_IS_STORED'\n",
        "WORK_DIR = 'PATH_TO_A_DIRECTORY_WHERE_SCRIPTS_FOR_THIS_TUTORIAL_ARE_SAVED'\n",
        "MODEL_CONFIG = \"punctuation_capitalization_config.yaml\"\n",
        "\n",
        "# model parameters\n",
        "TOKENS_IN_BATCH = 1024\n",
        "MAX_SEQ_LENGTH = 64\n",
        "LEARNING_RATE = 0.00002\n",
        "NUM_SAMPLES = 10000"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "UOeeeCGqI-9c"
      },
      "outputs": [],
      "source": [
        "## Tatoebaデータのダウンロードと前処理を行うget_tatoeba_data.pyスクリプトをダウンロードする。\n",
        "os.makedirs(WORK_DIR, exist_ok=True)\n",
        "if not os.path.exists(WORK_DIR + '/get_tatoeba_data.py'):\n",
        "    print('Downloading get_tatoeba_data.py...')\n",
        "    wget.download(f'https://raw.githubusercontent.com/NVIDIA/NeMo/{BRANCH}/examples/nlp/token_classification/data/get_tatoeba_data.py', WORK_DIR)\n",
        "else:\n",
        "    print ('get_tatoeba_data.py is already exists')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "H0ulD2TL13DR"
      },
      "outputs": [],
      "source": [
        "# データをダウンロードし、前処理を行う\n",
        "# --clean_dir フラグは、生の Tataoeba データを削除する。\n",
        "! python $WORK_DIR/get_tatoeba_data.py --data_dir $DATA_DIR --num_sample $NUM_SAMPLES --clean_dir"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pctMm2vsPlPT"
      },
      "source": [
        "上記のセルを実行すると、dataフォルダに学習に必要な以下の4つのファイルが格納されます（`--cleanan_dir`を使用しない場合は、Tatoebaの生データが存在する可能性があります）。\n",
        "- labels_dev.txt\n",
        "- labels_train.txt\n",
        "- text_dev.txt\n",
        "- text_train.txt\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "KKwxXXezPvXF"
      },
      "outputs": [],
      "source": [
        "! ls -l $DATA_DIR"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6UDPgadLN6SG"
      },
      "outputs": [],
      "source": [
        "# データを見てみよう\n",
        "print('Text:')\n",
        "! head -n 5 $DATA_DIR/text_train.txt\n",
        "\n",
        "print('\\nLabels:')\n",
        "! head -n 5 $DATA_DIR/labels_train.txt"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dwpZ8dpTsmdX"
      },
      "source": [
        "このように、`get_tatoeba_data.py`スクリプトは、Tatoebaをダウンロードするだけでなく、ラベルの作成も行っています。もし、自分のデータを前処理したい場合は、[examples/nlp/token_classification/data/prepare_data_for_punctuation_capitalization.py](https://github.com/NVIDIA/NeMo/blob/main/examples/nlp/token_classification/data/prepare_data_for_punctuation_capitalization.py)スクリプトを使用してください。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8eEx2ny5smdX"
      },
      "outputs": [],
      "source": [
        "cwd = os.getcwd()\n",
        "NEMO_ROOT = \"~/NeMo\"\n",
        "!python $NEMO_ROOT/examples/nlp/token_classification/data/prepare_data_for_punctuation_capitalization.py \\\n",
        "    --source_file $DATA_DIR/text_train.txt \\\n",
        "    --output_dir $DATA_DIR/my_train_preprocessed"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4WD5lx64smdX"
      },
      "outputs": [],
      "source": [
        "!ls $DATA_DIR/my_train_preprocessed -l"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Om-Iy8VvsmdX"
      },
      "source": [
        "# tarred dataset\n",
        "\n",
        "データセットが大きすぎてメモリに保存できない場合、tarred datasetを使うことができる。tarデータセットとは、モデルに渡すことのできるバッチを含むtarファイルの集合である。\n",
        "\n",
        "すべてのtarファイルには同じ数のバッチが含まれるので、もしデータセット内のバッチの数がパラメータ `--num_batches_per_tar_file` 値で均等に割り切れない場合、最大で `--num_batches_per_tar_file - 1` 個のバッチが失われる可能性があります。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "YsW5hXcOsmdX"
      },
      "outputs": [],
      "source": [
        "# テキストファイル、ラベルファイルの行数\n",
        "!wc -l $DATA_DIR/text_train.txt\n",
        "!wc -l $DATA_DIR/labels_train.txt"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ro9xQkZZsmdY"
      },
      "outputs": [],
      "source": [
        "NEMO_ROOT = \"~/NeMo\"\n",
        "!python $NEMO_ROOT/examples/nlp/token_classification/data/create_punctuation_capitalization_tarred_dataset.py \\\n",
        "    --text $DATA_DIR/text_train.txt \\\n",
        "    --labels $DATA_DIR/labels_train.txt \\\n",
        "    --output_dir $DATA_DIR/train_tarred \\\n",
        "    --num_batches_per_tarfile 5 \\\n",
        "    --tokens_in_batch 1024 \\\n",
        "    --lines_per_dataset_fragment 4000 \\\n",
        "    --tokenizer_name bert-base-uncased \\\n",
        "    --n_jobs 2"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ulQrR2H6smdY"
      },
      "outputs": [],
      "source": [
        "!ls $DATA_DIR/train_tarred -l"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XRGvzaQ8smdY"
      },
      "outputs": [],
      "source": [
        "!ls $DATA_DIR/train_tarred/*.tar | wc -l  # tarファイル数"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tJ5QF-OCsmdY"
      },
      "outputs": [],
      "source": [
        "!ls $DATA_DIR/train_tarred/ | grep -v '.tar'  # すべて非tarファイル"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OY-J6TqjsmdY"
      },
      "source": [
        "もし、圧縮されたデータセットを使用したい場合は、以下の設定が必要である。\n",
        "- 設定パラメータ `model.train_ds.tar_metadata_file` にメタデータの JSON ファイルを渡す。\n",
        "- model.train_ds.use_tarred_dataset=true` に設定してください。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "daludzzL2Jba"
      },
      "source": [
        "# モデル設定"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_whKCxfTMo6Y"
      },
      "source": [
        "Punctuation and Capitalization Modelでは、事前に学習した[BERT](https://arxiv.org/pdf/1810.04805.pdf)モデルの上に、2つのトークンレベルの分類器を共同で学習させています。\n",
        "- 句読点を予測する分類器と、大文字を予測する分類器です。\n",
        "- もう1つは大文字を予測する分類器です。\n",
        "\n",
        "このモデルは、複数の重要なセクションを宣言した設定ファイルで定義されています。それらは以下の通りです。\n",
        "- **モデル**。言語モデル、トークン分類器、オプティマイザ、スケジューラ、データセット、その他関連情報など、モデルに関連するすべての引数。\n",
        "\n",
        "- **trainer**: PyTorch Lightningに渡す任意の引数\n",
        "\n",
        "完全な設定の説明は[docs](https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/main/nlp/punctuation_and_capitalization.html#training-punctuation-and-capitalization-model)を参照してください。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "T1gA8PsJ13MJ"
      },
      "outputs": [],
      "source": [
        "# モデルの設定ファイルをダウンロードする \n",
        "config_dir = WORK_DIR + '/configs/'\n",
        "os.makedirs(config_dir, exist_ok=True)\n",
        "if not os.path.exists(config_dir + MODEL_CONFIG):\n",
        "    print('Downloading config file...')\n",
        "    wget.download(f'https://raw.githubusercontent.com/NVIDIA/NeMo/{BRANCH}/examples/nlp/token_classification/conf/' + MODEL_CONFIG, config_dir)\n",
        "else:\n",
        "    print ('config file is already exists')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "mX3KmWMvSUQw"
      },
      "outputs": [],
      "source": [
        "# この行は、モデルの全設定を表示します。\n",
        "config_path = f'{WORK_DIR}/configs/{MODEL_CONFIG}'\n",
        "print(config_path)\n",
        "config = OmegaConf.load(config_path)\n",
        "print(OmegaConf.to_yaml(config))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZCgWzNBkaQLZ"
      },
      "source": [
        "# コンフィグ内のデータのセットアップ\n",
        "\n",
        "設定ファイルには、`common_dataset_parameters`、`train_ds`、`validation_ds` という辞書が含まれています。これらは、対応するconfigのDatasetとDataLoaderを設定するために使用される設定ファイルである。\n",
        "\n",
        "パラメータ `train_ds.ds_item` と `validation_ds.ds_item` には、train と dev のデータセットが格納されたディレクトリを指定する。\n",
        "\n",
        "複数のデータセットで評価したい場合は、評価用ファイルが格納されているディレクトリを以下のように指定する。\n",
        "\n",
        "`model.validation_ds.ds_item=[PATH_TO_DEV1,PATH_TO_DEV2]` （パスと角括弧の間にスペースがないことに注意）。\n",
        "\n",
        "また、`model.train_ds.ds_item`を含むいくつかの設定項目には、値の代わりに`???`\n",
        "\n",
        "それでは、configにデータディレクトリのパスを追加してみましょう。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "LQHCJN-ZaoLp"
      },
      "outputs": [],
      "source": [
        "# このチュートリアルでは、train と dev のデータは同じフォルダーにあります。\n",
        "config.model.train_ds.ds_item = DATA_DIR\n",
        "config.model.validation_ds.ds_item=DATA_DIR\n",
        "del config.model.test_ds  # テストデータはなく、trainとdevのみです。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nB96-3sTc3yk"
      },
      "source": [
        "# PyTorch Lightningトレーナーの構築\n",
        "\n",
        "NeMoのモデルは主にPyTorch Lightningのモジュールなので、PyTorch Lightningのエコシステムと完全に互換性があります。\n",
        "\n",
        "まずはTrainerオブジェクトをインスタンス化しましょう。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1tG4FzZ4Ui60"
      },
      "outputs": [],
      "source": [
        "print(\"Trainer config - \\n\")\n",
        "print(OmegaConf.to_yaml(config.trainer))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "knF6QeQQdMrH"
      },
      "outputs": [],
      "source": [
        "# いくつかのトレーナー設定を変更する\n",
        "# GPU が利用可能かどうかを確認し、それを使用する\n",
        "accelerator = 'gpu' if torch.cuda.is_available() else 'cpu'\n",
        "config.trainer.devices = 1\n",
        "config.trainer.accelerator = accelerator\n",
        "config.trainer.precision = 16 if torch.cuda.is_available() else 32\n",
        "\n",
        "# 混合精度トレーニングの場合は、precision=16、amp_level=01としてください。\n",
        "\n",
        "# 最大エポック数を1に減らし、短時間での学習が可能\n",
        "config.trainer.max_epochs = 1\n",
        "\n",
        "# 分散学習フラグの削除\n",
        "config.trainer.strategy = None\n",
        "\n",
        "trainer = pl.Trainer(**config.trainer)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8IlEMdVxdr6p"
      },
      "source": [
        "# NeMoの実験をセットアップする¶。\n",
        "\n",
        "NeMoには実験マネージャがあり、ロギングやチェックポイントを処理してくれますので、それを使ってみましょう。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8uztqGAmdrYt"
      },
      "outputs": [],
      "source": [
        "exp_dir = exp_manager(trainer, config.get(\"exp_manager\", None))\n",
        "\n",
        "# exp_dir は、簡単にアクセスできるように、現在の実験へのパスを提供します。\n",
        "exp_dir = str(exp_dir)\n",
        "exp_dir"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6FI_nQsJo_11"
      },
      "source": [
        "# モデルトレーニング"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8tjLhUvL_o7_"
      },
      "source": [
        "モデルを初期化する前に、モデル設定のいくつかを変更したい場合があります。例えば、事前に学習したBERTモデルを変更したい場合があります。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Xeuc2i7Y_nP5"
      },
      "outputs": [],
      "source": [
        "# BERT-like モデルの全対応機種一覧\n",
        "print(nemo_nlp.modules.get_pretrained_lm_models_list())\n",
        "\n",
        "PRETRAINED_BERT_MODEL = \"bert-base-uncased\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "RK2xglXyAUOO"
      },
      "outputs": [],
      "source": [
        "# 指定された上記のモデルパラメータをコンフィグに追加する\n",
        "config.model.language_model.pretrained_model_name = PRETRAINED_BERT_MODEL\n",
        "config.model.train_ds.tokens_in_batch = TOKENS_IN_BATCH\n",
        "config.model.validation_ds.tokens_in_batch = TOKENS_IN_BATCH\n",
        "config.model.optim.lr = LEARNING_RATE\n",
        "config.model.train_ds.num_samples = NUM_SAMPLES\n",
        "config.model.validation_ds.num_samples = NUM_SAMPLES\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gYKcziSsiAAd"
      },
      "source": [
        "さて、これでモデルを初期化する準備が整いました。モデル初期化呼び出しの間、データセットとデータローダーは、訓練と評価のために準備されます。\n",
        "また、事前に学習された BERT モデルがダウンロードされますが、選択した BERT モデルのサイズによっては、最大で数分かかることがあります。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Yk2hJssviAAe",
        "pycharm": {
          "name": "#%%\n"
        },
        "scrolled": false
      },
      "outputs": [],
      "source": [
        "# モデルの初期化\n",
        "# この段階で、学習と評価のためのデータセットとデータローダーが準備される\n",
        "model = nemo_nlp.models.PunctuationCapitalizationModel(cfg=config.model, trainer=trainer)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kQ592Tx4pzyB"
      },
      "source": [
        "## トレーニングの進捗をモニタリングする\n",
        "オプションで、Tensorboardのビジュアライゼーションを作成し、トレーニングの進捗をモニタリングすることができます。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "mTJr16_pp0aS"
      },
      "outputs": [],
      "source": [
        "try:\n",
        "  from google import colab\n",
        "  COLAB_ENV = True\n",
        "except (ImportError, ModuleNotFoundError):\n",
        "  COLAB_ENV = False\n",
        "\n",
        "# TensorBoardノートブック拡張を読み込む\n",
        "if COLAB_ENV:\n",
        "  %load_ext tensorboard\n",
        "  %tensorboard --logdir {exp_dir}\n",
        "else:\n",
        "  print(\"To use tensorboard, please use this notebook in a Google Colab environment.\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hUvnSpyjp0Dh"
      },
      "outputs": [],
      "source": [
        "# トレーニング開始\n",
        "trainer.fit(model)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "m2PWpVGtsmdb"
      },
      "source": [
        "# tarred データセットによる学習"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "trl_BVtzsmdb"
      },
      "outputs": [],
      "source": [
        "config = OmegaConf.load(config_path)\n",
        "config.model.train_ds.ds_item = f'{DATA_DIR}/train_tarred'\n",
        "config.model.train_ds.use_tarred_dataset = True\n",
        "# `use_tarred_dataset=true` の場合は、メタデータファイル名のみが必要である。\n",
        "config.model.train_ds.tar_metadata_file = 'metadata.punctuation_capitalization.tokens1024.max_seq_length512.bert-base-uncased.json'\n",
        "config.model.validation_ds.ds_item = DATA_DIR\n",
        "del config.model.test_ds  # テストデータはなく、trainとdevのみです。\n",
        "\n",
        "# トレーナー\n",
        "accelerator = 'gpu' if torch.cuda.is_available() else 'cpu'\n",
        "config.trainer.devices = 1\n",
        "config.trainer.accelerator = accelerator\n",
        "config.trainer.precision = 16 if torch.cuda.is_available() else 32\n",
        "config.trainer.max_epochs = 1\n",
        "config.trainer.strategy = None\n",
        "\n",
        "# エクスプレス・マネージャー\n",
        "config.exp_manager.explicit_log_dir = 'tarred_experiment'\n",
        "\n",
        "config.model.language_model.pretrained_model_name = PRETRAINED_BERT_MODEL\n",
        "config.model.validation_ds.tokens_in_batch = TOKENS_IN_BATCH\n",
        "config.model.optim.lr = LEARNING_RATE\n",
        "config.model.validation_ds.num_samples = NUM_SAMPLES"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dZoOE5HAsmdb"
      },
      "outputs": [],
      "source": [
        "trainer = pl.Trainer(**config.trainer)\n",
        "exp_dir = exp_manager(trainer, config.get(\"exp_manager\", None))\n",
        "model = nemo_nlp.models.PunctuationCapitalizationModel(cfg=config.model, trainer=trainer)\n",
        "trainer.fit(model)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VPdzJVAgSFaJ"
      },
      "source": [
        "# 事前学習済みモデルを使った推論\n",
        "\n",
        "モデルの性能を見るために、いくつかの例で推論を実行してみましょう。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DQhsamclRtxJ"
      },
      "outputs": [],
      "source": [
        "print(f\"Available_models: {nemo_nlp.models.PunctuationCapitalizationModel.get_available_model_names()}\\n\")\n",
        "\n",
        "pretrained_model = nemo_nlp.models.PunctuationCapitalizationModel.from_pretrained(\"punctuation_en_distilbert\")\n",
        "# 推論に必要な行列のリストを定義する\n",
        "queries = [\n",
        "        'we bought four shirts and one mug from the nvidia gear store in santa clara',\n",
        "        'what can i do for you today',\n",
        "        'how are you',\n",
        "        'how is the weather in',\n",
        "    ]\n",
        "inference_results = pretrained_model.add_punctuation_capitalization(queries)\n",
        "print()\n",
        "\n",
        "for query, result in zip(queries, inference_results):\n",
        "    print(f'Query   : {query}')\n",
        "    print(f'Combined: {result.strip()}\\n')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "V15QVS4hsmdb"
      },
      "source": [
        "大量のテキストに対する推論は、スクリプト [examples/nlp/token_classification/punctuate_capitalize_infer.py](https://github.com/NVIDIA/NeMo/blob/stable/examples/nlp/token_classification/punctuate_capitalize_infer.py) によって行うことができます。\n",
        "\n",
        "```\n",
        "python punctuate_capitalize_infer.py \\\n",
        "    --input_manifest <PATH/TO/INPUT/MANIFEST> \\\n",
        "    --output_manifest <PATH/TO/OUTPUT/MANIFEST> \\\n",
        "    --pretrained_name punctuation_en_bert \\\n",
        "    --max_seq_length 64 \\\n",
        "    --margin 16 \\\n",
        "    --step 8\n",
        "```\n",
        "\n",
        "`<PATH/TO/INPUT/MANIFEST>` は NeMo [ASR manifest](https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/stable/asr/datasets.html) へのパスで、句読点や大文字を復元する必要があるテキストが格納されている場所です。マニフェストが `'pred_text'` キーを含んでいる場合、`'pred_text'` 要素が処理されます。そうでなければ、句読点と大文字は `'text'` 要素にリストアされます。\n",
        "\n",
        "`<PATH/TO/OUTPUT/MANIFEST>` は、結果を保存する NeMo ASR マニフェストへのパスです。句読点や大文字小文字が復元されたテキストは\n",
        "句読点と大文字が復元されたテキストは、入力マニフェストに `'pred_text'` キーが存在する場合、 `'pred_text'` 要素に保存されます。\n",
        "そうでなければ、結果は `'text'` 要素に保存されます。\n",
        "\n",
        "また、句読点や大文字小文字を復元するために、プレーンテキストとしてテキストを渡すこともできます。スクリプトのパラメータ `--input_text` と `--output_text` についてはヘルプを参照してください。\n",
        "[punctuate_capitalize_infer.py](https://github.com/NVIDIA/NeMo/blob/stable/examples/nlp/token_classification/punctuate_capitalize_infer.py) を参照してください。\n",
        "\n",
        "スクリプト `examples/nlp/token_classification/punctuate_capitalize_infer.py` は、任意の長さのテキストの句読点と大文字を復元することができます。長いシーケンスは、セグメントに分割されます。\n",
        "各セグメントは `--max_seq_length - 2` 個のトークンに分割されます。各セグメントは `[CLS]` と `[SEP]` トークンで始まり、`[CLS]` と `[SEP]` で終わります。\n",
        "各セグメントは、前のセグメントから `--step` トークンだけオフセットされています。例えば\n",
        "すべての文字がトークンで、 `--max_seq_length=5`, `--step=2` の場合、テキスト `\"hello\"` は次のように分割されます。\n",
        "セグメント `[['[CLS]', 'h', 'e', 'l', '[SEP]'], ['[CLS]', 'l', 'l', 'o', '[SEP]']]` に分割されることになります。\n",
        "\n",
        "セグメントが重複している場合は、複数のセグメントに存在するトークンの予測確率を掛け合わせてから、最適な候補を選択します。\n",
        "\n",
        "分割は、セグメントのエッジ付近でモデルのパフォーマンスを低下させる。パラメータ `--margin` を使って、セグメントの端にある `--margin` トークンに対して予測された確率を破棄することができる。例えば、全ての文字がトークンで、 `--max_seq_length=5`, `--step=1`, `--margin=1` の場合、テキスト `\"hello\"` はセグメント `[['[CLS]', 'h', 'e', 'l', '[SEP]'], ['[CLS]', 'e', 'l', 'l', '[SEP]'], ['[CLS]', 'l', 'l', 'o', '[SEP]']]` に分割されることになります。実際の予測値を計算する前に、アスタリスクでマークされたトークンの確率は削除される。`[['[CLS]', 'h', 'e', 'l'*, '[SEP]'*], ['[CLS]'*, 'e'*, 'l', 'l'*, '[SEP]'*], ['[CLS]'*, 'l'*, 'l', 'o', '[SEP]']]`.\n",
        "\n",
        "`add_punctuation_capitalization` メソッドでは、パラメータ `max_seq_length`, `step`, `margin` を使用することができます。\n",
        "\n",
        "以下の例のテキストはIWSLT 2019テストデータセットからのものです。\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0rrZDZq9smdc"
      },
      "outputs": [],
      "source": [
        "inference_results = pretrained_model.add_punctuation_capitalization(\n",
        "    [\n",
        "        \"each of these songs represents a scene a movement in some cases a sonic revolution that completely altered the \"\n",
        "        \"course of popular music they're all also calling cards almost for those cities songs totally linked with their \"\n",
        "        \"city's identity and might be why you probably consider them to be music cities now the magical mythical thing \"\n",
        "        \"the thing we kind of all love about stories like these is that those cities weren't doing anything in particular \"\n",
        "        \"to make those moments happen there's no formula for capturing lightning in a bottle a formula didn't give us \"\n",
        "        \"grunge music or introduce tupock to dr dray and there's definitely no blueprint for how to open your record \"\n",
        "        \"business in a south memphis neighborhood that turns out is home to booker t jones william bell and albert king \"\n",
        "        \"so this is just something that happens and right when the stars perfectly align great music just happens and \"\n",
        "        \"in the meantime new york and nashville can churn out the hits that come through our radios define our \"\n",
        "        \"generations and soundtrack our weddings and our funerals and everything in between i don't know about you but \"\n",
        "        \"the very idea of that is just deadly boring to me there are musicians all around you making powerful important \"\n",
        "        \"music and thanks to the internet and it's limitless possibilities for creators to create music and fans to \"\n",
        "        \"discover that music those zyite guy songs don't have to be handed down to us from some conference room full of \"\n",
        "        \"songwriters in a corporate high rise but also and more importantly we can't decide that it's just something that \"\n",
        "        \"happens because music is about so much more than hits those big iconic moments that change everything it's more \"\n",
        "        \"than just entertainment for so many of us music is truly a way to navigate life a means of self expression sure \"\n",
        "        \"but it also helps us find ourrselfel worse and figure out who we are it connects us with other people as almost \"\n",
        "        \"nothing else can across language barriers across social and cultural and economic divides music makes us smarter \"\n",
        "        \"and healthier and happier music is necessary what if you lived in a city that believed that that said we're not \"\n",
        "        \"waiting for that hit song to define us we're a music city because music is necessary by seeing music as \"\n",
        "        \"necessary a city can build two things 1st an ecosystem to support the development of professional musicians \"\n",
        "        \"and music business and 2nd a receptive and engaged audience to sustain them and those are the two critical \"\n",
        "        \"elements of a music city a city whose leaders recognize the importance of music for our development as \"\n",
        "        \"individuals our connection as a community and our viability as a vibrant place to live see smart cities music \"\n",
        "        \"cities know that thriving night lifef a creative class culture is what attracts young talented people to cities \"\n",
        "        \"it's what brings that lightning and no we can't predict the next egg that will hatch but we can create a city \"\n",
        "        \"that acts like an incubator to do that 1st we got to know what we've got that means identifying and quantifying \"\n",
        "        \"our assets we need to know them backward and forward from who and what and where they are to what their impact \"\n",
        "        \"is on the economy let's count our recording studios and our record labels our historic landmarks and our \"\n",
        "        \"hardcore punk clubs we should count monthly free jazz nights and weekly folk jams music schools artist \"\n",
        "        \"development instrument shops every lay than every luther music museums open once a open year round and music \"\n",
        "        \"festivals open just one weekend year now ideally through this process we'll create an actual asset map \"\n",
        "        \"dropping a pin for each one allowing us to see exactly what we've got and where organic momentum is already \"\n",
        "        \"happening because it's not enough to paint in broad strokes here when it comes to specific support for music \"\n",
        "        \"locally and a broad understanding of a music brand nationally you've got to have the receipts next we'll need \"\n",
        "        \"to identify our challenges now it's important to knowe that for the most part this won't be just the opposite \"\n",
        "        \"of step one we won't gain a whole lot by simply thinking about what's missing from our map instead we need to \"\n",
        "        \"approach this more holistically there are lots of music venues on our map awesome but are they struggling do \"\n",
        "        \"we have a venue ladder which just means can an artist starting out at a coffee house open mike see a clear path \"\n",
        "        \"for how they'll grow from that 25 seat room to a hundred seat room and so on or are we expecting them to go from \"\n",
        "        \"a coffee house to a colossum maybe our challenges lie in city infrastructure public transportation affordable \"\n",
        "        \"housing maybe like in london where the number of music venues went from 400 in 2010 to 100 in 2015 we need to \"\n",
        "        \"think about protections against gentrification the mayor of london in december of last year actually added \"\n",
        "        \"something called the agent of change principle to the city's comprehensive plan the name says it all if a real \"\n",
        "        \"estate developer wants to build condos next to an existing music venue the developer is the agent of change they \"\n",
        "        \"have to take the necessary steps for noise mitigation next and this is a very big one we need leadership and we \"\n",
        "        \"need a strategy now we know there's a lot of magic in this mix a lot of right people right place right time and \"\n",
        "        \"that will never stop being an important element of the way music is made the way some of the best most enduring \"\n",
        "        \"music is made but there cannot be a leadership vacuum in 2018 surriving music cities don't often happen and \"\n",
        "        \"don't have to happen accidentally we need elected officials who recognize the power of music and elevate the \"\n",
        "        \"voices of creatives and they're ready to put a strategy in place in music cities from berlin to paris to bogata \"\n",
        "        \"music advisory councils ensure that musicians have a seat at the table they're volunteer councils and they work \"\n",
        "        \"directly with a designated advocate inside of city hall or even the chamber of commerce the strongest strategies \"\n",
        "        \"will build music community supports like this one inward while also exporting music outward they go hand in hand \"\n",
        "        \"when we look inward we create that place that musicians want to live and when we look outward we build \"\n",
        "        \"opportunities for them to advance their career while also driving attention back to our city and leveraging music \"\n",
        "        \"as a talent attraction tool and here's something else that will help with that we've got to figure out who we are \"\n",
        "        \"now when i say austin you probably think live music capital and why because in 1991 leadership in austin saw \"\n",
        "        \"something percolating with an existing asset and they chose to own it by recognizing that momentum naming it and \"\n",
        "        \"claiming it they inevitably caused more live music venues to open existing spaces to add live music to their \"\n",
        "        \"repertoire and they created a swell of civic buy in around the idea which meant that it wasn't just a slogan in \"\n",
        "        \"some tourism pamphlet was something that locals really started to believe and take pride in now generally \"\n",
        "        \"speaking what austin created is just an assets based narrative and when we think back to step one we know that \"\n",
        "        \"every city will not tick every box many cities won't have recording studios like memphis or a songwriter and \"\n",
        "        \"publishing scene like nashville and that's not a deal breaker we simply have to find the momentum happening in \"\n",
        "        \"our city what are our unique assets in comparison to no other place so if all of that sounds like something \"\n",
        "        \"you'd like to happen where you live here are three things you can do to move the needle 1st you can use your \"\n",
        "        \"feet your ears and your dollars show up be that receptive and engaged audience that is so necessary for a music \"\n",
        "        \"city to thrive pay a cover charge buy a record discover new music and please take your friends two you can use \"\n",
        "        \"your voice buy in to the assets based narrative talk about and celebrate what your city has and three you can \"\n",
        "        \"use your vote seek out leadership that doesn't just pay lip service to your city's music but recognizes its \"\n",
        "        \"power and is prepared to put a strategy in place to elevate it grow it and build collaboration no there really \"\n",
        "        \"is no telling what city could be defined by a certain scene or a certain song in the next decade but as much \"\n",
        "        \"as we absolutely cannot predict that what we absolutely can predict is what happens when we treat music as \"\n",
        "        \"necessary and we work to build a music city and that is a place where i want to live thank you\"\n",
        "    ],\n",
        "    max_seq_length=128,\n",
        "    step=8,\n",
        "    margin=16,\n",
        "    batch_size=32,\n",
        ")\n",
        "print(inference_results[0])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ref1qSonGNhP"
      },
      "source": [
        "## 学習スクリプト\n",
        "\n",
        "NeMoがローカルにインストールされている場合、[nlp/token_classification/punctuation_capitalization_train_evaluate.py](https://github.com/NVIDIA/NeMo/blob/main/examples/nlp/token_classification/punctuation_capitalization_train_evaluate.py)でモデルを学習することも可能です。\n",
        "\n",
        "学習スクリプトを実行するには、以下を使用します。\n",
        "\n",
        "`python punctuation_capitalization_train_evaluate.py model.train_ds.ds_item=PATH_TO_TRAIN_DATA_DIR` とします。\n",
        "\n",
        "NUM_SAMPLES=-1に設定し、モデルの性能を向上させるために他のデータセットを含めることを検討する。\n",
        "\n",
        "# あなたのデータでモデルを微調整する\n",
        "\n",
        "ゼロからモデルを学習する場合、データセットはモデルの初期化時に学習用に準備されます。事前に学習した Punctuation and Capitalization モデルを使用する場合、学習前に、学習データと評価データを設定する必要があります。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4X1BahRlkaNf"
      },
      "outputs": [],
      "source": [
        "# let's reload our pretrained model\n",
        "pretrained_model = nemo_nlp.models.PunctuationCapitalizationModel.from_pretrained('punctuation_en_distilbert')\n",
        "\n",
        "# trainとvalidationの設定 Pytorch DataLoaders\n",
        "pretrained_model.update_config_after_restoring_from_checkpoint(\n",
        "    train_ds={\n",
        "        'ds_item': DATA_DIR,\n",
        "        'text_file': 'text_train.txt',\n",
        "        'labels_file': 'labels_train.txt',\n",
        "        'tokens_in_batch': 1024,\n",
        "    },\n",
        "    validation_ds={\n",
        "        'ds_item': DATA_DIR,\n",
        "        'text_file': 'text_dev.txt',\n",
        "        'labels_file': 'labels_dev.txt',\n",
        "        'tokens_in_batch': 1024,\n",
        "    },\n",
        ")\n",
        "\n",
        "# このチュートリアルでは、fast_dev_runをTrueに設定し、実際のモデル学習のために1つのトレーニングバッチと1つの検証バッチを実行します。\n",
        "fast_dev_run = True\n",
        "trainer = pl.Trainer(devices=1, accelerator='gpu', fast_dev_run=fast_dev_run)\n",
        "pretrained_model.set_trainer(trainer)\n",
        "pretrained_model.setup_training_data()\n",
        "pretrained_model.setup_validation_data()\n",
        "trainer.fit(pretrained_model)"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [
        "daYw_Xll2ZR9"
      ],
      "name": "Punctuation_and_Capitalization.ipynb のコピー",
      "private_outputs": true,
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3.8.10 64-bit",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.8.10"
    },
    "pycharm": {
      "stem_cell": {
        "cell_type": "raw",
        "metadata": {
          "collapsed": false
        },
        "source": []
      }
    },
    "vscode": {
      "interpreter": {
        "hash": "916dbcbb3f70747c44a77c7bcd40155683ae19c65e1c03b4aa3499c5328201f1"
      }
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}