prepare_vbx_tfdata.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6ypmOr_03EHP"
      },
      "source": [
        "Here's what this notebook does, step-by-step:\n",
        "\n",
        "1. Install the `montreal-force-aligner` (MFA) library using conda.\n",
        "2. Download the Viet Bible Vox dataset from Huggingface.\n",
        "3. Generate a normalized transcript file for each .mp3 file in the dataset.\n",
        "4. Train an MFA model, then align speech and phonemes (creating a timestamp for each phoneme).\n",
        "5. Generate a TensorFlow dataset (tfrecord files) that is ready for training TTS models."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SwhMuEANS0v4"
      },
      "outputs": [],
      "source": [
        "%pip install -Uqq tensorflow numpy torch tqdm librosa numba regex"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# ./data is our working directory\n",
        "!mkdir -p data\n",
        "%cd data"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "euumE7pMw6j8",
        "outputId": "0c54a090-f8c5-4729-fb86-f3e47bf6a3c8"
      },
      "outputs": [],
      "source": [
        "#### INSTALL MFA  ####\n",
        "!wget https://repo.anaconda.com/miniconda/Miniconda3-py311_23.5.2-0-Linux-x86_64.sh -qO ./miniconda.sh\n",
        "# !wget https://repo.anaconda.com/miniconda/Miniconda3-latest-MacOSX-arm64.sh -qO $PWD/miniconda.sh # for Apple M1\n",
        "!bash ./miniconda.sh -b -p ./miniconda\n",
        "!source ./miniconda/bin/activate && conda create -n aligner -c conda-forge montreal-forced-aligner=2.2.15 -y --quiet"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1whkRTu5NBam"
      },
      "outputs": [],
      "source": [
        "!git lfs install # LFS is needed to download the data\n",
        "!git clone https://huggingface.co/datasets/ntt123/VietBibleVox\n",
        "!cd VietBibleVox; unzip -qq data.zip\n",
        "!cd VietBibleVox; ls *.mp3 | wc -l\n",
        "!cd VietBibleVox; ls *.txt | wc -l"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sPFR2tbJV1au"
      },
      "outputs": [],
      "source": [
        "import regex\n",
        "import unicodedata\n",
        "from pathlib import Path"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "QgaPWypIWDQF"
      },
      "outputs": [],
      "source": [
        "vietnamese_characters = [\n",
        "    'a', 'à', 'á', 'ả', 'ã', 'ạ',\n",
        "    'ă', 'ằ', 'ắ', 'ẳ', 'ẵ', 'ặ',\n",
        "    'â', 'ầ', 'ấ', 'ẩ', 'ẫ', 'ậ',\n",
        "    'e', 'è', 'é', 'ẻ', 'ẽ', 'ẹ',\n",
        "    'ê', 'ề', 'ế', 'ể', 'ễ', 'ệ',\n",
        "    'i', 'ì', 'í', 'ỉ', 'ĩ', 'ị',\n",
        "    'o', 'ò', 'ó', 'ỏ', 'õ', 'ọ',\n",
        "    'ô', 'ồ', 'ố', 'ổ', 'ỗ', 'ộ',\n",
        "    'ơ', 'ờ', 'ớ', 'ở', 'ỡ', 'ợ',\n",
        "    'u', 'ù', 'ú', 'ủ', 'ũ', 'ụ',\n",
        "    'ư', 'ừ', 'ứ', 'ử', 'ữ', 'ự',\n",
        "    'y', 'ỳ', 'ý', 'ỷ', 'ỹ', 'ỵ',\n",
        "    'b', 'c', 'd', 'đ', 'g', 'h', \n",
        "    'k', 'l', 'm', 'n', 'p', 'q', \n",
        "    'r', 's', 't', 'v', 'x'\n",
        "]\n",
        "alphabet = \"\".join(vietnamese_characters)\n",
        "space_re = regex.compile(r\"\\s+\")\n",
        "number_re = regex.compile(\"([0-9]+)\")\n",
        "digits = [\"không\", \"một\", \"hai\", \"ba\", \"bốn\", \"năm\", \"sáu\", \"bảy\", \"tám\", \"chín\"]\n",
        "num_re = regex.compile(r\"([0-9.,]*[0-9])\")\n",
        "keep_text_and_num_re = regex.compile(rf'[^\\s{alphabet}.,0-9]')\n",
        "keep_text_re = regex.compile(rf'[^\\s{alphabet}]')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jXeBft4Ctk9p"
      },
      "outputs": [],
      "source": [
        "def read_number(num: str) -> str:\n",
        "    if len(num) == 1:\n",
        "        return digits[int(num)]\n",
        "    elif len(num) == 2 and num.isdigit():\n",
        "        n = int(num)\n",
        "        end = digits[n % 10]\n",
        "        if n == 10:\n",
        "            return \"mười\"\n",
        "        if n % 10 == 5:\n",
        "            end = \"lăm\"\n",
        "        if n % 10 == 0:\n",
        "            return digits[n // 10] + \" mươi\"\n",
        "        elif n < 20:\n",
        "            return \"mười \" + end\n",
        "        else:\n",
        "            if n % 10 == 1:\n",
        "                end = \"mốt\"\n",
        "            return digits[n // 10] + \" mươi \" + end\n",
        "    elif len(num) == 3 and num.isdigit():\n",
        "        n = int(num)\n",
        "        if n % 100 == 0:\n",
        "            return digits[n // 100] + \" trăm\"\n",
        "        elif num[1] == \"0\":\n",
        "            return digits[n // 100] + \" trăm lẻ \" + digits[n%100]\n",
        "        else:\n",
        "            return digits[n // 100] + \" trăm \" + read_number(num[1:])\n",
        "    elif \",\" in num:\n",
        "        n1, n2 = num.split(\",\")\n",
        "        return read_number(n1) + \" phẩy \" + read_number(n2)\n",
        "    elif \".\" in num:\n",
        "        parts = num.split(\".\")\n",
        "        if len(parts) == 2:\n",
        "            if parts[1] == \"000\":\n",
        "                return read_number(parts[0]) + \" ngàn\"\n",
        "            elif parts[1].startswith(\"00\"):\n",
        "                end = digits[int(parts[1][2:])]\n",
        "                return read_number(parts[0]) + \" ngàn lẻ \" + end\n",
        "            else:\n",
        "                return read_number(parts[0]) + \" ngàn \" + read_number(parts[1])\n",
        "        elif len(parts) == 3:\n",
        "            if num == \"1.000.000\":\n",
        "                return \"một triệu\"\n",
        "            elif num == \"1.100.000\":\n",
        "                return \"một triệu một trăm ngàn\"\n",
        "            elif num == \"1.820.000\":\n",
        "                return \"một triệu tám trăm hai mươi ngàn\"\n",
        "            elif num == \"3.640.000\":\n",
        "                return \"ba triệu sáu trăm bốn mươi ngàn\"\n",
        "            else:\n",
        "                raise ValueError(f\"Cannot convert {num}\")\n",
        "    else:\n",
        "        raise ValueError(f\"Cannot convert {num}\")\n",
        "    return num"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5taIbrZ5WGX8"
      },
      "outputs": [],
      "source": [
        "def normalize_text(x):\n",
        "    x = unicodedata.normalize('NFKC', x)\n",
        "    x = x.lower()\n",
        "    x = num_re.sub(r\" \\1 \", x)\n",
        "    x = keep_text_and_num_re.sub(\" \", x)\n",
        "    words = x.split()\n",
        "    words = [ read_number(w) if num_re.fullmatch(w) else w for w in words ]\n",
        "    x = \" \".join(words)\n",
        "    x = keep_text_re.sub(\" \", x)\n",
        "    x = space_re.sub(\" \", x)\n",
        "    x = x.strip()\n",
        "    return x"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-QWNUkoVV6WF"
      },
      "outputs": [],
      "source": [
        "all_text = []\n",
        "for fp in sorted(Path(\"VietBibleVox\").glob(\"*.txt\")):\n",
        "    with open(fp, \"r\", encoding=\"utf-8\") as f:\n",
        "        text = f.read()\n",
        "        text = normalize_text(text)\n",
        "        all_text.append(text)\n",
        "    # override the text file\n",
        "    with open(fp, \"w\", encoding=\"utf-8\") as f:\n",
        "        f.write(text)\n",
        "all_words = sorted(set((\" \".join(all_text)).split()))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "JFZDqnMzcH5g"
      },
      "outputs": [],
      "source": [
        "with open(\"lexicon.txt\", \"w\") as f:\n",
        "    for w in all_words:\n",
        "        w = w.strip()\n",
        "        p = list(w)\n",
        "        p = \" \".join(p)\n",
        "        f.write(f\"{w}\\t{p}\\n\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "D8s22KjOwtew"
      },
      "outputs": [],
      "source": [
        "# replace `nproc` by `sysctl -n hw.physicalcpu` if you are using MacOS\n",
        "!source miniconda/bin/activate aligner; \\\n",
        "mfa train \\\n",
        "    --num_jobs `nproc` \\\n",
        "    --use_mp \\\n",
        "    --clean \\\n",
        "    --overwrite \\\n",
        "    --no_textgrid_cleanup \\\n",
        "    --single_speaker \\\n",
        "    --output_format json \\\n",
        "    --output_directory VietBibleVox \\\n",
        "    VietBibleVox ./lexicon.txt vbx_mfa"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "uy88PuP4TtoU"
      },
      "outputs": [],
      "source": [
        "#### PREPARE TF DATASET ####\n",
        "# Note: We are using a sampling rate of 16k,\n",
        "# even though the original data has a sampling rate of 48k."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "p58RNMt8CISF"
      },
      "outputs": [],
      "source": [
        "import json\n",
        "from pathlib import Path\n",
        "import numpy as np\n",
        "import torch\n",
        "import json\n",
        "import librosa\n",
        "import tensorflow as tf\n",
        "from tqdm.auto import tqdm\n",
        "import random"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "YbF8kuar0zuC"
      },
      "outputs": [],
      "source": [
        "mel_basis = {}\n",
        "hann_window = {}\n",
        "\n",
        "def spectrogram_torch(y, n_fft, sampling_rate, hop_size, win_size, center=False):\n",
        "    global hann_window\n",
        "    if torch.min(y) < -1.0:\n",
        "        print(\"min value is \", torch.min(y))\n",
        "    if torch.max(y) > 1.0:\n",
        "        print(\"max value is \", torch.max(y))\n",
        "\n",
        "    dtype_device = str(y.dtype) + \"_\" + str(y.device)\n",
        "    wnsize_dtype_device = str(win_size) + \"_\" + dtype_device\n",
        "    if wnsize_dtype_device not in hann_window:\n",
        "        hann_window[wnsize_dtype_device] = torch.hann_window(win_size).to(\n",
        "            dtype=y.dtype, device=y.device\n",
        "        )\n",
        "\n",
        "    y = torch.nn.functional.pad(\n",
        "        y.unsqueeze(1),\n",
        "        (int((n_fft - hop_size) / 2), int((n_fft - hop_size) / 2)),\n",
        "        mode=\"reflect\",\n",
        "    )\n",
        "    y = y.squeeze(1)\n",
        "\n",
        "    spec = torch.stft(\n",
        "        y,\n",
        "        n_fft,\n",
        "        hop_length=hop_size,\n",
        "        win_length=win_size,\n",
        "        window=hann_window[wnsize_dtype_device],\n",
        "        center=center,\n",
        "        pad_mode=\"reflect\",\n",
        "        normalized=False,\n",
        "        onesided=True,\n",
        "        return_complex=True\n",
        "    )\n",
        "    spec = torch.view_as_real(spec)\n",
        "    spec = torch.sqrt(spec.pow(2).sum(-1) + 1e-6)\n",
        "    spec = spec.squeeze(0)\n",
        "    return torch.swapaxes(spec, 0, 1)\n",
        "\n",
        "\n",
        "def tensor_to_bytes(t):\n",
        "    t = tf.constant(t)\n",
        "    t = tf.io.serialize_tensor(t)\n",
        "    return tf.train.Feature(bytes_list=tf.train.BytesList(value=[t.numpy()]))\n",
        "\n",
        "\n",
        "def write_tfdata(data, out_file):\n",
        "    with tf.io.TFRecordWriter(out_file) as file_writer:\n",
        "        for wav_file, seq, _ in data:\n",
        "            phone_seq = []\n",
        "            for phone, duration in seq:\n",
        "                phone_idx = phone_set.index(phone)\n",
        "                phone_seq.append((phone_idx, duration))\n",
        "            phone_seq = np.array(phone_seq, dtype=np.float32)\n",
        "\n",
        "            # load wav\n",
        "            wav, sr = librosa.load(wav_file, sr=config[\"data\"][\"sampling_rate\"], dtype=np.float32)\n",
        "            wav = torch.from_numpy(wav)\n",
        "            # compute spec\n",
        "            spec = spectrogram_torch(\n",
        "                wav[None],\n",
        "                n_fft=config[\"data\"][\"filter_length\"],\n",
        "                sampling_rate=config[\"data\"][\"sampling_rate\"],\n",
        "                hop_size=config[\"data\"][\"hop_length\"],\n",
        "                win_size=config[\"data\"][\"win_length\"],\n",
        "                center=False\n",
        "            )\n",
        "            features = {\n",
        "                \"phone_idx\": tensor_to_bytes(phone_seq[:, 0].astype(np.int32)),\n",
        "                \"phone_duration\": tensor_to_bytes(phone_seq[:, 1]),\n",
        "                \"wav\": tensor_to_bytes(wav.half().numpy()),\n",
        "                \"spec\": tensor_to_bytes(spec.half().numpy())\n",
        "            }\n",
        "            example = tf.train.Example(features=tf.train.Features(feature=features))\n",
        "            file_writer.write(example.SerializeToString())\n",
        "\n",
        "def write_split(split, data, num_chunks):\n",
        "    data = np.array(data, dtype=object)\n",
        "    chunks = list(np.array_split(data, num_chunks))\n",
        "    for i, chunk in enumerate(tqdm(chunks)):\n",
        "        write_tfdata(chunk, f\"tfdata/{split}/part_{i:03d}.tfrecords\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "UIAxgVx-DkRI"
      },
      "outputs": [],
      "source": [
        "!mkdir -p tfdata/{train,test}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "STURjAulDm5B"
      },
      "outputs": [],
      "source": [
        "with open(\"../config.json\", \"rb\") as f:\n",
        "    config = json.load(f)\n",
        "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
        "data_dir = Path(\"VietBibleVox\")\n",
        "json_files = sorted(data_dir.glob(\"*.json\"))\n",
        "dataset = []\n",
        "phone_set = []\n",
        "\n",
        "for file_path in json_files:\n",
        "    with open(file_path, \"rb\") as f:\n",
        "        data = json.load(f)\n",
        "    seq = []\n",
        "    word_index = 0\n",
        "    words =  data[\"tiers\"][\"words\"][\"entries\"]\n",
        "    for start, end, phone in data[\"tiers\"][\"phones\"][\"entries\"]:\n",
        "        if start > words[word_index][1] - 1e-5:\n",
        "            seq.append( (\"<SEP>\", 0) )\n",
        "            word_index += 1\n",
        "        duration = end * 1000 - start * 1000 # ms\n",
        "        phone_set.append(phone)\n",
        "        seq.append( (phone, duration) )\n",
        "    wav_file = file_path.with_suffix(\".mp3\")\n",
        "    dataset.append((wav_file, seq, data[\"end\"]))\n",
        "\n",
        "phone_set = [\"<SEP>\"] + sorted(set(phone_set))\n",
        "assert len(phone_set) <= 256\n",
        "with open(\"phone_set.json\", \"w\", encoding=\"utf-8\") as f:\n",
        "    json.dump(phone_set, f)\n",
        "\n",
        "assert phone_set.index(\"<SEP>\") == 0"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2d86zvZBD0r6"
      },
      "outputs": [],
      "source": [
        "random.Random(42).shuffle(dataset)\n",
        "L = len(dataset) - 256\n",
        "train_data = dataset[:L]\n",
        "test_data = dataset[L:]\n",
        "print(\"Train data size:\", len(train_data))\n",
        "print(\"Test data size:\", len(test_data))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "UnDTA1sKD9Ug"
      },
      "outputs": [],
      "source": [
        "write_split(\"test\", test_data, 1)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "JdG7lqdEEQNz"
      },
      "outputs": [],
      "source": [
        "write_split(\"train\", train_data, 256)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "M7ZIyB_ANeSa"
      },
      "outputs": [],
      "source": []
    }
  ],
  "metadata": {
    "accelerator": "TPU",
    "colab": {
      "machine_shape": "hm",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "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.10.11"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}