ASR.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "48779dab",
   "metadata": {},
   "source": [
    "Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "3fa4076b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import (\n",
    "    AutoProcessor,\n",
    "    AutoModelForSpeechSeq2Seq,\n",
    "    Wav2Vec2BertForCTC,\n",
    "    Wav2Vec2BertProcessor,\n",
    "    Wav2Vec2CTCTokenizer,\n",
    "    SeamlessM4TFeatureExtractor,\n",
    "    AutoProcessor,\n",
    "    Trainer,\n",
    "    Seq2SeqTrainingArguments,\n",
    "    AutoConfig,\n",
    ")\n",
    "import json\n",
    "import os\n",
    "import random\n",
    "import torch\n",
    "from dataclasses import dataclass\n",
    "from typing import Dict, List, Union\n",
    "from accelerate import Accelerator\n",
    "import warnings\n",
    "from tqdm.auto import tqdm\n",
    "import huggingface_hub\n",
    "import datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "5a12a60b",
   "metadata": {},
   "outputs": [],
   "source": [
    "BATCH_SIZE = 64\n",
    "BASE_MODEL = \"flozi00/distilwhisper-german-canary\"\n",
    "OUT_MODEL = \"distilwhisper-german-canary-v2\"\n",
    "DATASET_NAME=\"flozi00/german-canary-asr-0324\"\n",
    "DATASET_SUBSET=\"default\"\n",
    "AUDIO_PATH=\"audio\"\n",
    "EPOCHS = 1\n",
    "SAVE_STEPS = 1000\n",
    "LR= 1e-5"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c738c2d8",
   "metadata": {},
   "source": [
    "Data Class for Collator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "31d7cc27",
   "metadata": {},
   "outputs": [],
   "source": [
    "@dataclass\n",
    "class ASRDataCollator:\n",
    "    processor: AutoProcessor\n",
    "    wav_key: str = os.getenv(\"AUDIO_PATH\", \"audio\")\n",
    "    locale_key: str = os.getenv(\"LOCALE_KEY\", \"german\")\n",
    "    text_key: str = os.getenv(\"TEXT_KEY\", \"transkription\")\n",
    "    max_audio_in_seconds: float = float(os.getenv(\"MAX_AUDIO_IN_SECONDS\", 20.0))\n",
    "\n",
    "    def __call__(\n",
    "        self, features: List[Dict[str, Union[List[int], torch.Tensor]]]\n",
    "    ) -> Dict[str, torch.Tensor]:\n",
    "        input_features = []\n",
    "        label_features = []\n",
    "\n",
    "        for i in range(len(features)):\n",
    "            try:\n",
    "                feature = features[i]\n",
    "\n",
    "                myaudio = feature[self.wav_key][\"array\"]\n",
    "                mytext = feature[self.text_key]\n",
    "            except Exception as e:\n",
    "                print(e)\n",
    "                continue\n",
    "\n",
    "            audio_len = int((len(myaudio) / 16000))\n",
    "            if audio_len > self.max_audio_in_seconds:\n",
    "                print(\"skipping audio\")\n",
    "                continue\n",
    "\n",
    "            # Extract the text from the feature and normalize it\n",
    "            mylang = self.locale_key\n",
    "\n",
    "            # Extract the audio features from the audio\n",
    "            extracted = self.processor.feature_extractor(\n",
    "                myaudio,\n",
    "                sampling_rate=16000,\n",
    "                return_tensors=\"pt\",\n",
    "            )\n",
    "\n",
    "            # check if feature extractor return input_features or input_values\n",
    "            ft = (\n",
    "                \"input_values\"\n",
    "                if hasattr(extracted, \"input_values\")\n",
    "                else \"input_features\"\n",
    "            )\n",
    "\n",
    "            # append to input_features\n",
    "            input_features.append(\n",
    "                {\n",
    "                    ft: getattr(\n",
    "                        extracted,\n",
    "                        ft,\n",
    "                    )[0]\n",
    "                }\n",
    "            )\n",
    "\n",
    "            # set prefix tokens if possible\n",
    "            try:\n",
    "                self.processor.tokenizer.set_prefix_tokens(mylang)\n",
    "            except Exception:\n",
    "                pass\n",
    "\n",
    "            # append to label_features and tokenize\n",
    "            label_features.append(\n",
    "                {\"input_ids\": self.processor.tokenizer(mytext).input_ids}\n",
    "            )\n",
    "\n",
    "        batch = self.processor.feature_extractor.pad(\n",
    "            input_features,\n",
    "            padding=\"longest\",\n",
    "            return_tensors=\"pt\",\n",
    "        )\n",
    "\n",
    "        labels_batch = self.processor.tokenizer.pad(\n",
    "            label_features,\n",
    "            padding=\"longest\",\n",
    "            return_tensors=\"pt\",\n",
    "        )\n",
    "\n",
    "        # replace padding with -100 to ignore loss correctly\n",
    "        labels = labels_batch[\"input_ids\"].masked_fill(\n",
    "            labels_batch.attention_mask.ne(1), -100\n",
    "        )\n",
    "\n",
    "        batch[\"labels\"] = labels\n",
    "\n",
    "        return batch"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "67ab4af6",
   "metadata": {},
   "source": [
    "preprocessing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "7b07e4fc",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_model(\n",
    "    model_name: str\n",
    "):\n",
    "    kwargs = {}\n",
    "    # get the config of the base model and extract the model type from it\n",
    "    conf = AutoConfig.from_pretrained(\n",
    "        pretrained_model_name_or_path=model_name,\n",
    "        trust_remote_code=True,\n",
    "    )\n",
    "\n",
    "\n",
    "    processor = AutoProcessor.from_pretrained(\n",
    "        model_name,\n",
    "        legacy=False,\n",
    "        trust_remote_code=True,\n",
    "    )\n",
    "\n",
    "    kwargs[\"attn_implementation\"] = \"sdpa\"\n",
    "\n",
    "    model = AutoModelForSpeechSeq2Seq.from_pretrained(\n",
    "        model_name,\n",
    "        config=conf,\n",
    "        low_cpu_mem_usage=True,\n",
    "        torch_dtype=torch.bfloat16,\n",
    "        **kwargs,\n",
    "    )\n",
    "    \n",
    "    model = model.cuda()\n",
    "\n",
    "    return model, processor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "7d31877c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "77f2983095794094a700f939519845c5",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Resolving data files:   0%|          | 0/84 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4b49868be9614e28a2aa94de4712af76",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Resolving data files:   0%|          | 0/84 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "cv_data = datasets.load_dataset(\n",
    "    DATASET_NAME,\n",
    "    DATASET_SUBSET,\n",
    "    split=\"train\",\n",
    "    streaming=True,\n",
    ").cast_column(\n",
    "    AUDIO_PATH,\n",
    "    datasets.Audio(sampling_rate=16000, decode=True),\n",
    ").with_format(\n",
    "    \"torch\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3235b8f7",
   "metadata": {},
   "source": [
    "Main function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "9752cb52",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n"
     ]
    }
   ],
   "source": [
    "model, processor = get_model(\n",
    "    model_name=BASE_MODEL\n",
    ")\n",
    "\n",
    "try:\n",
    "    model.config.forced_decoder_ids = None\n",
    "    model.config.suppress_tokens = []\n",
    "except Exception:\n",
    "    pass\n",
    "\n",
    "cv_data = cv_data.shuffle(seed=random.randint(0, 1000))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "46ac97a6",
   "metadata": {},
   "outputs": [],
   "source": [
    "training_args = Seq2SeqTrainingArguments(output_dir=OUT_MODEL, per_device_train_batch_size=BATCH_SIZE, push_to_hub=True, gradient_checkpointing=True, \n",
    "                                        dataloader_pin_memory=False, hub_private_repo=True, bf16=True,\n",
    "                                        learning_rate=LR, optim=\"paged_lion_8bit\", max_steps=int(985000/BATCH_SIZE),\n",
    "                                        gradient_accumulation_steps=1, logging_steps=10, save_strategy=\"steps\",\n",
    "                                        save_steps=SAVE_STEPS, optim_target_modules=[\"attn\", \"mlp\"], remove_unused_columns=False,\n",
    "                                        gradient_checkpointing_kwargs={\"use_reentrant\": False})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "e487ee8e",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.10/dist-packages/accelerate/accelerator.py:432: FutureWarning: Passing the following arguments to `Accelerator` is deprecated and will be removed in version 1.0 of Accelerate: dict_keys(['dispatch_batches', 'split_batches', 'even_batches', 'use_seedable_sampler']). Please pass an `accelerate.DataLoaderConfiguration` instead: \n",
      "dataloader_config = DataLoaderConfiguration(dispatch_batches=None, split_batches=False, even_batches=True, use_seedable_sampler=True)\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "trainer = Trainer(\n",
    "    model=model,\n",
    "    args=training_args,\n",
    "    train_dataset=cv_data,\n",
    "    eval_dataset=None,\n",
    "    tokenizer=processor.feature_extractor,\n",
    "    data_collator=ASRDataCollator(processor=processor),\n",
    "    compute_metrics=None,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "9573a752",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Fri Apr  5 07:20:49 2024       \n",
      "+---------------------------------------------------------------------------------------+\n",
      "| NVIDIA-SMI 545.29.06              Driver Version: 545.29.06    CUDA Version: 12.3     |\n",
      "|-----------------------------------------+----------------------+----------------------+\n",
      "| GPU  Name                 Persistence-M | Bus-Id        Disp.A | Volatile Uncorr. ECC |\n",
      "| Fan  Temp   Perf          Pwr:Usage/Cap |         Memory-Usage | GPU-Util  Compute M. |\n",
      "|                                         |                      |               MIG M. |\n",
      "|=========================================+======================+======================|\n",
      "|   0  NVIDIA L40S                    Off | 00000000:00:10.0 Off |                    0 |\n",
      "| N/A   36C    P0              70W / 350W |   1997MiB / 46068MiB |      0%      Default |\n",
      "|                                         |                      |                  N/A |\n",
      "+-----------------------------------------+----------------------+----------------------+\n",
      "                                                                                         \n",
      "+---------------------------------------------------------------------------------------+\n",
      "| Processes:                                                                            |\n",
      "|  GPU   GI   CI        PID   Type   Process name                            GPU Memory |\n",
      "|        ID   ID                                                             Usage      |\n",
      "|=======================================================================================|\n",
      "+---------------------------------------------------------------------------------------+\n"
     ]
    }
   ],
   "source": [
    "!nvidia-smi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "130bb675",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[34m\u001b[1mwandb\u001b[0m: Currently logged in as: \u001b[33mflozi00\u001b[0m. Use \u001b[1m`wandb login --relogin`\u001b[0m to force relogin\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "wandb version 0.16.6 is available!  To upgrade, please run:\n",
       " $ pip install wandb --upgrade"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Tracking run with wandb version 0.16.5"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Run data is saved locally in <code>/workdir/wandb/run-20240405_072050-d5zd9klc</code>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Syncing run <strong><a href='https://wandb.ai/flozi00/huggingface/runs/d5zd9klc/workspace' target=\"_blank\">captain-directive-424</a></strong> to <a href='https://wandb.ai/flozi00/huggingface' target=\"_blank\">Weights & Biases</a> (<a href='https://wandb.me/run' target=\"_blank\">docs</a>)<br/>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       " View project at <a href='https://wandb.ai/flozi00/huggingface' target=\"_blank\">https://wandb.ai/flozi00/huggingface</a>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       " View run at <a href='https://wandb.ai/flozi00/huggingface/runs/d5zd9klc/workspace' target=\"_blank\">https://wandb.ai/flozi00/huggingface/runs/d5zd9klc/workspace</a>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "`use_cache = True` is incompatible with gradient checkpointing. Setting `use_cache = False`...\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "    <div>\n",
       "      \n",
       "      <progress value='115' max='61562' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      [  115/61562 02:56 < 26:40:53, 0.64 it/s, Epoch 0.00/9223372036854775807]\n",
       "    </div>\n",
       "    <table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       " <tr style=\"text-align: left;\">\n",
       "      <th>Step</th>\n",
       "      <th>Training Loss</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>100</td>\n",
       "      <td>0.053700</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table><p>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "train_result = trainer.train()\n",
    "trainer.save_model()  # Saves the feature extractor too for easy upload"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9a5aaab4",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}