g2p_for_OpenUtau_training_notebook.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "gpuType": "T4",
      "authorship_tag": "ABX9TyNAyupNoYdAAL8dPRZUSQFb",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/MLo7Ghinsan/DiffSinger_colab_notebook_MLo7/blob/main/g2p_for_OpenUtau_training_notebook.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "This is a notebook for training G2P files for OpenUtau.\n",
        "\n",
        "The original code can be found [here](https://github.com/stakira/OpenUtau/tree/master/py).\n",
        "\n",
        "This notebook is an edited copy of Lotte (how do i even tag you here help)"
      ],
      "metadata": {
        "id": "kJ2JAE6pYVrZ"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "#Setup"
      ],
      "metadata": {
        "id": "WN8XxMnEZZOA"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "cellView": "form",
        "id": "PH6eQUDIXpKC"
      },
      "outputs": [],
      "source": [
        "#@title # Mount Google Drive and Setup\n",
        "\n",
        "from google.colab import drive\n",
        "from IPython.display import clear_output\n",
        "\n",
        "!rm -rf /content/sample_data\n",
        "drive.mount(\"/content/drive\")\n",
        "!git clone https://github.com/stakira/OpenUtau.git\n",
        "#thank you lotte <3\n",
        "!pip install antlr4-python3-runtime==4.9.*\n",
        "!pip install hydra-core==1.3.2 omegaconf==2.3.0 -q #that warning popup is annoying\n",
        "!pip install torch==2.1.0 torchaudio==2.1.0\n",
        "!pip install editdistance==0.6.2\n",
        "!pip install tqdm==4.65.0\n",
        "!pip install onnx\n",
        "!pip install PyYAML\n",
        "#moving it cus like i dont think we need the whole ou folder\n",
        "!mv /content/OpenUtau/py/g2p /content\n",
        "!rm -rf OpenUtau\n",
        "%cd /content/g2p\n",
        "\n",
        "clear_output()"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Training"
      ],
      "metadata": {
        "id": "9j8n5pfgb1ml"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Input dictionary\n",
        "from IPython.display import clear_output\n",
        "%cd /content/g2p\n",
        "import os\n",
        "if not os.path.exists(\"/content/user_g2p_data\"):\n",
        "    os.makedirs(\"/content/user_g2p_data\")\n",
        "clear_output()\n",
        "import sys\n",
        "import re\n",
        "import torch\n",
        "import hydra\n",
        "import yaml\n",
        "from omegaconf import OmegaConf\n",
        "sys.path.append(os.path.abspath('.'))\n",
        "from dataset import SphinxDataset\n",
        "from trainer import G2pTrainer\n",
        "from model import GreedyG2p\n",
        "\n",
        "def train(trainer):\n",
        "    print('training...')\n",
        "    trainer.train()\n",
        "\n",
        "def export(trainer, model_path, onnx_path):\n",
        "    print('exporting model...')\n",
        "    trainer.model.load_state_dict(torch.load(model_path))\n",
        "    trainer.model.cpu()\n",
        "    greedy = GreedyG2p(trainer.model.max_len,\n",
        "                       trainer.model.encoder, trainer.model.decoder)\n",
        "    greedy.export(onnx_path)\n",
        "\n",
        "    print('testing...')\n",
        "    trainer.test('test_log.txt')\n",
        "\n",
        "#@markdown input dictionary (can be dict.txt or dsdict.yaml)\n",
        "dict_path = \"\" # @param {type:\"string\"}\n",
        "\n",
        "#@markdown input model save location\n",
        "save_path = \"\" # @param {type:\"string\"}\n",
        "if not save_path:\n",
        "    raise ValueError(\"Empty save_path, please specify a path to save your model\")\n",
        "else:\n",
        "    pass\n",
        "\n",
        "user_dict_path = save_path +\"/user_dictionary.txt\" #cleaned dict\n",
        "user_config = save_path + \"/config.yaml\" #generated config\n",
        "user_phonemes = save_path + \"/phones.txt\" #yea phoneme list yea poosay\n",
        "\n",
        "default_config = {\n",
        "    \"_target_\": \"model.G2p\",\n",
        "    \"max_len\": 48,\n",
        "    \"encoder\": {\n",
        "        \"_target_\": \"model.Encoder\",\n",
        "        \"graphemes\": [],\n",
        "        \"d_model\": 64,\n",
        "        \"d_hidden\": 128,\n",
        "        \"num_layers\": 2,\n",
        "        \"dropout\": 0.1\n",
        "    },\n",
        "    \"decoder\": {\n",
        "        \"_target_\": \"model.Decoder\",\n",
        "        \"phonemes\": [],\n",
        "        \"d_model\": 64,\n",
        "        \"d_hidden\": 128,\n",
        "        \"num_layers\": 2,\n",
        "        \"dropout\": 0.1\n",
        "    }\n",
        "}\n",
        "\n",
        "with open(user_config, \"w\") as cfg:\n",
        "    yaml.dump(default_config, cfg)\n",
        "\n",
        "if dict_path:\n",
        "    pass\n",
        "else:\n",
        "    raise ValueError(\"Please input path to your dictionary\")\n",
        "\n",
        "if dict_path.endswith(\".txt\"):\n",
        "    print(\"using txt format\")\n",
        "    with open(dict_path, \"r\", encoding = \"utf-8\") as dict:\n",
        "        content = dict.read()\n",
        "        if \"\\t\" in content:\n",
        "            no_tab_content = content.replace(\"\\t\", \"  \")\n",
        "            with open(user_dict_path, \"w\", encoding = \"utf-8\") as file:\n",
        "                file.write(no_tab_content)\n",
        "        else:\n",
        "            with open(user_dict_path, \"w\", encoding = \"utf-8\") as file:\n",
        "                file.write(content)\n",
        "\n",
        "elif dict_path.endswith(\".yaml\"):\n",
        "    print(\"using yaml format\")\n",
        "    #because they're troublesome >:( (not sure about true and false tho, but ill just include them cus why not)\n",
        "    bool2string = {\n",
        "        \"yes\": \"'yes'\",\n",
        "        \"no\": \"'no'\",\n",
        "        \"on\": \"'on'\",\n",
        "        \"off\": \"'off'\",\n",
        "        \"true\": \"'true'\",\n",
        "        \"false\": \"'false'\"\n",
        "    }\n",
        "\n",
        "    #ik its extra but oh well\n",
        "    with open(dict_path, \"r\", encoding = \"utf-8\") as file:\n",
        "        content = file.read()\n",
        "        for word, replacement in bool2string.items():\n",
        "            pattern = r\"(?<!\\')\\b{}\\b(?!\\')\".format(re.escape(word))\n",
        "            content = re.sub(pattern, replacement, content)\n",
        "    with open(dict_path, \"w\", encoding = \"utf-8\") as file:\n",
        "        file.write(content)\n",
        "\n",
        "    with open(dict_path, \"r\", encoding = \"utf-8\") as dict:\n",
        "        content = yaml.safe_load(dict)\n",
        "        entries = content.get(\"entries\", [])\n",
        "        with open(user_dict_path, \"w\", encoding = \"utf-8\") as file:\n",
        "            for entry in entries:\n",
        "                grapheme = entry[\"grapheme\"]\n",
        "                phonemes = \" \".join(entry[\"phonemes\"])\n",
        "                file.write(f\"{grapheme}  {phonemes}\\n\")\n",
        "else:\n",
        "    raise TypeError(\"format not supported\")\n",
        "\n",
        "graphemes = set()\n",
        "phonemes = set()\n",
        "\n",
        "with open(user_dict_path, \"r\", encoding = \"utf-8\") as file:\n",
        "    for line in file:\n",
        "        line = line.strip()\n",
        "        if line:\n",
        "            grapheme = line.split()[0]\n",
        "            graphemes.add(grapheme)\n",
        "            phoneme = line.split(\"  \")[1]\n",
        "            phonemes.update(set(phoneme.split()))\n",
        "\n",
        "#add the necessary stuff ig based on readme\n",
        "\n",
        "required_char = [\"<unk>\", \"<pad>\", \"<bos>\", \"<eos>\"]\n",
        "\n",
        "#edited this for when theres like special character in dict'''\\\\\n",
        "for char in required_char:\n",
        "    graphemes.discard(char)\n",
        "    phonemes.discard(char)\n",
        "\n",
        "graphemes = required_char + sorted(graphemes)\n",
        "phonemes = required_char + sorted(phonemes)\n",
        "\n",
        "vowel_types = {\"a\", \"i\", \"u\", \"e\", \"o\", \"N\", \"M\", \"NG\"}\n",
        "with open(user_phonemes, \"w\") as f:\n",
        "    for phoneme in phonemes:\n",
        "        if phoneme in required_char:\n",
        "            continue\n",
        "        if phoneme in vowel_types:\n",
        "            f.write(f\"{phoneme}\\tvowel\\n\")\n",
        "        else:\n",
        "            f.write(f\"{phoneme}\\t-\\n\")\n",
        "\n",
        "with open(user_config, \"r\", encoding = \"utf-8\") as cfg:\n",
        "    training_config = yaml.safe_load(cfg)\n",
        "training_config[\"encoder\"][\"graphemes\"] = graphemes\n",
        "training_config[\"decoder\"][\"phonemes\"] = phonemes\n",
        "with open(user_config, \"w\", encoding = \"utf-8\") as cfg:\n",
        "    yaml.dump(training_config, cfg, allow_unicode = True)\n",
        "\n",
        "cfg = user_config\n",
        "cfg = OmegaConf.load(cfg)\n",
        "\n",
        "dataset = user_dict_path\n",
        "dataset = SphinxDataset(dataset, cfg,\n",
        "                        comment_prefix=';;;',\n",
        "                        # \"RECORDS(1)\" -> \"RECORDS\"\n",
        "                        remove_word_digits=True,\n",
        "                        # \"R EH1 K ER0 D Z\" -> \"R EH K ER D Z\"\n",
        "                        remove_phoneme_digits=True)\n",
        "\n",
        "# @markdown you may need to adjust the batch size and epochs. <br> If there's too much loss, you can try decreasing the batch size. <br> Between 50 and 150 epochs is generally recommended for training, although you can play around with this a bit. <br> You can continue training from the latest checkpoint at a later moment if you so desire.\n",
        "loss_device = \"cuda\" # @param [\"cpu\", \"cuda\"]\n",
        "batch_size = 150 # @param {type:\"slider\", min:1, max:300, step:1}\n",
        "epochs = 120 # @param {type:\"slider\", min:1, max:300, step:1}\n",
        "\n",
        "\n",
        "#@markdown the model that get saves to save_path will be the g2p-best.ptsd, you can check in g2p folder under /content for other model if applicable"
      ],
      "metadata": {
        "cellView": "form",
        "id": "2Ht_ufjfb4T9"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Start training\n",
        "\n",
        "#@markdown This option to test the finished model for its word error and phoneme error rate.... very slow though, especially with big model\n",
        "test_for_error = False # @param {type:\"boolean\"}\n",
        "\n",
        "trainer = G2pTrainer(\n",
        "    device=torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\"),\n",
        "    loss_device=torch.device(loss_device),\n",
        "    model=hydra.utils.instantiate(cfg),\n",
        "    dataset=dataset,\n",
        "    batch_size=batch_size,\n",
        "    epochs=epochs)\n",
        "\n",
        "train(trainer)\n",
        "\n",
        "!cp /content/g2p/g2p-best.ptsd $save_path\n",
        "\n",
        "if test_for_error:\n",
        "    trainer.test(test_log = save_path +\"/test_log.txt\")\n",
        "else:\n",
        "    pass"
      ],
      "metadata": {
        "cellView": "form",
        "id": "8ffFaHTeUwNi"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "#Export ONNX"
      ],
      "metadata": {
        "id": "vr0st7XYV1MK"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Convert best checkpoint (not here yet lmao if you wanna use it then edit the code in this cell to your path(s), ill come back later to make it look a lil nicer)\n",
        "model_path = \"\"\n",
        "output_path = \"\" #+filename ill add that later cus im not so sure about the code above yet\n",
        "export(trainer, model_path, output_path)\n"
      ],
      "metadata": {
        "cellView": "form",
        "id": "J84E3kR3VbQs"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}