ESPnetEZ/SLU/SLU_finetune_owsm.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# OWSM Fine-tuning for Spoken Language Understanding\n",
    "\n",
    "This Jupyter notebook provides a step-by-step guide on using the ESPnetEZ module to fine-tune the OWSM model. In this demonstration, we will use the `SLURP` dataset (intent classification task) to fine-tune an OWSM model for a Spoken Language Understanding (SLU) task.\n",
    "\n",
    "This demo will focus on how to add new tokens to the pre-trained OWSM model for the intent classification task. For other sections such as data preparation, training, and evaluation, please refer to the other notebooks.\n",
    "\n",
    "In this notebook, we assume that you have already downloaded the `SLURP` dataset and created the dump file using the recipe. If you haven't done this before and are unfamiliar with the recipes provided in ESPnet, you can refer to the data preparation sections in the `train_from_scratch.ipynb` or `finetune_owsm.ipynb` notebooks in the ASR demos.\n",
    "\n",
    "Author: Masao Someki [@Masao-Someki](https://github.com/Masao-Someki)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, let's install ESPnet if you haven't already:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "plaintext"
    }
   },
   "outputs": [],
   "source": [
    "!pip install -U espnet"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then import necessary libraries and set several hyper parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import numpy as np\n",
    "import librosa\n",
    "from pathlib import Path\n",
    "from espnet2.layers.create_adapter_fn import create_lora_adapter\n",
    "import argparse\n",
    "\n",
    "import espnetez as ez\n",
    "\n",
    "\n",
    "FINETUNE_MODEL = \"espnet/owsm_v3.1_ebf_base\"\n",
    "\n",
    "DATA_PATH = \"./data\"\n",
    "DUMP_DIR = \"./dump/raw\"\n",
    "STATS_DIR = \"./exp/stats_owsm\"\n",
    "\n",
    "ADDITIONAL_SPECIAL_TOKENS = [\n",
    "    \"<intent>\"\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Adding a New Token to the Pre-trained OWSM Tokenizer\n",
    "\n",
    "In this section, we will add a new `<intent>` token to the pre-trained OWSM tokenizer. We have prepared the `add_special_tokens` function from the `espnet2.preprocess` module to add new tokens to the pre-trained tokenizer, converter, and the Embedding layer.\n",
    "\n",
    "We will use the Embedding layer in the `build_model_fn` function to replace the pre-trained Embedding layer with the new one."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from espnet2.bin.s2t_inference import Speech2Text\n",
    "pretrained_model = Speech2Text.from_pretrained(\n",
    "    FINETUNE_MODEL,\n",
    "    # category_sym=\"<en>\",\n",
    "    beam_size=10,\n",
    "    device=\"cpu\"\n",
    ")\n",
    "tokenizer = pretrained_model.tokenizer\n",
    "converter = pretrained_model.converter\n",
    "\n",
    "# Add new <intent_cls> token after ST-related tokens\n",
    "tokenizer, converter, _ = ez.preprocess.add_special_tokens(\n",
    "    tokenizer, converter, pretrained_model.s2t_model.decoder.embed[0],\n",
    "    ADDITIONAL_SPECIAL_TOKENS, insert_after=\"<st_zho>\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# And load configuration of pre-trained model. before deleting it.\n",
    "training_config = vars(pretrained_model.s2t_train_args)\n",
    "del pretrained_model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data Preparation\n",
    "\n",
    "To create text with the new `<intent>` token, we need to modify the dump file generated by the recipe. Specifically, we want to change the `text` format to:\n",
    "\n",
    "`<eng><intent><notimestamps> intent`\n",
    "\n",
    "To achieve this, we need to write a custom dataset class and `data_info` with appropriate functions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CustomDataset:\n",
    "    def __init__(self, data_path, is_train=True):\n",
    "        self.data_path = data_path\n",
    "        if is_train:\n",
    "            data_path = f\"{data_path}/train\"\n",
    "        else:\n",
    "            data_path = f\"{data_path}/devel\"\n",
    "        \n",
    "        self.data = {}\n",
    "        with open(f\"{data_path}/wav.scp\", \"r\") as f:\n",
    "            for line in f.readlines():\n",
    "                audio_id, audio_path = line.strip().split(maxsplit=1)\n",
    "                self.data[audio_id] = {\n",
    "                    'audio_path': audio_path\n",
    "                }\n",
    "\n",
    "        with open(f\"{data_path}/transcript\", \"r\") as f:\n",
    "            for line in f.readlines():\n",
    "                audio_id, translated = line.strip().split(maxsplit=1)\n",
    "                self.data[audio_id]['transcript'] = translated\n",
    "        \n",
    "        with open(f\"{data_path}/text\", \"r\") as f:\n",
    "            for line in f.readlines():\n",
    "                audio_id, intent, _ = line.strip().split(maxsplit=2)\n",
    "                self.data[audio_id]['intent'] = intent\n",
    "        \n",
    "        self.keys = list(self.data.keys())\n",
    "    \n",
    "    def __len__(self):\n",
    "        return len(self.keys)\n",
    "    \n",
    "    def __getitem__(self, idx):\n",
    "        return {\n",
    "            'audio_path': self.data[idx]['audio_path'],\n",
    "            'intent': self.data[idx]['intent'],\n",
    "            'transcript': self.data[idx]['transcript']\n",
    "        }\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def tokenize(text):\n",
    "    return np.array(converter.tokens2ids(tokenizer.text2tokens(text)))\n",
    "\n",
    "data_info = {\n",
    "    \"speech\": lambda d : librosa.load(d['audio_path'], sr=16000)[0],\n",
    "    \"text\": lambda d : tokenize(f\"<eng><intent><notimestamps>{d['intent']}\"),\n",
    "    \"text_prev\": lambda d : tokenize(\"<na>\"),\n",
    "    \"text_ctc\": lambda d : tokenize(d['transcript'].lower()),\n",
    "} "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Model Preparation\n",
    "\n",
    "Let's prepare the `build_model_fn` function for the Trainer. Inside the `build_model_fn` function, we will replace the pre-trained Embedding layer to the new one."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def count_parameters(model):\n",
    "    return sum(p.numel() for p in model.parameters() if p.requires_grad)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def build_model_fn(args):\n",
    "    pretrained_model = Speech2Text.from_pretrained(\n",
    "        FINETUNE_MODEL,\n",
    "        beam_size=10,\n",
    "    )\n",
    "    model = pretrained_model.s2t_model\n",
    "    model.train()\n",
    "    print(f'Trainable parameters: {count_parameters(model)}')\n",
    "    \n",
    "    # Add new <intent> token\n",
    "    _, _, new_embedding = ez.preprocess.add_special_tokens(\n",
    "        pretrained_model.tokenizer,\n",
    "        pretrained_model.converter,\n",
    "        model.decoder.embed[0],\n",
    "        ADDITIONAL_SPECIAL_TOKENS,\n",
    "        insert_after=\"<st_zho>\"\n",
    "    )\n",
    "    new_embedding.weight.requires_grad = True\n",
    "    model.decoder.embed[0] = new_embedding\n",
    "    # apply lora if you want.\n",
    "    # create_lora_adapter(model, target_modules=LORA_TARGET)\n",
    "    # print(f'Trainable parameters after LORA: {count_parameters(model)}')\n",
    "    return model\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training Configuration\n",
    "\n",
    "Now let's set up the training configuration for the OWSM finetuning. Basically all configurations are the same as the OWSM training, but we will change some parameters for this finetuning."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# When you load your configuration from yaml file, you can write with update_finetune_config\n",
    "finetune_config = ez.config.update_finetune_config(\n",
    "    \"s2t\",\n",
    "    training_config,\n",
    "    \"owsm_finetune_base.yaml\",\n",
    ")\n",
    "\n",
    "# When you don't use yaml file, you can load finetune_config in the following way.\n",
    "# task_class = ez.task.get_ez_task(\"s2t\")\n",
    "# default_config = task_class.get_default_config()\n",
    "# finetune_config = default_config.update(your_config_in_dict)\n",
    "\n",
    "finetune_config['multiple_iterator'] = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataset = CustomDataset(data_path=\"./dump/raw\", is_train=True)\n",
    "dev_dataset = CustomDataset(data_path=\"./dump/raw\", is_train=False)\n",
    "\n",
    "train_dataset = ez.dataset.ESPnetEZDataset(train_dataset, data_info=data_info)\n",
    "dev_dataset = ez.dataset.ESPnetEZDataset(dev_dataset, data_info=data_info)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training\n",
    "\n",
    "Now we have everything prepared, we can start training the OWSM model for ST task."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer = ez.Trainer(\n",
    "    task=\"s2t\",\n",
    "    train_config=finetune_config,\n",
    "    train_dataset=train_dataset,\n",
    "    valid_dataset=dev_dataset,\n",
    "    data_info=data_info,\n",
    "    build_model_fn=build_model_fn,\n",
    "    output_dir=EXP_DIR,\n",
    "    stats_dir=STATS_DIR,\n",
    "    ngpu=1,\n",
    ")\n",
    "trainer.collect_stats()\n",
    "trainer.train()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Finetuning configuration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "yaml"
    }
   },
   "outputs": [],
   "source": [
    "seed: 2022\n",
    "num_workers: 8\n",
    "batch_type: unsorted\n",
    "batch_size: 1\n",
    "batch_bins: 500000\n",
    "accum_grad: 1\n",
    "max_epoch: 10\n",
    "patience: none\n",
    "init: none\n",
    "best_model_criterion:\n",
    "-   - valid\n",
    "    - acc\n",
    "    - max\n",
    "keep_nbest_models: 3\n",
    "use_amp: true\n",
    "ngpu: 1\n",
    "\n",
    "optim: adamw\n",
    "optim_conf:\n",
    "    lr: 0.0001\n",
    "    weight_decay: 0.000001\n",
    "scheduler: warmuplr\n",
    "scheduler_conf:\n",
    "    warmup_steps: 15000\n",
    "\n",
    "specaug: specaug\n",
    "specaug_conf:\n",
    "    apply_time_warp: true\n",
    "    time_warp_window: 5\n",
    "    time_warp_mode: bicubic\n",
    "    apply_freq_mask: true\n",
    "    freq_mask_width_range:\n",
    "    - 0\n",
    "    - 27\n",
    "    num_freq_mask: 2\n",
    "    apply_time_mask: true\n",
    "    time_mask_width_ratio_range:\n",
    "    - 0.\n",
    "    - 0.05\n",
    "    num_time_mask: 5\n",
    "\n",
    "\n",
    "use_preprocessor: false\n",
    "preprocessor: default\n",
    "preprocessor_conf:\n",
    "    fs: 16000\n",
    "    text_name:\n",
    "        - \"text\""
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}