training.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "7fa2feb8",
   "metadata": {},
   "source": [
    "#  Shared Task: Mozilla Common Voice Spontaneous Speech ASR\n",
    "#  Training\n",
    "*https://www.codabench.org/competitions/10820/*  \n",
    "*1st place solution*  \n",
    "*Copyright (c) 2025 Igor Ivanov*  \n",
    "*Email: vecxoz@gmail.com*  \n",
    "*License: MIT*  \n",
    "*I will be happy to answer any questions.*"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6dfa145d",
   "metadata": {},
   "source": [
    "# Contents\n",
    "\n",
    "1. Solution summary \n",
    "2. Installation\n",
    "3. Create corpus and KenLM models\n",
    "4. Training\n",
    "5. Quantization"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "43ff13da-4175-4345-a880-d59e26bbb8d7",
   "metadata": {},
   "source": [
    "# 1. Summary\n",
    "\n",
    "In this notebook we present a training code for all 4 tasks of the Shared Task: Mozilla Common Voice Spontaneous Speech ASR. We did not use external data. Only Common Voice datasets were used: spontaneous speech for 21 languages, and scripted speech for 5 unseen languages. We fine-tuned the MMS model with adapter layers per language.  \n",
    "\n",
    "Our best single model features the following improvements over the baseline. (1) More data. We used both training and validation subsets for fine-tuning. (2) Different pretrained checkpoint `facebook/mms-1b-l1107`. (3) Longer training for 30 epochs. (4) Learning rate schedules tailored for each language. (5) Beam search decoding with [KenLM](https://kheafield.com/code/kenlm/) language model. For the small model subtask we used the same MMS models with pruning and 4-bit quantization.  \n",
    "\n",
    "Our overall best submission is an ensemble of 4 models. (1) `facebook/mms-1b-l1107` fine-tuned using training data only. (2) `facebook/mms-1b-l1107` fine-tuned using all data (training and validation subsets). (3) `facebook/mms-1b-all` fine-tuned using all data. (4) `facebook/mms-1b-fl102` fine-tuned using all data. We applied the [ROVER](https://github.com/usnistgov/SCTK) ensembling method, which outperformed each single model.  \n",
    "\n",
    "Please find all details in the paper."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0ddd35e1-bbd4-4f56-a143-e8430cd7238c",
   "metadata": {},
   "source": [
    "# 2. Installation\n",
    "\n",
    "## Directory structure\n",
    "```\n",
    "solution_training\n",
    "|-- cv-corpus-23.0-2025-09-05             # Common Voice scripted corpus\n",
    "    |-- ady\n",
    "    |-- bas\n",
    "    |-- kbd\n",
    "    |-- qxp\n",
    "    |-- ush\n",
    "|-- kenlm_dist                            # KenLM (source and binaries)\n",
    "|-- mcv-sps-st-09-2025                    # Common Voice spontaneous corpus\n",
    "    |-- sps-corpus-1.0-2025-09-05-aln\n",
    "    |-- sps-corpus-1.0-2025-09-05-bew\n",
    "    |-- ...\n",
    "|-- collect_models.py\n",
    "|-- create_corpus.py\n",
    "|-- create_lm.py\n",
    "|-- LICENSE.txt\n",
    "|-- prune_quantize.py\n",
    "|-- README.md\n",
    "|-- requirements.txt\n",
    "|-- train_alldata_script.py               # Train using all data, scripted corpus\n",
    "|-- train_alldata_spont.py                # Train using all data, spontaneous corpus\n",
    "|-- training.ipynb\n",
    "|-- train_trdata_script.py                # Train using training subset only, scripted corpus\n",
    "|-- train_trdata_spont.py                 # Train using training subset only, spontaneous corpus\n",
    "```\n",
    "\n",
    "We included all datasets in the distribution archive. You can use the following links to re-download some of the data.  \n",
    "\n",
    "1) Spontaneous corpus dedicated to the Shared Task, 26 languages, corresponding to `mcv-sps-st-09-2025` directory.  \n",
    "Mozilla Common Voice Spontaneous Speech ASR Shared Task Train/Dev Data  \n",
    "https://datacollective.mozillafoundation.org/datasets/cmfzu8u8wa555eq8onrk334h4\n",
    "3) Scripted corpus, 5 languages, corresponding to `cv-corpus-23.0-2025-09-05` directory. We used version 23. The current version is 24 and direct links do not work. We didn't find the way to get version 23 and did not estimate the difference between v23 and v24.  \n",
    "`ady`: Common Voice Scripted Speech 23.0 - Adyghe  \n",
    "`bas`: Common Voice Scripted Speech 23.0 - Basaa  \n",
    "`kbd`: Common Voice Scripted Speech 23.0 - Kabardian  \n",
    "`qxp`: Common Voice Scripted Speech 23.0 - Puno Quechua  \n",
    "`ush`: Common Voice Scripted Speech 23.0 - Ushojo  \n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f734ecae-d352-439a-aa18-96fd2546702e",
   "metadata": {},
   "source": [
    "## KenLM binary check\n",
    "\n",
    "**License note.** KenLM is licensed under LGPL. We did not modify it and use it via \"dynamic linking\" i.e. calling binary from Python. In this scenario LGPL terms allow to use arbitrary license for the code which calls LGPL binary. Specifically our code is licensed under MIT.\n",
    "\n",
    "We built a binary from source on Ubuntu 22.04. If you are using a similar system it should work out of the box, if you have the following system libs installed, especially `libboost`. Please run `sudo apt install ...` command below and then `lmplz --help` command. If a help message is displayed, then everything works. If not, you have to build it from source, as shown below or in official documentation: https://kheafield.com/code/kenlm/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bb2b904d-f9a7-4d0d-b2b4-0c6867c24614",
   "metadata": {},
   "outputs": [],
   "source": [
    "!sudo apt install -y build-essential cmake libboost-system-dev \\\n",
    "libboost-thread-dev libboost-program-options-dev libboost-test-dev \\\n",
    "libeigen3-dev zlib1g-dev libbz2-dev liblzma-dev"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ef5405c8-c6fc-4248-b2bc-4a6662650180",
   "metadata": {},
   "outputs": [],
   "source": [
    "!./kenlm_dist/build/bin/lmplz --help"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "35b537f7-36a3-40cd-9d79-892b150cd925",
   "metadata": {},
   "source": [
    "## KenLM installation (if needed)\n",
    "\n",
    "Compilation takes about 2 minutes."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cdaba278-3472-4289-8dfb-b1017310064e",
   "metadata": {},
   "source": [
    "```\n",
    "!mv kenlm_dist kenlm_dist_prebuilt\n",
    "\n",
    "!wget -O - https://kheafield.com/code/kenlm.tar.gz | tar xz\n",
    "# Rename to avoid import conflict with \"kenlm\" Python package which is installed independently\n",
    "!mv kenlm kenlm_dist\n",
    "!mkdir kenlm_dist/build\n",
    "%cd kenlm_dist/build\n",
    "!cmake ..\n",
    "!make -j 4\n",
    "\n",
    "%cd ../..\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8108da48-12aa-4663-9c20-3028b9b8117e",
   "metadata": {},
   "source": [
    "## Package installation\n",
    "\n",
    "**Hardware:**\n",
    "* Core-i5 CPU\n",
    "* 32 GB RAM\n",
    "* 500 GB SSD\n",
    "* RTX-3090-24GB GPU\n",
    "\n",
    "**System:**\n",
    "* Ubuntu 22.04\n",
    "* Python 3.12\n",
    "* CUDA 12.8\n",
    "* PyTorch 2.9.0\n",
    "\n",
    "**Note.** Flash Attention 2 (`flash_attn==2.7.4.post1`) is included in the `requirements.txt`. Installation (compilation) of this version takes about 2 hours on Core-i5, 12th Gen. If you don't need it, please remove it from `requirements.txt` and set parameter `--attn_implementation=sdpa` for all training scripts."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b5444ca5-e719-4e40-a851-1a2fc4900542",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install -r requirements.txt"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "240115c1-34d3-4748-95d6-0f5f221338e3",
   "metadata": {},
   "source": [
    "## Check Flash Attention 2\n",
    "\n",
    "If import is successful, then Flash Attention 2 was installed correctly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1d4e8525-8db1-4f66-ad1b-41c9b81ec93c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from flash_attn import flash_attn_qkvpacked_func, flash_attn_func"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d4b7630c-0475-4345-9643-e033878889f2",
   "metadata": {},
   "source": [
    "# 3. Create corpus and KenLM models\n",
    "\n",
    "**Note.**  We create KenLM models from all available data including validation subsets. It means that if you try to predict the validation subsets, predictions will be overfitted when using these models."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a2dbb97f-430f-4fc7-9152-4af63ef17df8",
   "metadata": {},
   "source": [
    "#### Create corpus"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f52a81a2-33d1-4bb4-a61f-c628ad0d63b4",
   "metadata": {},
   "outputs": [],
   "source": [
    "!python create_corpus.py \\\n",
    "--input_dir=./ \\\n",
    "--output_dir=./kenlm_corpus"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "326fcfad-cc36-40fa-bd39-391e3370392c",
   "metadata": {},
   "source": [
    "#### Create models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e900ba63-30ba-48f0-8cca-760d3605549e",
   "metadata": {},
   "outputs": [],
   "source": [
    "!python create_lm.py \\\n",
    "--input_dir=./kenlm_corpus \\\n",
    "--output_dir=./kenlm_models_order_3 \\\n",
    "--bin_path=./kenlm_dist/build/bin/lmplz \\\n",
    "--ngram_order=3"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eda488d4-8e09-4c22-9ff4-3675ed3bc9dc",
   "metadata": {},
   "source": [
    "# 4. Training\n",
    "\n",
    "We used an `RTX-3090-24GB` GPU for training.  \n",
    "There are 104 models total: 4 methods by 26 languages.  \n",
    "Training takes about 3 hours per language for training only data, and about 4 hours per language for all data.  \n",
    "About 100 hours for the best single method (all languages).  \n",
    "About 400 hours total (all 4 methods, all 26 languages).  \n",
    "\n",
    "Training scripts will produce the following structure. These model directories are not ready for inference. We need to run `collect_models.py` as shown below.\n",
    "```\n",
    "|-- models-1\n",
    "    |-- ady\n",
    "    |-- aln\n",
    "    |-- ...\n",
    "|-- models-2\n",
    "|-- models-3\n",
    "|-- models-4\n",
    "```\n",
    "\n",
    "**Notes:**\n",
    "1. There are 4 training scripts. Logic and hyperparameters are the same. The main difference is the dataset and learning rate scheduler.\n",
    "   1. `train_trdata_spont.py` Spontaneous corpus, training subset only, maximum cleaning (remove examples if the same audio has several different transcriptions, remove reported examples, etc.), `ReduceLROnPlateau` scheduler.\n",
    "   2. `train_trdata_script.py` Scripted corpus, training subset only, maximum cleaning, `ReduceLROnPlateau` scheduler.\n",
    "   3. `train_alldata_spont.py` Spontaneous corpus, all data (training + validation), no cleaning (use all examples), `MultiStepLR` scheduler.\n",
    "   4. `train_alldata_script.py` Scripted corpus, all data, no cleaning, `MultiStepLR` scheduler.\n",
    "3. In general we found that the best hyperparameters are: learning rate 1e-3 and batch size 2. For some languages there are very long examples which do not allow natural batch size 2 on 24 GB vRAM. In these cases we used batch size 1 and gradient accumulation 2 (effective batch size is still 2).\n",
    "4. Some examples are so long that they do not allow even batch size 1, in this case we used truncation to maximum length of 4_800_000 (5 minutes). With this approach we truncate only audio, ground truth transcription remains the same, creating partial mismatch. Given that the amount of such examples is small (probably only one), there is no significant negative effect on the model, but if you have many such examples in your dataset, you probably need to remove them by length before training.\n",
    "5. In our development stage we have to run `Method 1` set to obtain epoch milestones for other methods. If you need only the best single models per language, you can run only `Method 2` set, because epoch milestones are already known."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "682002db-89e4-433a-82de-58840c826bd4",
   "metadata": {},
   "source": [
    "## Method 1. Training subset only, `facebook/mms-1b-l1107`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d637ac46-ecf0-4b19-8597-e49be028e9c7",
   "metadata": {},
   "outputs": [],
   "source": [
    "!python train_trdata_spont.py --output_dir=models-1 --model_name=facebook/mms-1b-l1107 --lang=aln\n",
    "!python train_trdata_spont.py --output_dir=models-1 --model_name=facebook/mms-1b-l1107 --lang=bew\n",
    "!python train_trdata_spont.py --output_dir=models-1 --model_name=facebook/mms-1b-l1107 --lang=bxk\n",
    "!python train_trdata_spont.py --output_dir=models-1 --model_name=facebook/mms-1b-l1107 --lang=cgg\n",
    "!python train_trdata_spont.py --output_dir=models-1 --model_name=facebook/mms-1b-l1107 --lang=el-CY --batch_size=1 --accum=2\n",
    "!python train_trdata_spont.py --output_dir=models-1 --model_name=facebook/mms-1b-l1107 --lang=hch\n",
    "!python train_trdata_spont.py --output_dir=models-1 --model_name=facebook/mms-1b-l1107 --lang=kcn\n",
    "!python train_trdata_spont.py --output_dir=models-1 --model_name=facebook/mms-1b-l1107 --lang=koo\n",
    "!python train_trdata_spont.py --output_dir=models-1 --model_name=facebook/mms-1b-l1107 --lang=led\n",
    "!python train_trdata_spont.py --output_dir=models-1 --model_name=facebook/mms-1b-l1107 --lang=lke\n",
    "!python train_trdata_spont.py --output_dir=models-1 --model_name=facebook/mms-1b-l1107 --lang=lth --batch_size=1 --accum=2 --max_length=4_800_000 --truncation=1\n",
    "!python train_trdata_spont.py --output_dir=models-1 --model_name=facebook/mms-1b-l1107 --lang=meh\n",
    "!python train_trdata_spont.py --output_dir=models-1 --model_name=facebook/mms-1b-l1107 --lang=mmc --batch_size=1 --accum=2\n",
    "!python train_trdata_spont.py --output_dir=models-1 --model_name=facebook/mms-1b-l1107 --lang=pne\n",
    "!python train_trdata_spont.py --output_dir=models-1 --model_name=facebook/mms-1b-l1107 --lang=ruc\n",
    "!python train_trdata_spont.py --output_dir=models-1 --model_name=facebook/mms-1b-l1107 --lang=rwm\n",
    "!python train_trdata_spont.py --output_dir=models-1 --model_name=facebook/mms-1b-l1107 --lang=sco --batch_size=1 --accum=2\n",
    "!python train_trdata_spont.py --output_dir=models-1 --model_name=facebook/mms-1b-l1107 --lang=tob\n",
    "!python train_trdata_spont.py --output_dir=models-1 --model_name=facebook/mms-1b-l1107 --lang=top --batch_size=1 --accum=2\n",
    "!python train_trdata_spont.py --output_dir=models-1 --model_name=facebook/mms-1b-l1107 --lang=ttj\n",
    "!python train_trdata_spont.py --output_dir=models-1 --model_name=facebook/mms-1b-l1107 --lang=ukv\n",
    "\n",
    "!python train_trdata_script.py  --output_dir=models-1 --model_name=facebook/mms-1b-l1107 --lang=ady\n",
    "!python train_trdata_script.py  --output_dir=models-1 --model_name=facebook/mms-1b-l1107 --lang=bas\n",
    "!python train_trdata_script.py  --output_dir=models-1 --model_name=facebook/mms-1b-l1107 --lang=kbd\n",
    "!python train_trdata_script.py  --output_dir=models-1 --model_name=facebook/mms-1b-l1107 --lang=qxp\n",
    "!python train_trdata_script.py  --output_dir=models-1 --model_name=facebook/mms-1b-l1107 --lang=ush"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "26afd416-dbaa-49bb-807f-45e8b18a1249",
   "metadata": {},
   "source": [
    "## Method 2. All data, `facebook/mms-1b-l1107`. Best single model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24b36355-2e59-43b1-b2fb-bf6d8da6e2a3",
   "metadata": {},
   "outputs": [],
   "source": [
    "!python train_alldata_spont.py --output_dir=models-2 --model_name=facebook/mms-1b-l1107 --lang=aln --epoch_milestones=\"[16, 23]\"\n",
    "!python train_alldata_spont.py --output_dir=models-2 --model_name=facebook/mms-1b-l1107 --lang=bew --epoch_milestones=\"[19, 23, 27]\" \n",
    "!python train_alldata_spont.py --output_dir=models-2 --model_name=facebook/mms-1b-l1107 --lang=bxk --epoch_milestones=\"[12, 28]\" \n",
    "!python train_alldata_spont.py --output_dir=models-2 --model_name=facebook/mms-1b-l1107 --lang=cgg --epoch_milestones=\"[16, 28]\" \n",
    "!python train_alldata_spont.py --output_dir=models-2 --model_name=facebook/mms-1b-l1107 --lang=el-CY --epoch_milestones=\"[15, 24, 27]\" --batch_size=1 --accum=2\n",
    "!python train_alldata_spont.py --output_dir=models-2 --model_name=facebook/mms-1b-l1107 --lang=hch --epoch_milestones=\"[10, 14, 22, 29]\" \n",
    "!python train_alldata_spont.py --output_dir=models-2 --model_name=facebook/mms-1b-l1107 --lang=kcn --epoch_milestones=\"[24, 29]\" \n",
    "!python train_alldata_spont.py --output_dir=models-2 --model_name=facebook/mms-1b-l1107 --lang=koo --epoch_milestones=\"[13, 21, 26, 29]\" \n",
    "!python train_alldata_spont.py --output_dir=models-2 --model_name=facebook/mms-1b-l1107 --lang=led --epoch_milestones=\"[18, 25]\" \n",
    "!python train_alldata_spont.py --output_dir=models-2 --model_name=facebook/mms-1b-l1107 --lang=lke --epoch_milestones=\"[12, 17, 22]\" \n",
    "!python train_alldata_spont.py --output_dir=models-2 --model_name=facebook/mms-1b-l1107 --lang=lth --epoch_milestones=\"[19, 23, 26]\" --batch_size=1 --accum=2 --max_length=4_800_000 --truncation=1\n",
    "!python train_alldata_spont.py --output_dir=models-2 --model_name=facebook/mms-1b-l1107 --lang=meh --epoch_milestones=\"[ 8, 15, 19, 22]\" \n",
    "!python train_alldata_spont.py --output_dir=models-2 --model_name=facebook/mms-1b-l1107 --lang=mmc --epoch_milestones=\"[13, 20, 23, 28]\" --batch_size=1 --accum=2\n",
    "!python train_alldata_spont.py --output_dir=models-2 --model_name=facebook/mms-1b-l1107 --lang=pne --epoch_milestones=\"[27]\" \n",
    "!python train_alldata_spont.py --output_dir=models-2 --model_name=facebook/mms-1b-l1107 --lang=ruc --epoch_milestones=\"[18, 25]\" \n",
    "!python train_alldata_spont.py --output_dir=models-2 --model_name=facebook/mms-1b-l1107 --lang=rwm --epoch_milestones=\"[17, 27]\" \n",
    "!python train_alldata_spont.py --output_dir=models-2 --model_name=facebook/mms-1b-l1107 --lang=sco --epoch_milestones=\"[29]\" --batch_size=1 --accum=2\n",
    "!python train_alldata_spont.py --output_dir=models-2 --model_name=facebook/mms-1b-l1107 --lang=tob --epoch_milestones=\"[ 5, 8, 11, 14, 17, 20, 23, 26, 29]\" \n",
    "!python train_alldata_spont.py --output_dir=models-2 --model_name=facebook/mms-1b-l1107 --lang=top --epoch_milestones=\"[ 6, 11, 18, 21, 24, 27]\" --batch_size=1 --accum=2\n",
    "!python train_alldata_spont.py --output_dir=models-2 --model_name=facebook/mms-1b-l1107 --lang=ttj --epoch_milestones=\"[25]\" \n",
    "!python train_alldata_spont.py --output_dir=models-2 --model_name=facebook/mms-1b-l1107 --lang=ukv --epoch_milestones=\"[ 8, 22, 25, 28]\" \n",
    "\n",
    "!python train_alldata_script.py --output_dir=models-2 --model_name=facebook/mms-1b-l1107 --lang=ady --epoch_milestones=\"[16, 26]\" \n",
    "!python train_alldata_script.py --output_dir=models-2 --model_name=facebook/mms-1b-l1107 --lang=bas --epoch_milestones=\"[19, 28]\" \n",
    "!python train_alldata_script.py --output_dir=models-2 --model_name=facebook/mms-1b-l1107 --lang=kbd --epoch_milestones=\"[11, 18, 25, 28]\" \n",
    "!python train_alldata_script.py --output_dir=models-2 --model_name=facebook/mms-1b-l1107 --lang=qxp --epoch_milestones=\"[ 7, 11, 21]\" \n",
    "!python train_alldata_script.py --output_dir=models-2 --model_name=facebook/mms-1b-l1107 --lang=ush --epoch_milestones=\"[10, 14, 17, 20, 23, 26, 29]\" "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f2236c29-a164-4a8a-ad98-b65da5c84f1d",
   "metadata": {},
   "source": [
    "## Method 3. All data, `facebook/mms-1b-all`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "63dda148-4efd-4673-92c5-721b22686ce6",
   "metadata": {},
   "outputs": [],
   "source": [
    "!python train_alldata_spont.py --output_dir=models-3 --model_name=facebook/mms-1b-all --lang=aln --epoch_milestones=\"[16, 23]\"\n",
    "!python train_alldata_spont.py --output_dir=models-3 --model_name=facebook/mms-1b-all --lang=bew --epoch_milestones=\"[19, 23, 27]\" \n",
    "!python train_alldata_spont.py --output_dir=models-3 --model_name=facebook/mms-1b-all --lang=bxk --epoch_milestones=\"[12, 28]\" \n",
    "!python train_alldata_spont.py --output_dir=models-3 --model_name=facebook/mms-1b-all --lang=cgg --epoch_milestones=\"[16, 28]\" \n",
    "!python train_alldata_spont.py --output_dir=models-3 --model_name=facebook/mms-1b-all --lang=el-CY --epoch_milestones=\"[15, 24, 27]\" --batch_size=1 --accum=2\n",
    "!python train_alldata_spont.py --output_dir=models-3 --model_name=facebook/mms-1b-all --lang=hch --epoch_milestones=\"[10, 14, 22, 29]\" \n",
    "!python train_alldata_spont.py --output_dir=models-3 --model_name=facebook/mms-1b-all --lang=kcn --epoch_milestones=\"[24, 29]\" \n",
    "!python train_alldata_spont.py --output_dir=models-3 --model_name=facebook/mms-1b-all --lang=koo --epoch_milestones=\"[13, 21, 26, 29]\" \n",
    "!python train_alldata_spont.py --output_dir=models-3 --model_name=facebook/mms-1b-all --lang=led --epoch_milestones=\"[18, 25]\" \n",
    "!python train_alldata_spont.py --output_dir=models-3 --model_name=facebook/mms-1b-all --lang=lke --epoch_milestones=\"[12, 17, 22]\" \n",
    "!python train_alldata_spont.py --output_dir=models-3 --model_name=facebook/mms-1b-all --lang=lth --epoch_milestones=\"[19, 23, 26]\" --batch_size=1 --accum=2 --max_length=4_800_000 --truncation=1\n",
    "!python train_alldata_spont.py --output_dir=models-3 --model_name=facebook/mms-1b-all --lang=meh --epoch_milestones=\"[ 8, 15, 19, 22]\" \n",
    "!python train_alldata_spont.py --output_dir=models-3 --model_name=facebook/mms-1b-all --lang=mmc --epoch_milestones=\"[13, 20, 23, 28]\" --batch_size=1 --accum=2\n",
    "!python train_alldata_spont.py --output_dir=models-3 --model_name=facebook/mms-1b-all --lang=pne --epoch_milestones=\"[27]\" \n",
    "!python train_alldata_spont.py --output_dir=models-3 --model_name=facebook/mms-1b-all --lang=ruc --epoch_milestones=\"[18, 25]\" \n",
    "!python train_alldata_spont.py --output_dir=models-3 --model_name=facebook/mms-1b-all --lang=rwm --epoch_milestones=\"[17, 27]\" \n",
    "!python train_alldata_spont.py --output_dir=models-3 --model_name=facebook/mms-1b-all --lang=sco --epoch_milestones=\"[29]\" --batch_size=1 --accum=2\n",
    "!python train_alldata_spont.py --output_dir=models-3 --model_name=facebook/mms-1b-all --lang=tob --epoch_milestones=\"[ 5, 8, 11, 14, 17, 20, 23, 26, 29]\" \n",
    "!python train_alldata_spont.py --output_dir=models-3 --model_name=facebook/mms-1b-all --lang=top --epoch_milestones=\"[ 6, 11, 18, 21, 24, 27]\" --batch_size=1 --accum=2\n",
    "!python train_alldata_spont.py --output_dir=models-3 --model_name=facebook/mms-1b-all --lang=ttj --epoch_milestones=\"[25]\" \n",
    "!python train_alldata_spont.py --output_dir=models-3 --model_name=facebook/mms-1b-all --lang=ukv --epoch_milestones=\"[ 8, 22, 25, 28]\" \n",
    "\n",
    "!python train_alldata_script.py --output_dir=models-3 --model_name=facebook/mms-1b-all --lang=ady --epoch_milestones=\"[16, 26]\" \n",
    "!python train_alldata_script.py --output_dir=models-3 --model_name=facebook/mms-1b-all --lang=bas --epoch_milestones=\"[19, 28]\" \n",
    "!python train_alldata_script.py --output_dir=models-3 --model_name=facebook/mms-1b-all --lang=kbd --epoch_milestones=\"[11, 18, 25, 28]\" \n",
    "!python train_alldata_script.py --output_dir=models-3 --model_name=facebook/mms-1b-all --lang=qxp --epoch_milestones=\"[ 7, 11, 21]\" \n",
    "!python train_alldata_script.py --output_dir=models-3 --model_name=facebook/mms-1b-all --lang=ush --epoch_milestones=\"[10, 14, 17, 20, 23, 26, 29]\" "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b299aeaf-895a-419a-a0a0-ea00c2e344f8",
   "metadata": {},
   "source": [
    "## Method 4. All data, `facebook/mms-1b-fl102`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d032a9f0-eb50-43c9-9db6-35d09c2d5c01",
   "metadata": {},
   "outputs": [],
   "source": [
    "!python train_alldata_spont.py --output_dir=models-4 --model_name=facebook/mms-1b-fl102 --lang=aln --epoch_milestones=\"[16, 23]\"\n",
    "!python train_alldata_spont.py --output_dir=models-4 --model_name=facebook/mms-1b-fl102 --lang=bew --epoch_milestones=\"[19, 23, 27]\" \n",
    "!python train_alldata_spont.py --output_dir=models-4 --model_name=facebook/mms-1b-fl102 --lang=bxk --epoch_milestones=\"[12, 28]\" \n",
    "!python train_alldata_spont.py --output_dir=models-4 --model_name=facebook/mms-1b-fl102 --lang=cgg --epoch_milestones=\"[16, 28]\" \n",
    "!python train_alldata_spont.py --output_dir=models-4 --model_name=facebook/mms-1b-fl102 --lang=el-CY --epoch_milestones=\"[15, 24, 27]\" --batch_size=1 --accum=2\n",
    "!python train_alldata_spont.py --output_dir=models-4 --model_name=facebook/mms-1b-fl102 --lang=hch --epoch_milestones=\"[10, 14, 22, 29]\" \n",
    "!python train_alldata_spont.py --output_dir=models-4 --model_name=facebook/mms-1b-fl102 --lang=kcn --epoch_milestones=\"[24, 29]\" \n",
    "!python train_alldata_spont.py --output_dir=models-4 --model_name=facebook/mms-1b-fl102 --lang=koo --epoch_milestones=\"[13, 21, 26, 29]\" \n",
    "!python train_alldata_spont.py --output_dir=models-4 --model_name=facebook/mms-1b-fl102 --lang=led --epoch_milestones=\"[18, 25]\" \n",
    "!python train_alldata_spont.py --output_dir=models-4 --model_name=facebook/mms-1b-fl102 --lang=lke --epoch_milestones=\"[12, 17, 22]\" \n",
    "!python train_alldata_spont.py --output_dir=models-4 --model_name=facebook/mms-1b-fl102 --lang=lth --epoch_milestones=\"[19, 23, 26]\" --batch_size=1 --accum=2 --max_length=4_800_000 --truncation=1\n",
    "!python train_alldata_spont.py --output_dir=models-4 --model_name=facebook/mms-1b-fl102 --lang=meh --epoch_milestones=\"[ 8, 15, 19, 22]\" \n",
    "!python train_alldata_spont.py --output_dir=models-4 --model_name=facebook/mms-1b-fl102 --lang=mmc --epoch_milestones=\"[13, 20, 23, 28]\" --batch_size=1 --accum=2\n",
    "!python train_alldata_spont.py --output_dir=models-4 --model_name=facebook/mms-1b-fl102 --lang=pne --epoch_milestones=\"[27]\" \n",
    "!python train_alldata_spont.py --output_dir=models-4 --model_name=facebook/mms-1b-fl102 --lang=ruc --epoch_milestones=\"[18, 25]\" \n",
    "!python train_alldata_spont.py --output_dir=models-4 --model_name=facebook/mms-1b-fl102 --lang=rwm --epoch_milestones=\"[17, 27]\" \n",
    "!python train_alldata_spont.py --output_dir=models-4 --model_name=facebook/mms-1b-fl102 --lang=sco --epoch_milestones=\"[29]\" --batch_size=1 --accum=2\n",
    "!python train_alldata_spont.py --output_dir=models-4 --model_name=facebook/mms-1b-fl102 --lang=tob --epoch_milestones=\"[ 5, 8, 11, 14, 17, 20, 23, 26, 29]\" \n",
    "!python train_alldata_spont.py --output_dir=models-4 --model_name=facebook/mms-1b-fl102 --lang=top --epoch_milestones=\"[ 6, 11, 18, 21, 24, 27]\" --batch_size=1 --accum=2\n",
    "!python train_alldata_spont.py --output_dir=models-4 --model_name=facebook/mms-1b-fl102 --lang=ttj --epoch_milestones=\"[25]\" \n",
    "!python train_alldata_spont.py --output_dir=models-4 --model_name=facebook/mms-1b-fl102 --lang=ukv --epoch_milestones=\"[ 8, 22, 25, 28]\" \n",
    "\n",
    "!python train_alldata_script.py --output_dir=models-4 --model_name=facebook/mms-1b-fl102 --lang=ady --epoch_milestones=\"[16, 26]\" \n",
    "!python train_alldata_script.py --output_dir=models-4 --model_name=facebook/mms-1b-fl102 --lang=bas --epoch_milestones=\"[19, 28]\" \n",
    "!python train_alldata_script.py --output_dir=models-4 --model_name=facebook/mms-1b-fl102 --lang=kbd --epoch_milestones=\"[11, 18, 25, 28]\" \n",
    "!python train_alldata_script.py --output_dir=models-4 --model_name=facebook/mms-1b-fl102 --lang=qxp --epoch_milestones=\"[ 7, 11, 21]\" \n",
    "!python train_alldata_script.py --output_dir=models-4 --model_name=facebook/mms-1b-fl102 --lang=ush --epoch_milestones=\"[10, 14, 17, 20, 23, 26, 29]\" "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b3739d15-d566-4eaa-9336-e9e8305e5a90",
   "metadata": {},
   "source": [
    "## Collect adapters and vocabs, and create final models\n",
    "\n",
    "Script `collect_models.py` will create the following directories, corresponding to the ones used in the `infer.py` script.  \n",
    "\n",
    "```\n",
    "|-- models-1-final\n",
    "|-- models-2-final\n",
    "|-- models-3-final\n",
    "|-- models-4-final\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e275aca9-d10c-4058-9d19-d61f64970e1a",
   "metadata": {},
   "outputs": [],
   "source": [
    "!python collect_models.py \\\n",
    "--input_dir=./ \\\n",
    "--output_dir=./"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6794cc3e-8566-4204-8407-98b95a81fa20",
   "metadata": {},
   "source": [
    "# 5. Quantization\n",
    "\n",
    "Script `prune_quantize.py` will create the `models-5-final` directory, corresponding to the one used in the `infer.py` script.  \n",
    "\n",
    "**Note.** For quantization we need a dedicated model per language (instead of one model + 26 adapters), because the adapter mechanism does not work for quantized models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1e422f8c-651e-4b7d-9ef3-99b33c71c64e",
   "metadata": {},
   "outputs": [],
   "source": [
    "!python prune_quantize.py \\\n",
    "--input_dir=./models-2 \\\n",
    "--output_dir=./models-5-final \\\n",
    "--n_top_layers_to_remove=3 \\\n",
    "--bnb_4bit_quant_type=nf4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "f3efa9c1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# END"
   ]
  }
 ],
 "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.12.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}