File size: 22,834 Bytes

b386992

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "SUOXg71A3w78"
   },
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "You can run either this notebook locally (if you have all the dependencies and a GPU) or on Google Colab.\n",
    "\n",
    "Instructions for setting up Colab are as follows:\n",
    "1. Open a new Python 3 notebook.\n",
    "2. Import this notebook from GitHub (File -> Upload Notebook -> \"GITHUB\" tab -> copy/paste GitHub URL)\n",
    "3. Connect to an instance with a GPU (Runtime -> Change runtime type -> select \"GPU\" for hardware accelerator)\n",
    "4. Run this cell to set up dependencies.\n",
    "5. Restart the runtime (Runtime -> Restart Runtime) for any upgraded packages to take effect\n",
    "\n",
    "\n",
    "NOTE: User is responsible for checking the content of datasets and the applicable licenses and determining if suitable for the intended use.\n",
    "\"\"\"\n",
    "# If you're using Google Colab and not running locally, run this cell.\n",
    "import os\n",
    "\n",
    "# # Install dependencies\n",
    "!pip install wget\n",
    "!apt-get install sox libsndfile1 ffmpeg\n",
    "!pip install text-unidecode\n",
    "!pip install matplotlib>=3.3.2\n",
    "\n",
    "## Install NeMo\n",
    "BRANCH = 'main'\n",
    "!python -m pip install git+https://github.com/NVIDIA/NeMo.git@$BRANCH#egg=nemo_toolkit[all]\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "RGNuJWr66C38"
   },
   "source": [
    "# Automatic Speech Recognition with Transducer Models using HF Datasets\n",
    "\n",
    "We have discussed training various ASR models in NeMo using custom datasets, either for fine-tuning or for scratch-training. In this tutorial, we will showcase how to use Hugging Face datasets library in order to finetune a Transducer ASR model on a small dataset from for the Telugu language. \n",
    "It includes discussions relevant to preparing datasets with HF and how to use them to finetune NeMo models. The same method applies to training from scratch. However, for training, we recommend using scripts directly from the `examples/asr/` folder.\n",
    "\n",
    "In this tutorial, we demonstrate the usage of HF datasets for the Telugu language, where we use the Fluers dataset for training, validation, and testing. However, the same procedure can be used for other languages or domains and finetuned for specific use cases accordingly. \n",
    "\n",
    "For scripts, refer to [speech_to_text_finetune.py]('https://github.com/NVIDIA/NeMo/blob/main/examples/asr/speech_to_text_finetune.py') for training from scratch. \n",
    "\n",
    "--------\n",
    "\n",
    "**Note**: It is assumed that the previous tutorial - \"Intro-to-Transducers\" has been reviewed, and there is some familiarity with the config components of transducer models."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "eAqCcJ-T6C6k"
   },
   "source": [
    "# Preparing the dataset\n",
    "\n",
    "In this tutorial, we will be utilizing the `Telugu` language dataset part of fluers set. More details about dataset can be found at: https://huggingface.co/datasets/google/fleurs/viewer/te_in \n",
    "\n",
    "Dataset consists of 3 splits: `train`, `validation`, and `test`. We will use the `train` and `validation` splits for training and validation, respectively. The `test` split will be used for testing the model. Train set consists of 2,300 utterances and validation set consists of 311 examples. \n",
    "\n",
    "You may listen to some of the samples from the dataset using the following code:\n",
    "\n",
    "```python\n",
    "\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import datasets\n",
    "import IPython.display as ipd\n",
    "# sample from the validation set\n",
    "sample = datasets.load_dataset('google/fleurs', 'te_in', split='validation')[0]\n",
    " \n",
    "# Let's listen to the audio file and its transcription\n",
    "ipd.Audio(sample['audio']['array'], rate=16000)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Lets look at the transcription\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Transcription:\", sample['transcription'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Audio samples are in the `audio` field and the corresponding transcriptions are in the `transcription` field. We will use these fields to prepare data and train our model.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Preparing tokenizer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Since we are finetuning Parakeet model, which is an English language model, we need to update the tokenizer and update the decoder to support the new language. \n",
    "\n",
    "First, we will extract text transcriptions from the dataset and use them to train a tokenizer. We will use the scripts from NeMo first to get the data from HF dataset using `get_hf_dataset.py` script. Next we use `process_asr_text_tokenizer.py` script to prepare the tokenizer from [scripts](https://github.com/NVIDIA/NeMo/tree/main/scripts/tokenizers) folder. \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Download the `get_hf_text_data.py` script from the [scripts](https://github.com/NVIDIA/NeMo/blob/main/scripts/tokenizers) folder and run the following command to get the data from HF dataset. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "i2hD4LkoJvrx"
   },
   "outputs": [],
   "source": [
    "if not os.path.exists(\"scripts/get_hf_text_data.py\"):\n",
    "  !wget -P scripts/ https://raw.githubusercontent.com/NVIDIA/NeMo/$BRANCH/scripts/tokenizers/get_hf_text_data.py"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Major difference from NeMo dataset configs to training using HF-datasets is for using hf-datasets, users need to mention the hf-dataset information through hf data config and pass to the script for downloading necessary data. Users can switch to another dataset by changing necessary fields in the hf data config. \n",
    "Let's create that config here."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from omegaconf import OmegaConf, open_dict\n",
    "train_split = {\n",
    "            'path': 'google/fleurs',\n",
    "            'name': 'te_in',\n",
    "            'split': 'train', # we will update this accordingly based on the split\n",
    "            'streaming': False}\n",
    "print(OmegaConf.to_yaml(train_split))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Since we need clean data for training tokenizer and models, we need to filter the data based on how dataset was constructed and how we would like the ASR model output to be. \n",
    "\n",
    "Based on prior analysis of text transcripts of the current hf dataset, we skip all non-alphanumeric characters except `full-stop` using `normalize_text` option of the `get_hf_text_data.py` script, based on `huggingface_data_tokenizer.yaml` config file. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "if not os.path.exists('configs/huggingface_data_tokenizer.yaml'):\n",
    "    !wget -P configs/ https://raw.githubusercontent.com/NVIDIA/NeMo/$BRANCH/scripts/tokenizers/conf/huggingface_data_tokenizer.yaml\n",
    "\n",
    "\n",
    "!export HYDRA_FULL_ERROR=1;python scripts/get_hf_text_data.py \\\n",
    "    --config-path=\"../configs\" \\\n",
    "    --config-name=\"huggingface_data_tokenizer\" \\\n",
    "    normalize_text=True \\\n",
    "    symbols_to_keep=[\".\"] \\\n",
    "    text_key=\"transcription\" \\\n",
    "    output_file='telugu_train_corpus.txt' \\\n",
    "    +hf_data_cfg='[{train_split}]'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "From the above command, we were able to use the `huggingface_data_tokenizer.yaml` config file to download the data from HF dataset. The download data is saved to `telugu_train_corpus.txt` file, which we will use to train the tokenizer. Before that, let's look at some utterances from the normalized (preprocessed) text file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('telugu_train_corpus.txt', 'r') as f:\n",
    "    for i in range(5):\n",
    "        print(f.readline().strip())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As can be seen, there are characters from English and also numerical characters as well. For this tutorial we will retain those characters and prepare the tokenizer. But if you would like to filter out those characters, you can use your own function in the `get_hf_text_data.py` script to filter out the characters. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, we will use the `process_asr_text_tokenizer.py` script to prepare the tokenizer, for finetuning we can limit the vocab size to 256 for telugu language."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "if not os.path.exists(\"scripts/process_asr_text_tokenizer.py\"):\n",
    "  !wget -P scripts/ https://raw.githubusercontent.com/NVIDIA/NeMo/$BRANCH/scripts/tokenizers/process_asr_text_tokenizer.py\n",
    "\n",
    "# Now this downloads the text corpus of data to tokenizers script\n",
    "VOCAB_SIZE = 256  # can be any value above 29\n",
    "TOKENIZER_TYPE = \"spe\"  # can be wpe or spe\n",
    "SPE_TYPE = \"bpe\"  # can be bpe or unigram\n",
    "\n",
    "# ------------------------------------------------------------------- #\n",
    "!rm -r tokenizers/\n",
    "\n",
    "if not os.path.exists(\"tokenizers\"):\n",
    "  os.makedirs(\"tokenizers\")\n",
    "\n",
    "!python scripts/process_asr_text_tokenizer.py \\\n",
    "    --data_file='telugu_train_corpus.txt' \\\n",
    "    --data_root=\"tokenizers\" \\\n",
    "    --tokenizer=$TOKENIZER_TYPE \\\n",
    "    --spe_type=$SPE_TYPE \\\n",
    "    --no_lower_case \\\n",
    "    --log \\\n",
    "    --vocab_size=$VOCAB_SIZE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "JHDZswN6LIBJ"
   },
   "outputs": [],
   "source": [
    "# Tokenizer path\n",
    "if TOKENIZER_TYPE == 'spe':\n",
    "  TOKENIZER = os.path.join(\"tokenizers\", f\"tokenizer_spe_{SPE_TYPE}_v{VOCAB_SIZE}\")\n",
    "  TOKENIZER_TYPE_CFG = \"bpe\"\n",
    "else:\n",
    "  TOKENIZER = os.path.join(\"tokenizers\", f\"tokenizer_wpe_v{VOCAB_SIZE}\")\n",
    "  TOKENIZER_TYPE_CFG = \"wpe\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Tokenizers are saved at:\", TOKENIZER)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "j8tx2m2w6C87"
   },
   "source": [
    "# Preparing a Transducer Model\n",
    "\n",
    "Now that we have the dataset and tokenizer prepared, let us begin by setting up the config of the Transducer model! In this tutorial, we will finetune [Parakeet RNNT 0.6B](https://huggingface.co/nvidia/parakeet-rnnt-0.6b) model, but the same procedure can be used for other RNNT models as well.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "pQETBCSML0Us"
   },
   "source": [
    "## Prepare the config \n",
    "For finetuning the model, we need to update the config file to include the tokenizer and the dataset information. We will use the `parakeet-rnnt-6b` model and update the config file to include the tokenizer and the dataset information. For this we use `speech_to_text_hf_finetune.yaml` config file, and training script `speech_to_text_finetune.py` from the `examples/asr` folder.\n",
    "\n",
    "For this demo, we will replicate only a portion of the script - in order to show just the necessary components for training the model on single GPU. However, we recommend users to use the scripts directly from the `examples/asr` folder for training the model on multiple GPUs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "6a_vedo0Lyo8"
   },
   "outputs": [],
   "source": [
    "## Grab the config we'll use in this example\n",
    "!mkdir -p configs\n",
    "if not os.path.exists('configs/speech_to_text_hf_finetune.yaml'):\n",
    "    !wget -P configs/ https://raw.githubusercontent.com/NVIDIA/NeMo/$BRANCH/examples/asr/conf/asr_finetune/speech_to_text_hf_finetune.yaml\n",
    "\n",
    "config = OmegaConf.load(\"configs/speech_to_text_hf_finetune.yaml\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Update tokenizer info"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "config.model.tokenizer.update_tokenizer = True\n",
    "config.model.tokenizer.dir = TOKENIZER\n",
    "config.model.tokenizer.type = TOKENIZER_TYPE_CFG\n",
    "print(OmegaConf.to_yaml(config.model.tokenizer))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "N-t7fRl6GS3A"
   },
   "source": [
    "Need to mention the same data preprocessing scripts here as HF datasets are processed batchwise and the same preprocessing scripts are used for applying the same preprocessing to the dataset, as was used for training the tokenizer. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "B9nY5JQaIhKz"
   },
   "outputs": [],
   "source": [
    "\n",
    "config.model.train_ds.normalize_text=False\n",
    "config.model.normalize_text = True # same as the normalize_text in the get_hf_text_data.py\n",
    "config.model.symbols_to_keep=[\".\"] # same as the symbols_to_keep in the get_hf_text_data.py\n",
    "config.model.data_path = 'google/fleurs'\n",
    "config.model.data_name = 'te_in'\n",
    "config.model.streaming = False\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "d0tyOaqNLq-4"
   },
   "source": [
    "Setting up dataset config for validation and test datasets is similar to the training dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "mJbqMEkjMTfM"
   },
   "outputs": [],
   "source": [
    "config.model.train_ds.hf_data_cfg=train_split\n",
    "config.model.train_ds.text_key='transcription'\n",
    "config.model.train_ds.batch_size=16 # change this based on your GPU memory.\n",
    "config.model.train_ds.normalize_text=True\n",
    "\n",
    "config.model.validation_ds.hf_data_cfg=train_split\n",
    "config.model.validation_ds.hf_data_cfg.split='validation' # updated this based on the split\n",
    "config.model.validation_ds.text_key='transcription'\n",
    "config.model.validation_ds.normalize_text=True\n",
    "\n",
    "config.model.test_ds.hf_data_cfg=train_split\n",
    "config.model.test_ds.hf_data_cfg.split='test' # updated this based on the split\n",
    "config.model.test_ds.text_key='transcription'\n",
    "config.model.test_ds.normalize_text=True\n",
    "\n",
    "print(OmegaConf.to_yaml(config.model))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Update the trainer config to include the necessary parameters for training the model."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training the model\n",
    "\n",
    "For finetuning the model, basic parameters we need to primarily focus on are optimizer, scheduler, learning rate and the number of steps to finetune.\n",
    "\n",
    "For this example we will stick to AdamW optimizer and CosineAnnealing scheduler, for 2500 steps. \n",
    "These parameters can be updated based on the dataset and the model being used."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(OmegaConf.to_yaml(config.model.optim))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Current optim looks fine except for the warmup_steps\n",
    "config.model.optim.sched.warmup_steps = 500 # 10% of the total steps\n",
    "config.model.optim.lr = 3e-4\n",
    "\n",
    "del config.model.spec_augment #For this example, we are not using SpecAugment"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's load the model and also the dataloader for training the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Unv8-GvOWhad"
   },
   "outputs": [],
   "source": [
    "from nemo.collections.asr.models import ASRModel\n",
    "asr_model=ASRModel.from_pretrained('nvidia/parakeet-rnnt-0.6b')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After loading the pretrained model, decoder of the model needs to updated to support the new language."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "decoder=asr_model.decoder.state_dict()\n",
    "joint_state = asr_model.joint.state_dict()\n",
    "prev_vocab_size = asr_model.tokenizer.vocab_size\n",
    "asr_model.change_vocabulary(new_tokenizer_dir=TOKENIZER, new_tokenizer_type=TOKENIZER_TYPE_CFG)\n",
    "if asr_model.tokenizer.vocab_size == prev_vocab_size: # checking new tokenizer vocab size\n",
    "    asr_model.decoder.load_state_dict(decoder)\n",
    "    asr_model.joint.load_state_dict(joint_state)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Loading dataloaders"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "VsunP99saF5c"
   },
   "outputs": [],
   "source": [
    "from nemo.utils import logging, model_utils\n",
    "cfg = model_utils.convert_model_config_to_dict_config(config)\n",
    "asr_model.setup_training_data(cfg.model.train_ds)\n",
    "asr_model.setup_validation_data(cfg.model.validation_ds)\n",
    "if hasattr(cfg.model, 'test_ds') and cfg.model.test_ds.manifest_filepath is not None:\n",
    "    asr_model.setup_test_data(cfg.model.test_ds)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "yT1vJH9OkS0u"
   },
   "outputs": [],
   "source": [
    "# We will also reduce the hidden dimension of the joint and the prediction networks to preserve some memory\n",
    "asr_model.setup_optimization(cfg.model.optim)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "SH4ZfXHOdGhX"
   },
   "source": [
    "------\n",
    "\n",
    "Setup a Pytorch Lightning Trainer:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can reduce the number of steps to 5000 for this small dataset.\n",
    "and also update precision to float16 for faster training. # Change this bf16 training on Ampere based GPUs. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Fmf0iSY-a6LC"
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "from lightning.pytorch import Trainer\n",
    "\n",
    "if torch.cuda.is_available():\n",
    "  accelerator = 'gpu'\n",
    "else:\n",
    "  accelerator = 'gpu'\n",
    "\n",
    "MAX_STEPS = 5000\n",
    "\n",
    "# Initialize a Trainer for the Transducer model\n",
    "trainer = Trainer(devices=1, accelerator=accelerator, max_epochs=-1, max_steps=MAX_STEPS,\n",
    "                  enable_checkpointing=False, logger=False,\n",
    "                  log_every_n_steps=100, check_val_every_n_epoch=10, precision='bf16')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "RheLsmA1cRz0"
   },
   "outputs": [],
   "source": [
    "# Build the model\n",
    "trainer.fit(asr_model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Once trained users can save the model and use it for inference."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "asr_model.save_to(\"telugu_asr_model.nemo\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Running inference on  a sample from the test dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "del asr_model\n",
    "asr_model = ASRModel.restore_from(\"telugu_asr_model.nemo\")\n",
    "\n",
    "!rm -r telugu_asr_model.nemo # remove the model to save space\n",
    "\n",
    "# Let's test the model on a sample from the test set\n",
    "sample = datasets.load_dataset('google/fleurs', 'te_in', split='test')[0]\n",
    "transcription = sample['transcription']\n",
    "\n",
    "# Let's listen to the audio file and its transcription\n",
    "ipd.Audio(sample['audio']['array'], rate=16000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "asr_model.eval()\n",
    "# save audio file to disk\n",
    "audio_path = \"sample.wav\"\n",
    "import soundfile as sf\n",
    "sf.write(audio_path, sample['audio']['array'], 16000)\n",
    "transcription = asr_model.transcribe([audio_path])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "print(\"Predicted:\", transcription[0])\n",
    "print(\"Target:\", sample['transcription'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "os.remove(audio_path) # remove the audio file to save space"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [
    "V5sMoFHmVvhg"
   ],
   "name": "ASR-with-Transducers.ipynb",
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}