From 7dba51ad3b19f05047710b8ba3c36453d53bc6ea Mon Sep 17 00:00:00 2001
From: Shikhar Bharadwaj <sbharad2@andrew.cmu.edu>
Date: Mon, 23 Sep 2024 09:30:29 -0400
Subject: [PATCH 01/47] add clotho_v2 aac recipe with beats encoder

---
 .gitignore                                    |    1 +
 egs2/TEMPLATE/asr1/db.sh                      |    1 +
 .../Untitled-checkpoint.ipynb                 |    6 +
 egs2/clotho_v2/asr1/asr.sh                    |    1 +
 egs2/clotho_v2/asr1/cmd.sh                    |  110 ++
 egs2/clotho_v2/asr1/conf/branchformer.yaml    |   65 +
 egs2/clotho_v2/asr1/conf/dcase.yaml           |   60 +
 egs2/clotho_v2/asr1/conf/fbank.conf           |    2 +
 egs2/clotho_v2/asr1/conf/pbs.conf             |   11 +
 egs2/clotho_v2/asr1/conf/pitch.conf           |    1 +
 egs2/clotho_v2/asr1/conf/queue.conf           |   12 +
 egs2/clotho_v2/asr1/conf/slurm.conf           |   14 +
 egs2/clotho_v2/asr1/db.sh                     |    1 +
 .../Checker-checkpoint.ipynb                  |   91 ++
 .../.ipynb_checkpoints/data-checkpoint.sh     |   61 +
 .../data_prep-checkpoint.py                   |   49 +
 egs2/clotho_v2/asr1/local/Checker.ipynb       |   91 ++
 egs2/clotho_v2/asr1/local/data.sh             |   75 +
 egs2/clotho_v2/asr1/local/data_prep.py        |   80 +
 egs2/clotho_v2/asr1/local/evaluation.py       |   91 ++
 egs2/clotho_v2/asr1/local/path.sh             |    0
 egs2/clotho_v2/asr1/local/score.sh            |   29 +
 egs2/clotho_v2/asr1/path.sh                   |    1 +
 egs2/clotho_v2/asr1/pyscripts                 |    1 +
 egs2/clotho_v2/asr1/run.sh                    |   31 +
 egs2/clotho_v2/asr1/scripts                   |    1 +
 egs2/clotho_v2/asr1/steps                     |    1 +
 egs2/clotho_v2/asr1/utils                     |    1 +
 espnet2/asr/encoder/beats_encoder.py          | 1310 +++++++++++++++++
 espnet2/tasks/asr.py                          |    4 +-
 30 files changed, 2201 insertions(+), 1 deletion(-)
 create mode 100644 egs2/clotho_v2/.ipynb_checkpoints/Untitled-checkpoint.ipynb
 create mode 120000 egs2/clotho_v2/asr1/asr.sh
 create mode 100644 egs2/clotho_v2/asr1/cmd.sh
 create mode 100644 egs2/clotho_v2/asr1/conf/branchformer.yaml
 create mode 100644 egs2/clotho_v2/asr1/conf/dcase.yaml
 create mode 100644 egs2/clotho_v2/asr1/conf/fbank.conf
 create mode 100644 egs2/clotho_v2/asr1/conf/pbs.conf
 create mode 100644 egs2/clotho_v2/asr1/conf/pitch.conf
 create mode 100644 egs2/clotho_v2/asr1/conf/queue.conf
 create mode 100644 egs2/clotho_v2/asr1/conf/slurm.conf
 create mode 120000 egs2/clotho_v2/asr1/db.sh
 create mode 100644 egs2/clotho_v2/asr1/local/.ipynb_checkpoints/Checker-checkpoint.ipynb
 create mode 100755 egs2/clotho_v2/asr1/local/.ipynb_checkpoints/data-checkpoint.sh
 create mode 100644 egs2/clotho_v2/asr1/local/.ipynb_checkpoints/data_prep-checkpoint.py
 create mode 100644 egs2/clotho_v2/asr1/local/Checker.ipynb
 create mode 100755 egs2/clotho_v2/asr1/local/data.sh
 create mode 100644 egs2/clotho_v2/asr1/local/data_prep.py
 create mode 100644 egs2/clotho_v2/asr1/local/evaluation.py
 create mode 100644 egs2/clotho_v2/asr1/local/path.sh
 create mode 100755 egs2/clotho_v2/asr1/local/score.sh
 create mode 120000 egs2/clotho_v2/asr1/path.sh
 create mode 120000 egs2/clotho_v2/asr1/pyscripts
 create mode 100755 egs2/clotho_v2/asr1/run.sh
 create mode 120000 egs2/clotho_v2/asr1/scripts
 create mode 120000 egs2/clotho_v2/asr1/steps
 create mode 120000 egs2/clotho_v2/asr1/utils
 create mode 100644 espnet2/asr/encoder/beats_encoder.py

diff --git a/.gitignore b/.gitignore
index ba5301179ca..ed37f59e70c 100644
--- a/.gitignore
+++ b/.gitignore
@@ -103,3 +103,4 @@ tools/ffmpeg-*-static
 tools/ffmpeg-release
 tools/ffmpeg-release-*.tar.xz
 tools/Miniconda3-latest-*.sh
+egs2/clotho_v2/asr1/local/CLOTHO_v2.1/*
diff --git a/egs2/TEMPLATE/asr1/db.sh b/egs2/TEMPLATE/asr1/db.sh
index 6ad19ebc2c5..e94e95f2733 100755
--- a/egs2/TEMPLATE/asr1/db.sh
+++ b/egs2/TEMPLATE/asr1/db.sh
@@ -219,6 +219,7 @@ SPRING_INX=downloads
 VOXCELEB=
 KSPONSPEECH=
 HIFITTS=downloads
+CLOTHO_V2=/data/user_data/sbharad2/espnet/egs2/clotho_v2/asr1/local/CLOTHO_v2.1 # TODO(sbharad2): Downloads
 
 # For only CMU TIR environment
 if [[ "$(hostname)" == tir* ]]; then
diff --git a/egs2/clotho_v2/.ipynb_checkpoints/Untitled-checkpoint.ipynb b/egs2/clotho_v2/.ipynb_checkpoints/Untitled-checkpoint.ipynb
new file mode 100644
index 00000000000..363fcab7ed6
--- /dev/null
+++ b/egs2/clotho_v2/.ipynb_checkpoints/Untitled-checkpoint.ipynb
@@ -0,0 +1,6 @@
+{
+ "cells": [],
+ "metadata": {},
+ "nbformat": 4,
+ "nbformat_minor": 5
+}
diff --git a/egs2/clotho_v2/asr1/asr.sh b/egs2/clotho_v2/asr1/asr.sh
new file mode 120000
index 00000000000..60b05122cfd
--- /dev/null
+++ b/egs2/clotho_v2/asr1/asr.sh
@@ -0,0 +1 @@
+../../TEMPLATE/asr1/asr.sh
\ No newline at end of file
diff --git a/egs2/clotho_v2/asr1/cmd.sh b/egs2/clotho_v2/asr1/cmd.sh
new file mode 100644
index 00000000000..2aae6919fef
--- /dev/null
+++ b/egs2/clotho_v2/asr1/cmd.sh
@@ -0,0 +1,110 @@
+# ====== About run.pl, queue.pl, slurm.pl, and ssh.pl ======
+# Usage: <cmd>.pl [options] JOB=1:<nj> <log> <command...>
+# e.g.
+#   run.pl --mem 4G JOB=1:10 echo.JOB.log echo JOB
+#
+# Options:
+#   --time <time>: Limit the maximum time to execute.
+#   --mem <mem>: Limit the maximum memory usage.
+#   -–max-jobs-run <njob>: Limit the number parallel jobs. This is ignored for non-array jobs.
+#   --num-threads <ngpu>: Specify the number of CPU core.
+#   --gpu <ngpu>: Specify the number of GPU devices.
+#   --config: Change the configuration file from default.
+#
+# "JOB=1:10" is used for "array jobs" and it can control the number of parallel jobs.
+# The left string of "=", i.e. "JOB", is replaced by <N>(Nth job) in the command and the log file name,
+# e.g. "echo JOB" is changed to "echo 3" for the 3rd job and "echo 8" for 8th job respectively.
+# Note that the number must start with a positive number, so you can't use "JOB=0:10" for example.
+#
+# run.pl, queue.pl, slurm.pl, and ssh.pl have unified interface, not depending on its backend.
+# These options are mapping to specific options for each backend and
+# it is configured by "conf/queue.conf" and "conf/slurm.conf" by default.
+# If jobs failed, your configuration might be wrong for your environment.
+#
+#
+# The official documentation for run.pl, queue.pl, slurm.pl, and ssh.pl:
+#   "Parallelization in Kaldi": http://kaldi-asr.org/doc/queue.html
+# =========================================================~
+
+
+# Select the backend used by run.sh from "local", "stdout", "sge", "slurm", or "ssh"
+cmd_backend='local'
+
+# Local machine, without any Job scheduling system
+if [ "${cmd_backend}" = local ]; then
+
+    # The other usage
+    export train_cmd="run.pl"
+    # Used for "*_train.py": "--gpu" is appended optionally by run.sh
+    export cuda_cmd="run.pl"
+    # Used for "*_recog.py"
+    export decode_cmd="run.pl"
+
+# Local machine logging to stdout and log file, without any Job scheduling system
+elif [ "${cmd_backend}" = stdout ]; then
+
+    # The other usage
+    export train_cmd="stdout.pl"
+    # Used for "*_train.py": "--gpu" is appended optionally by run.sh
+    export cuda_cmd="stdout.pl"
+    # Used for "*_recog.py"
+    export decode_cmd="stdout.pl"
+
+
+# "qsub" (Sun Grid Engine, or derivation of it)
+elif [ "${cmd_backend}" = sge ]; then
+    # The default setting is written in conf/queue.conf.
+    # You must change "-q g.q" for the "queue" for your environment.
+    # To know the "queue" names, type "qhost -q"
+    # Note that to use "--gpu *", you have to setup "complex_value" for the system scheduler.
+
+    export train_cmd="queue.pl"
+    export cuda_cmd="queue.pl"
+    export decode_cmd="queue.pl"
+
+
+# "qsub" (Torque/PBS.)
+elif [ "${cmd_backend}" = pbs ]; then
+    # The default setting is written in conf/pbs.conf.
+
+    export train_cmd="pbs.pl"
+    export cuda_cmd="pbs.pl"
+    export decode_cmd="pbs.pl"
+
+
+# "sbatch" (Slurm)
+elif [ "${cmd_backend}" = slurm ]; then
+    # The default setting is written in conf/slurm.conf.
+    # You must change "-p cpu" and "-p gpu" for the "partition" for your environment.
+    # To know the "partion" names, type "sinfo".
+    # You can use "--gpu * " by default for slurm and it is interpreted as "--gres gpu:*"
+    # The devices are allocated exclusively using "${CUDA_VISIBLE_DEVICES}".
+
+    export train_cmd="slurm.pl"
+    export cuda_cmd="slurm.pl"
+    export decode_cmd="slurm.pl"
+
+elif [ "${cmd_backend}" = ssh ]; then
+    # You have to create ".queue/machines" to specify the host to execute jobs.
+    # e.g. .queue/machines
+    #   host1
+    #   host2
+    #   host3
+    # Assuming you can login them without any password, i.e. You have to set ssh keys.
+
+    export train_cmd="ssh.pl"
+    export cuda_cmd="ssh.pl"
+    export decode_cmd="ssh.pl"
+
+# This is an example of specifying several unique options in the JHU CLSP cluster setup.
+# Users can modify/add their own command options according to their cluster environments.
+elif [ "${cmd_backend}" = jhu ]; then
+
+    export train_cmd="queue.pl --mem 2G"
+    export cuda_cmd="queue-freegpu.pl --mem 2G --gpu 1 --config conf/queue.conf"
+    export decode_cmd="queue.pl --mem 4G"
+
+else
+    echo "$0: Error: Unknown cmd_backend=${cmd_backend}" 1>&2
+    return 1
+fi
diff --git a/egs2/clotho_v2/asr1/conf/branchformer.yaml b/egs2/clotho_v2/asr1/conf/branchformer.yaml
new file mode 100644
index 00000000000..b092371029a
--- /dev/null
+++ b/egs2/clotho_v2/asr1/conf/branchformer.yaml
@@ -0,0 +1,65 @@
+batch_type: folded
+batch_size: 64
+accum_grad: 2    # gradient accumulation steps
+max_epoch: 100
+patience: none
+init: xavier_uniform
+best_model_criterion:  # criterion to save best models
+-   - valid
+    - acc
+    - max
+keep_nbest_models: 10  # save nbest models and average these checkpoints
+use_amp: true    # whether to use automatic mixed precision
+num_att_plot: 0  # do not save attention plots to save time in the demo
+num_workers: 2   # number of workers in dataloader
+
+encoder: transformer
+encoder_conf:
+    output_size: 256
+    attention_heads: 4
+    linear_units: 1024
+    num_blocks: 12
+    dropout_rate: 0.1
+    positional_dropout_rate: 0.1
+    attention_dropout_rate: 0.1
+    input_layer: conv2d
+    normalize_before: true
+    # use_flash_attn: false
+
+decoder: transformer
+decoder_conf:
+    attention_heads: 4
+    linear_units: 1024
+    num_blocks: 6
+    dropout_rate: 0.1
+    positional_dropout_rate: 0.1
+    self_attention_dropout_rate: 0.1
+    src_attention_dropout_rate: 0.1
+    # use_flash_attn: false
+
+model_conf:
+    ctc_weight: 0.3  # joint CTC/attention training
+    lsm_weight: 0.1  # label smoothing weight
+    length_normalized_loss: false
+
+optim: adam
+optim_conf:
+    lr: 0.001
+scheduler: warmuplr  # linearly increase and exponentially decrease
+scheduler_conf:
+    warmup_steps: 800
+specaug: specaug
+specaug_conf:
+    apply_time_warp: true
+    time_warp_window: 5
+    time_warp_mode: bicubic
+    apply_freq_mask: true
+    freq_mask_width_range:
+    - 0
+    - 30
+    num_freq_mask: 2
+    apply_time_mask: true
+    time_mask_width_range:
+    - 0
+    - 40
+    num_time_mask: 2
diff --git a/egs2/clotho_v2/asr1/conf/dcase.yaml b/egs2/clotho_v2/asr1/conf/dcase.yaml
new file mode 100644
index 00000000000..ee5c7944bed
--- /dev/null
+++ b/egs2/clotho_v2/asr1/conf/dcase.yaml
@@ -0,0 +1,60 @@
+batch_type: folded
+batch_size: 80
+accum_grad: 1    # gradient accumulation steps
+max_epoch: 20
+patience: none
+best_model_criterion:  # criterion to save best models
+-   - valid
+    - acc
+    - max
+keep_nbest_models: 5  # save nbest models and average these checkpoints
+use_amp: true    # whether to use automatic mixed precision
+num_att_plot: 0  # do not save attention plots to save time in the demo
+num_workers: 2   # number of workers in dataloader
+
+input_size: 1
+# beats code takes care of generating mel spectrogram and specaug
+frontend: none
+normalize: none
+init: xavier_normal # For the decoder
+
+encoder: beats
+encoder_conf:
+    beats_ckpt_path: /home/sbharad2/dcase2024_cp/dcase24task6/pretrained_weights/BEATs_iter3_plus_AS2M_finetuned_on_AS2M_cpt1.pt
+    beats_config:
+        input_patch_size: 16
+        embed_dim: 512
+    specaug_config:
+        apply_freq_mask: true
+        freq_mask_width_range:
+        - 0
+        - 64
+        num_freq_mask: 2
+        apply_time_mask: true
+        time_mask_width_ratio_range:
+        - 0
+        - 0.12
+        num_time_mask: 5
+
+decoder: transformer
+decoder_conf:
+    attention_heads: 4
+    linear_units: 1024
+    num_blocks: 6
+    dropout_rate: 0.1
+    positional_dropout_rate: 0.1
+    self_attention_dropout_rate: 0.1
+    src_attention_dropout_rate: 0.1
+    # use_flash_attn: false
+
+model_conf:
+    ctc_weight: 0.3  # joint CTC/attention training
+    lsm_weight: 0.1  # label smoothing weight
+    length_normalized_loss: false
+
+optim: adam
+optim_conf:
+    lr: 0.001
+scheduler: warmuplr  # linearly increase and exponentially decrease
+scheduler_conf:
+    warmup_steps: 500
\ No newline at end of file
diff --git a/egs2/clotho_v2/asr1/conf/fbank.conf b/egs2/clotho_v2/asr1/conf/fbank.conf
new file mode 100644
index 00000000000..75232358639
--- /dev/null
+++ b/egs2/clotho_v2/asr1/conf/fbank.conf
@@ -0,0 +1,2 @@
+--sample-frequency=16000
+--num-mel-bins=80
diff --git a/egs2/clotho_v2/asr1/conf/pbs.conf b/egs2/clotho_v2/asr1/conf/pbs.conf
new file mode 100644
index 00000000000..119509938ce
--- /dev/null
+++ b/egs2/clotho_v2/asr1/conf/pbs.conf
@@ -0,0 +1,11 @@
+# Default configuration
+command qsub -V -v PATH -S /bin/bash
+option name=* -N $0
+option mem=* -l mem=$0
+option mem=0          # Do not add anything to qsub_opts
+option num_threads=* -l ncpus=$0
+option num_threads=1  # Do not add anything to qsub_opts
+option num_nodes=* -l nodes=$0:ppn=1
+default gpu=0
+option gpu=0
+option gpu=* -l ngpus=$0
diff --git a/egs2/clotho_v2/asr1/conf/pitch.conf b/egs2/clotho_v2/asr1/conf/pitch.conf
new file mode 100644
index 00000000000..e959a19d5b8
--- /dev/null
+++ b/egs2/clotho_v2/asr1/conf/pitch.conf
@@ -0,0 +1 @@
+--sample-frequency=16000
diff --git a/egs2/clotho_v2/asr1/conf/queue.conf b/egs2/clotho_v2/asr1/conf/queue.conf
new file mode 100644
index 00000000000..500582fab31
--- /dev/null
+++ b/egs2/clotho_v2/asr1/conf/queue.conf
@@ -0,0 +1,12 @@
+# Default configuration
+command qsub -v PATH -cwd -S /bin/bash -j y -l arch=*64*
+option name=* -N $0
+option mem=* -l mem_free=$0,ram_free=$0
+option mem=0          # Do not add anything to qsub_opts
+option num_threads=* -pe smp $0
+option num_threads=1  # Do not add anything to qsub_opts
+option max_jobs_run=* -tc $0
+option num_nodes=* -pe mpi $0  # You must set this PE as allocation_rule=1
+default gpu=0
+option gpu=0
+option gpu=* -l gpu=$0 -q g.q
diff --git a/egs2/clotho_v2/asr1/conf/slurm.conf b/egs2/clotho_v2/asr1/conf/slurm.conf
new file mode 100644
index 00000000000..3a2e19f78df
--- /dev/null
+++ b/egs2/clotho_v2/asr1/conf/slurm.conf
@@ -0,0 +1,14 @@
+# Default configuration
+command sbatch --export=PATH
+option name=* --job-name $0
+option time=* --time $0
+default mem=10G
+option mem=* --mem-per-cpu $0
+option mem=10G --mem-per-cpu 10G
+option mem=0
+option num_threads=* --cpus-per-task $0
+option num_threads=1 --cpus-per-task 1
+option num_nodes=* --nodes $0
+default gpu=0
+option gpu=0 -p cpu
+option gpu=* -p general --gres=gpu:$0 -c $0
\ No newline at end of file
diff --git a/egs2/clotho_v2/asr1/db.sh b/egs2/clotho_v2/asr1/db.sh
new file mode 120000
index 00000000000..50d86130898
--- /dev/null
+++ b/egs2/clotho_v2/asr1/db.sh
@@ -0,0 +1 @@
+../../TEMPLATE/asr1/db.sh
\ No newline at end of file
diff --git a/egs2/clotho_v2/asr1/local/.ipynb_checkpoints/Checker-checkpoint.ipynb b/egs2/clotho_v2/asr1/local/.ipynb_checkpoints/Checker-checkpoint.ipynb
new file mode 100644
index 00000000000..231fa5bf132
--- /dev/null
+++ b/egs2/clotho_v2/asr1/local/.ipynb_checkpoints/Checker-checkpoint.ipynb
@@ -0,0 +1,91 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 14,
+   "id": "1b6f0ae7-c574-4da9-8ed2-01513728be7e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import pandas as pd\n",
+    "import os\n",
+    "from scipy.io import wavfile\n",
+    "from collections import defaultdict as dd"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 18,
+   "id": "fe8ed4fc-4561-4e71-aa1d-8f41e61c7633",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "base_path = \"/data/user_data/sbharad2/espnet/egs2/clotho_v2/asr1/local/CLOTHO_v2.1/\"\n",
+    "splits = ['development', 'evaluation','test','validation']\n",
+    "\n",
+    "def get_audio_duration(fname):\n",
+    "    samplerate, data = wavfile.read(fname)\n",
+    "    duration_s = len(data)/samplerate\n",
+    "    return duration_s\n",
+    "\n",
+    "# Audio length\n",
+    "audio_s = dd(lambda : [])\n",
+    "for split in splits:\n",
+    "    for fname in os.listdir(os.path.join(base_path,'clotho_audio_files',split)):\n",
+    "        ffname = os.path.join(base_path,'clotho_audio_files',split,fname)\n",
+    "        audio_s[split].append(get_audio_duration(ffname))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 23,
+   "id": "59ea5e97-0f9d-4931-b28c-699513e38318",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "development 30.0 15.0\n",
+      "evaluation 30.0 15.003174603174603\n",
+      "test 30.0 15.000090702947846\n",
+      "validation 29.95390022675737 15.05\n"
+     ]
+    }
+   ],
+   "source": [
+    "for split in splits:\n",
+    "    print(split,max(audio_s[split]),min(audio_s[split]))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b3884f12-1a4a-41ef-b4c9-e7c99af27086",
+   "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.8.19"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}
diff --git a/egs2/clotho_v2/asr1/local/.ipynb_checkpoints/data-checkpoint.sh b/egs2/clotho_v2/asr1/local/.ipynb_checkpoints/data-checkpoint.sh
new file mode 100755
index 00000000000..a44a32ef0d7
--- /dev/null
+++ b/egs2/clotho_v2/asr1/local/.ipynb_checkpoints/data-checkpoint.sh
@@ -0,0 +1,61 @@
+#!/usr/bin/env bash
+# Set bash to 'debug' mode, it will exit on :
+# -e 'error', -u 'undefined variable', -o ... 'error in pipeline', -x 'print commands',
+set -e
+set -u
+set -o pipefail
+
+log() {
+    local fname=${BASH_SOURCE[1]##*/}
+    echo -e "$(date '+%Y-%m-%dT%H:%M:%S') (${fname}:${BASH_LINENO[0]}:${FUNCNAME[1]}) $*"
+}
+SECONDS=0
+
+stage=1
+stop_stage=100
+
+log "$0 $*"
+. utils/parse_options.sh
+
+. ./db.sh
+. ./path.sh
+. ./cmd.sh
+
+if [ $# -ne 0 ]; then
+    log "Error: No positional arguments are required."
+    exit 2
+fi
+
+if [ -z "${CLOTHO_V2}" ]; then
+    log "Fill the value of 'CLOTHO_V2' of db.sh"
+    exit 1
+fi
+
+SPLITS=(development validation evaluation)
+
+if [ ${stage} -le 1 ] && [ ${stop_stage} -ge 1 ]; then
+    log "stage 1: Data preparation"
+    for split_name in ${SPLITS[@]}; do
+        mkdir -p "data/${split_name}"
+    done
+
+    if [ ! -d ${CLOTHO_V2} ]; then
+        echo Cannot find CLOTHO_V2 root! Exiting...
+        exit 1
+    fi
+    
+    # Prepare data in the Kaldi format, including three files:
+    # text, wav.scp, utt2spk
+    echo "$(which python)"
+    python3 local/data_prep.py ${CLOTHO_V2}
+
+    for split_name in ${SPLITS[@]}; do
+        for f in text wav.scp utt2spk; do
+            sort data/${split_name}/${f} -o data/${split_name}/${f}
+        done
+        echo "Running spk2utt"
+        utils/utt2spk_to_spk2utt.pl data/${split_name}/utt2spk > "data/${split_name}/spk2utt"
+    done
+fi
+
+log "Successfully finished. [elapsed=${SECONDS}s]"
\ No newline at end of file
diff --git a/egs2/clotho_v2/asr1/local/.ipynb_checkpoints/data_prep-checkpoint.py b/egs2/clotho_v2/asr1/local/.ipynb_checkpoints/data_prep-checkpoint.py
new file mode 100644
index 00000000000..fb6e5810301
--- /dev/null
+++ b/egs2/clotho_v2/asr1/local/.ipynb_checkpoints/data_prep-checkpoint.py
@@ -0,0 +1,49 @@
+"""Prepares data for Clotho_v2 dataset."""
+
+import os
+import glob
+import sys
+import pandas as pd
+
+
+if __name__ == "__main__":
+    ROOT_DATA_DIR = sys.argv[1]
+    AUDIO_DIR = "clotho_audio_files"
+    CSV_DIR = "clotho_csv_files"
+
+    SPLIT_NAMES = ["development", "validation", "evaluation"]
+
+    for data_split in SPLIT_NAMES:
+        missing_audio = []
+        audio_regex = os.path.join(ROOT_DATA_DIR, AUDIO_DIR, data_split, "*.wav")
+        all_audio_list = glob.glob(audio_regex)
+        all_captions_path = os.path.join(
+            ROOT_DATA_DIR, CSV_DIR, f"clotho_captions_{data_split}.csv"
+        )
+        captions_df = pd.read_csv(all_captions_path)
+        N_PROCESSED = 0
+        with open(
+            os.path.join("data", data_split, "text"), "w", encoding="utf-8"
+        ) as text_f, open(
+            os.path.join("data", data_split, "wav.scp"), "w", encoding="utf-8"
+        ) as wav_scp_f, open(
+            os.path.join("data", data_split, "utt2spk"), "w", encoding="utf-8"
+        ) as utt2spk_f:
+            for uttid, row in captions_df.iterrows():
+                uttid = f"{data_split}_uttid"
+                text = row["caption_1"].strip()
+                file_name = row["file_name"]
+                audio_path = os.path.join(
+                    ROOT_DATA_DIR, AUDIO_DIR, data_split, f"{file_name}"
+                )
+                if audio_path not in all_audio_list:
+                    # Missing audio file
+                    missing_audio += [audio_path]
+                    continue
+                print(f"{uttid} dummy", file=utt2spk_f)
+                print(f"{uttid} {audio_path}", file=wav_scp_f)
+                print(f"{uttid} {text}", file=text_f)
+                N_PROCESSED += 1
+        print(
+            f"For split {data_split}: Processed {N_PROCESSED} audio and their captions. {len(missing_audio)} audio files were missing."
+        )
diff --git a/egs2/clotho_v2/asr1/local/Checker.ipynb b/egs2/clotho_v2/asr1/local/Checker.ipynb
new file mode 100644
index 00000000000..231fa5bf132
--- /dev/null
+++ b/egs2/clotho_v2/asr1/local/Checker.ipynb
@@ -0,0 +1,91 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 14,
+   "id": "1b6f0ae7-c574-4da9-8ed2-01513728be7e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import pandas as pd\n",
+    "import os\n",
+    "from scipy.io import wavfile\n",
+    "from collections import defaultdict as dd"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 18,
+   "id": "fe8ed4fc-4561-4e71-aa1d-8f41e61c7633",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "base_path = \"/data/user_data/sbharad2/espnet/egs2/clotho_v2/asr1/local/CLOTHO_v2.1/\"\n",
+    "splits = ['development', 'evaluation','test','validation']\n",
+    "\n",
+    "def get_audio_duration(fname):\n",
+    "    samplerate, data = wavfile.read(fname)\n",
+    "    duration_s = len(data)/samplerate\n",
+    "    return duration_s\n",
+    "\n",
+    "# Audio length\n",
+    "audio_s = dd(lambda : [])\n",
+    "for split in splits:\n",
+    "    for fname in os.listdir(os.path.join(base_path,'clotho_audio_files',split)):\n",
+    "        ffname = os.path.join(base_path,'clotho_audio_files',split,fname)\n",
+    "        audio_s[split].append(get_audio_duration(ffname))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 23,
+   "id": "59ea5e97-0f9d-4931-b28c-699513e38318",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "development 30.0 15.0\n",
+      "evaluation 30.0 15.003174603174603\n",
+      "test 30.0 15.000090702947846\n",
+      "validation 29.95390022675737 15.05\n"
+     ]
+    }
+   ],
+   "source": [
+    "for split in splits:\n",
+    "    print(split,max(audio_s[split]),min(audio_s[split]))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b3884f12-1a4a-41ef-b4c9-e7c99af27086",
+   "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.8.19"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}
diff --git a/egs2/clotho_v2/asr1/local/data.sh b/egs2/clotho_v2/asr1/local/data.sh
new file mode 100755
index 00000000000..69714c59a60
--- /dev/null
+++ b/egs2/clotho_v2/asr1/local/data.sh
@@ -0,0 +1,75 @@
+#!/usr/bin/env bash
+# Set bash to 'debug' mode, it will exit on :
+# -e 'error', -u 'undefined variable', -o ... 'error in pipeline', -x 'print commands',
+set -e
+set -u
+set -o pipefail
+
+log() {
+    local fname=${BASH_SOURCE[1]##*/}
+    echo -e "$(date '+%Y-%m-%dT%H:%M:%S') (${fname}:${BASH_LINENO[0]}:${FUNCNAME[1]}) $*"
+}
+SECONDS=0
+
+stage=1
+stop_stage=100
+
+log "$0 $*"
+. utils/parse_options.sh
+
+. ./db.sh
+. ./path.sh
+. ./cmd.sh
+
+if [ $# -ne 0 ]; then
+    log "Error: No positional arguments are required."
+    exit 2
+fi
+
+if [ -z "${CLOTHO_V2}" ]; then
+    log "Fill the value of 'CLOTHO_V2' of db.sh"
+    exit 1
+fi
+
+SPLITS=(development validation evaluation)
+N_REF=5
+
+if [ ${stage} -le 1 ] && [ ${stop_stage} -ge 1 ]; then
+    log "stage 1: Data preparation"
+    for split_name in ${SPLITS[@]}; do
+        mkdir -p "data/${split_name}"
+    done
+
+    if [ ! -d ${CLOTHO_V2} ]; then
+        echo Cannot find CLOTHO_V2 root! Exiting...
+        exit 1
+    fi
+    
+    # Prepare data in the Kaldi format, including three files:
+    # text, wav.scp, utt2spk
+    echo "$(which python)"
+    python3 local/data_prep.py ${CLOTHO_V2} ${N_REF}
+
+    for split_name in ${SPLITS[@]}; do
+        for f in wav.scp utt2spk; do
+            sort data/${split_name}/${f} -o data/${split_name}/${f}
+        done
+        # Sort all text files
+        if [ -f data/${split_name}/text ]; then
+            sort data/${split_name}/text -o data/${split_name}/text
+        fi
+        for i in $(seq ${N_REF}); do
+            if [ -f data/${split_name}/text_spk${i} ]; then
+                sort data/${split_name}/text_spk${i} -o data/${split_name}/text_spk${i}
+            fi
+        done
+        echo "Running spk2utt for ${split_name}"
+        utils/utt2spk_to_spk2utt.pl data/${split_name}/utt2spk > "data/${split_name}/spk2utt"
+    done
+fi
+
+log "Successfully finished. [elapsed=${SECONDS}s]"
+
+
+# local/data.sh ${post_process_local_data_opts} --asr_data_dir "${data_feats}/${train_set}"
+# TODO(sbharad2): Copy files for own case
\ No newline at end of file
diff --git a/egs2/clotho_v2/asr1/local/data_prep.py b/egs2/clotho_v2/asr1/local/data_prep.py
new file mode 100644
index 00000000000..3a9e5abde3a
--- /dev/null
+++ b/egs2/clotho_v2/asr1/local/data_prep.py
@@ -0,0 +1,80 @@
+"""Prepares data for Clotho_v2 dataset."""
+
+import os
+import glob
+import sys
+import pandas as pd
+
+
+if __name__ == "__main__":
+    ROOT_DATA_DIR = sys.argv[1]
+    N_REF = int(sys.argv[2])
+    AUDIO_DIR = "clotho_audio_files"
+    CSV_DIR = "clotho_csv_files"
+
+    SPLIT_NAMES = ["development", "validation", "evaluation"]
+
+    for data_split in SPLIT_NAMES:
+        missing_audio = []
+        audio_regex = os.path.join(ROOT_DATA_DIR, AUDIO_DIR, data_split, "*.wav")
+        all_audio_list = glob.glob(audio_regex)
+        all_captions_path = os.path.join(
+            ROOT_DATA_DIR, CSV_DIR, f"clotho_captions_{data_split}.csv"
+        )
+        captions_df = pd.read_csv(all_captions_path)
+        N_PROCESSED = 0
+        with open(
+            os.path.join("data", data_split, "wav.scp"), "w", encoding="utf-8"
+        ) as wav_scp_f, open(
+            os.path.join("data", data_split, "utt2spk"), "w", encoding="utf-8"
+        ) as utt2spk_f:
+            text_f = []
+            try:
+                if data_split == "evaluation":
+                    text_f = [
+                        open(
+                            os.path.join("data", data_split, f"text_spk{n_ref}"),
+                            "w",
+                            encoding="utf-8",
+                        )
+                        for n_ref in range(1, N_REF + 1)
+                    ]
+                else:
+                    text_f = [
+                        open(
+                            os.path.join("data", data_split, "text"),
+                            "w",
+                            encoding="utf-8",
+                        )
+                    ]
+                for uttid, row in captions_df.iterrows():
+                    uttid = f"{data_split}_{uttid}"
+                    file_name = row["file_name"]
+                    audio_path = os.path.join(
+                        ROOT_DATA_DIR, AUDIO_DIR, data_split, f"{file_name}"
+                    )
+                    if audio_path not in all_audio_list:
+                        # Missing audio file
+                        missing_audio += [audio_path]
+                        continue
+                    print(f"{uttid} dummy", file=utt2spk_f)
+                    print(f"{uttid} {audio_path}", file=wav_scp_f)
+                    if data_split == "evaluation":
+                        for i in range(1, N_REF + 1):
+                            text_i = row[f"caption_{i}"].strip()
+                            print(f"{uttid} {text_i}", file=text_f[i - 1])
+                    else:
+                        # TODO(sbharad2): Add options for concatenating data from multiple reference columns,
+                        # choosing particular ref column, randomly choosing col for each example.
+                        text = row["caption_1"].strip()
+                        print(f"{uttid} {text}", file=text_f[0])
+                    N_PROCESSED += 1
+                    if N_PROCESSED % 1000 == 0:
+                        print(f"Processed {N_PROCESSED} audio files.")
+                        continue
+            finally:
+                for f in text_f:
+                    f.close()
+        print(
+            f"For split {data_split}: Processed {N_PROCESSED} audio and their captions. {len(missing_audio)} audio files were missing."
+        )
diff --git a/egs2/clotho_v2/asr1/local/evaluation.py b/egs2/clotho_v2/asr1/local/evaluation.py
new file mode 100644
index 00000000000..04f2062febd
--- /dev/null
+++ b/egs2/clotho_v2/asr1/local/evaluation.py
@@ -0,0 +1,91 @@
+"""Wrapper around aac_metrics (https://github.com/Labbeti/aac-metrics) to evaluate the predictions."""
+
+import argparse
+from aac_metrics import Evaluate
+
+parser = argparse.ArgumentParser()
+parser.add_argument(
+    "--decode_file",
+    required=True,
+    help="Path to file containing predictions.",
+)
+parser.add_argument(
+    "--split",
+    default="evaluation",
+    choices=["evaluation", "validation"],
+    help="Split. Could be either evaluation or validation. Default: evaluation",
+)
+
+args = parser.parse_args()
+
+
+def _parse_and_read_file(file_path: str):
+    read_data = {}
+    try:
+        with open(file_path, "r", encoding="utf-8") as f:
+            for line in f:
+                uttid, prediction = line.strip().split(maxsplit=1)
+                read_data[uttid] = prediction.strip()
+    except Exception as exc:
+        raise ValueError(f"Error reading file {file_path}") from exc
+    if len(read_data) == 0:
+        raise ValueError(f"File {file_path} is empty")
+    return read_data
+
+
+def _convert_to_evaluation_format(predictions, references):
+    candidates = []
+    mult_references = []
+    for key in predictions.keys():
+        candidates.append(predictions[key])
+        if key not in references:
+            raise ValueError(f"Key {key} not found in references")
+        mult_references.append(references[key])
+    return candidates, mult_references
+
+
+def _evaluate(decode_fname: str, split: str):
+    # load all predictions from decode_file
+    predictions = _parse_and_read_file(decode_fname)
+
+    # modify ref file pattern based on split
+    ref_file_pattern = "data/{split}/text{suffix}"
+    ref_file_paths = []
+    if split == "evaluation":
+        ref_file_paths = [
+            ref_file_pattern.format(split="evaluation", suffix=f"_spk{i}")
+            for i in range(1, 6)
+        ]
+    elif split == "validation":
+        ref_file_paths = [ref_file_pattern.format(split="validation", suffix="")]
+    else:
+        raise ValueError(f"Invalid split: {split}")
+
+    # load all references
+    references = {}
+    for ref_file_path in ref_file_paths:
+        reference_ = _parse_and_read_file(ref_file_path)
+        # add reference_ values to reference based on keys
+        for key, value in reference_.items():
+            if key in references:
+                references[key].append(value)
+            else:
+                references[key] = [value]
+
+    candidates, mult_references = _convert_to_evaluation_format(predictions, references)
+    candidates = candidates[:10]
+    mult_references = mult_references[:10]  # TODO(sbharad2): Only for testing
+    evaluate = Evaluate(metrics=["spider", "fense", "vocab"])
+    corpus_scores, _ = evaluate(candidates, mult_references)
+
+    header_str = f" Split: {split} Evaluating over {len(candidates)} predictions. "
+    header_str_l = len(header_str)
+    print("=" * header_str_l)
+    print(header_str)
+    print("=" * header_str_l)
+    for metric, value in corpus_scores.items():
+        print(f" {metric:<20}: {value} ")
+    print("=" * header_str_l)
+
+
+_evaluate(args.decode_file, args.split)
diff --git a/egs2/clotho_v2/asr1/local/path.sh b/egs2/clotho_v2/asr1/local/path.sh
new file mode 100644
index 00000000000..e69de29bb2d
diff --git a/egs2/clotho_v2/asr1/local/score.sh b/egs2/clotho_v2/asr1/local/score.sh
new file mode 100755
index 00000000000..20b81278462
--- /dev/null
+++ b/egs2/clotho_v2/asr1/local/score.sh
@@ -0,0 +1,29 @@
+#!/usr/bin/env bash
+# Copyright 2021 Carnegie Mellon University (Author : Roshan Sharma)
+
+# TODO(sbharad2): Change for your own use case.
+## begin configuration section.
+# end configuration section.
+
+
+[ -f ./path.sh ] && . ./path.sh
+. parse_options.sh || exit 1;
+
+
+if [ $# -lt 1 ]; then
+  echo "Usage: local/score.sh <asr-inference-dir>"
+  exit 1;
+fi
+
+# pip install aac-metrics
+# aac-metrics-download
+asr_expdir=$1
+splits=(evaluation validation)
+
+set -ex
+for split in ${splits[@]}; do
+	for decode_file in $((ls -d ${asr_expdir}/*/*/* && ls -d ${asr_expdir}/*/*) | grep "${split}/text"); do
+		echo "Decode file: ${decode_file} for split:${split}"
+		python local/evaluation.py --decode_file=${decode_file} --split=${split}
+	done
+done
\ No newline at end of file
diff --git a/egs2/clotho_v2/asr1/path.sh b/egs2/clotho_v2/asr1/path.sh
new file mode 120000
index 00000000000..c9ac0a75bc6
--- /dev/null
+++ b/egs2/clotho_v2/asr1/path.sh
@@ -0,0 +1 @@
+../../TEMPLATE/asr1/path.sh
\ No newline at end of file
diff --git a/egs2/clotho_v2/asr1/pyscripts b/egs2/clotho_v2/asr1/pyscripts
new file mode 120000
index 00000000000..ac68ad75b60
--- /dev/null
+++ b/egs2/clotho_v2/asr1/pyscripts
@@ -0,0 +1 @@
+../../TEMPLATE/asr1/pyscripts
\ No newline at end of file
diff --git a/egs2/clotho_v2/asr1/run.sh b/egs2/clotho_v2/asr1/run.sh
new file mode 100755
index 00000000000..a5fc61f626d
--- /dev/null
+++ b/egs2/clotho_v2/asr1/run.sh
@@ -0,0 +1,31 @@
+# !/usr/bin/env bash
+# Set bash to 'debug' mode, it will exit on :
+# -e 'error', -u 'undefined variable', -o ... 'error in pipeline', -x 'print commands',
+set -e
+set -u
+set -o pipefail
+
+timestamp=$(date "+%Y%m%d.%H%M%S")
+    # --dumpdir "dump.${timestamp}" \
+    # --expdir "exp.${timestamp}" \
+
+./asr.sh \
+    --feats_normalize uttmvn \
+    --stage 11 \
+    --stop_stage 11 \
+    --asr_config conf/dcase.yaml \
+    --local_score_opts exp/asr_branchformer_raw_en_word/inference_beam_size10_ctc_weight0.3_asr_model_valid.acc.ave \
+    --ngpu 1 \
+    --gpu_inference true \
+    --nj 1 \
+    --inference_nj 1 \
+    --use_lm false \
+    --lang en \
+    --token_type word \
+    --max_wav_duration 30 \
+    --inference_args "--beam_size 10 --ctc_weight 0.3" \
+    --train_set  development \
+    --valid_set validation \
+    --test_sets "validation evaluation" \
+    --bpe_train_text "data/development/text" \
+    --lm_train_text "data/development/text" "$@"
\ No newline at end of file
diff --git a/egs2/clotho_v2/asr1/scripts b/egs2/clotho_v2/asr1/scripts
new file mode 120000
index 00000000000..b25829705dc
--- /dev/null
+++ b/egs2/clotho_v2/asr1/scripts
@@ -0,0 +1 @@
+../../TEMPLATE/asr1/scripts
\ No newline at end of file
diff --git a/egs2/clotho_v2/asr1/steps b/egs2/clotho_v2/asr1/steps
new file mode 120000
index 00000000000..69ab7056139
--- /dev/null
+++ b/egs2/clotho_v2/asr1/steps
@@ -0,0 +1 @@
+../../TEMPLATE/asr1/steps
\ No newline at end of file
diff --git a/egs2/clotho_v2/asr1/utils b/egs2/clotho_v2/asr1/utils
new file mode 120000
index 00000000000..e18ae14b549
--- /dev/null
+++ b/egs2/clotho_v2/asr1/utils
@@ -0,0 +1 @@
+../../TEMPLATE/asr1/utils
\ No newline at end of file
diff --git a/espnet2/asr/encoder/beats_encoder.py b/espnet2/asr/encoder/beats_encoder.py
new file mode 100644
index 00000000000..7729fc6b562
--- /dev/null
+++ b/espnet2/asr/encoder/beats_encoder.py
@@ -0,0 +1,1310 @@
+# --------------------------------------------------------
+# BEATs: Audio Pre-Training with Acoustic Tokenizers (https://arxiv.org/abs/2212.09058)
+# Adapted from Github source: https://github.com/microsoft/unilm/tree/master/beats
+# Copyright (c) 2022 Microsoft
+# Licensed under The MIT License [see LICENSE for details]
+# Based on fairseq code bases
+# https://github.com/pytorch/fairseq
+# --------------------------------------------------------
+
+from typing import Optional, Dict, Tuple
+import logging
+import math
+import warnings
+import numpy as np
+import torch
+import torch.nn as nn
+from torch.nn import LayerNorm, Parameter
+import torch.nn.functional as F
+import torchaudio.compliance.kaldi as ta_kaldi
+
+from espnet2.asr.encoder.abs_encoder import AbsEncoder
+from espnet2.asr.specaug.specaug import SpecAug
+
+from espnet.nets.pytorch_backend.nets_utils import make_pad_mask
+
+from packaging.version import parse as V
+
+if V(torch.__version__) >= V("1.6.0"):
+    from torch.cuda.amp import autocast
+else:
+    # Nothing to do if torch<1.6.0
+    @contextmanager
+    def autocast(enabled=True):
+        yield
+
+
+class BEATsConfig:
+    def __init__(self, cfg=None):
+        self.input_patch_size: int = 16  # patch size of patch embedding
+        self.embed_dim: int = 512  # patch embedding dimension
+        self.conv_bias: bool = False  # include bias in conv encoder
+
+        self.encoder_layers: int = 12  # num encoder layers in the transformer
+        self.encoder_embed_dim: int = 768  # encoder embedding dimension
+        self.encoder_ffn_embed_dim: int = 3072  # encoder embedding dimension for FFN
+        self.encoder_attention_heads: int = 12  # num encoder attention heads
+        self.activation_fn: str = "gelu"  # activation function to use
+
+        self.layer_wise_gradient_decay_ratio: float = (
+            1.0  # ratio for layer-wise gradient decay
+        )
+        self.layer_norm_first: bool = False  # apply layernorm first in the transformer
+        self.deep_norm: bool = False  # apply deep_norm first in the transformer
+
+        # dropouts
+        self.dropout: float = 0.1  # dropout probability for the transformer
+        self.attention_dropout: float = 0.1  # dropout probability for attention weights
+        self.activation_dropout: float = (
+            0.0  # dropout probability after activation in FFN
+        )
+        self.encoder_layerdrop: float = (
+            0.0  # probability of dropping a tarnsformer layer
+        )
+        self.dropout_input: float = (
+            0.0  # dropout to apply to the input (after feat extr)
+        )
+
+        # positional embeddings
+        self.conv_pos: int = (
+            128  # number of filters for convolutional positional embeddings
+        )
+        self.conv_pos_groups: int = (
+            16  # number of groups for convolutional positional embedding
+        )
+
+        # relative position embedding
+        self.relative_position_embedding: bool = (
+            False  # apply relative position embedding
+        )
+        self.num_buckets: int = 320  # number of buckets for relative position embedding
+        self.max_distance: int = (
+            1280  # maximum distance for relative position embedding
+        )
+        self.gru_rel_pos: bool = False  # apply gated relative position embedding
+
+        # label predictor
+        self.finetuned_model: bool = False  # whether the model is a fine-tuned model.
+        self.predictor_dropout: float = 0.1  # dropout probability for the predictor
+        self.predictor_class: int = 527  # target class number for the predictor
+
+        if cfg is not None:
+            self.update(cfg)
+
+    def update(self, cfg: dict):
+        self.__dict__.update(cfg)
+
+
+class BeatsEncoder(AbsEncoder):
+    """BEATs: Audio Pre-Training with Acoustic Tokenizers.
+
+    (https://arxiv.org/abs/2212.09058)
+    Args:
+        input_size: Input size for xs_pad in `forward`.
+        beats_ckpt_path: Path to a pretrained BEATs checkpoint. If `beats_config` is provided and it does
+            not match the config in the checkpoint, code might throw an error.
+        max_layer: Use all layers for encoding if None. Otherwise use upto `max_layer`.
+        beats_config: `BEATsConfig` object. If provided, we will try to override the config in the
+            checkpoint. This can be used to change dropouts etc for fine-tuning the model while starting
+            from a pretrained checkpoint.
+        specaug_config: Dictionary containing parameters for SpecAugment. If provided, SpecAugment will be applied.
+    """
+
+    def __init__(
+        self,
+        input_size: int,
+        beats_ckpt_path: str = None,
+        max_layer: int = None,
+        beats_config: Optional[BEATsConfig] = None,
+        specaug_config: Optional[Dict] = None,
+    ) -> None:
+        super().__init__()
+
+        self.fbank_mean = 15.41663
+        self.fbank_std = 6.55582
+        self.max_layer = max_layer
+
+        if beats_ckpt_path and beats_config:
+            logging.warning(
+                "Both pretrained checkpoint and config are provided. Config will override."
+            )
+        config = BEATsConfig()
+        self.loaded_state_dict_ = None
+        if beats_ckpt_path is not None:
+            self.loaded_state_dict_ = torch.load(beats_ckpt_path)
+            config = BEATsConfig(self.loaded_state_dict_["cfg"])
+            logging.info(f"Loaded BEATs pretrained config from {beats_ckpt_path}.")
+        if beats_config is not None:
+            config.update(beats_config)
+
+        self.specaug = None
+        if specaug_config is not None:
+            self.specaug = SpecAug(**specaug_config)
+
+        self._output_size = config.encoder_embed_dim
+
+        self.embed = config.embed_dim
+        self.input_patch_size = config.input_patch_size
+        self.post_extract_proj = (
+            nn.Linear(self.embed, config.encoder_embed_dim)
+            if self.embed != config.encoder_embed_dim
+            else None
+        )
+        self.patch_embedding = nn.Conv2d(
+            1,
+            self.embed,
+            kernel_size=self.input_patch_size,
+            stride=self.input_patch_size,
+            bias=config.conv_bias,
+        )
+        self.dropout_input = nn.Dropout(config.dropout_input)
+        assert not config.deep_norm or not config.layer_norm_first
+        self.encoder = TransformerEncoder(config)
+        self.layer_norm = LayerNorm(self.embed)
+
+        self.espnet_initialization_fn()
+        # TODO(sbharad2): add label predictor for fine-tuned model
+
+    def espnet_initialization_fn(self):
+        logging.info("BEATs Intialization function called.")
+        if self.post_extract_proj:
+            torch.nn.init.xavier_normal_(self.post_extract_proj.weight)
+            if self.post_extract_proj.bias is not None:
+                torch.nn.init.constant_(self.post_extract_proj.bias, 0)
+        torch.nn.init.xavier_normal_(self.patch_embedding.weight)
+        if self.patch_embedding.bias is not None:
+            torch.nn.init.constant_(self.patch_embedding.bias, 0)
+
+        # BEATs has different initialization from ESPNet for other modules, so override.
+        self.encoder.apply(init_bert_params)
+        if self.loaded_state_dict_ is not None:
+            load_info = self.load_state_dict(
+                self.loaded_state_dict_["model"], strict=False
+            )
+            # strict=False to ignore Weights in the predictor
+            logging.info(
+                f"Loaded BEATs pretrained model. Following keys were missing: {load_info.missing_keys}. "
+                f"Follwing keys were not found: {load_info.unexpected_keys}."
+            )
+
+    def forward_padding_mask(
+        self,
+        features: torch.Tensor,
+        padding_mask: torch.Tensor,
+    ) -> torch.Tensor:
+        """Forward padding mask."""
+        extra = padding_mask.size(1) % features.size(1)
+        if extra > 0:
+            padding_mask = padding_mask[:, :-extra]
+        padding_mask = padding_mask.view(padding_mask.size(0), features.size(1), -1)
+        padding_mask = padding_mask.all(-1)  # remove empty sequences
+        return padding_mask
+
+    def preprocess(
+        self,
+        source: torch.Tensor,
+    ) -> torch.Tensor:
+        """Preprocess raw audio."""
+        fbanks = []
+        for waveform in source:
+            waveform = waveform.unsqueeze(0) * 2**15  # float32 to int16
+            fbank = ta_kaldi.fbank(
+                waveform,
+                num_mel_bins=128,
+                sample_frequency=16000,
+                frame_length=25,
+                frame_shift=10,
+            )
+            fbanks.append(fbank)
+        fbank = torch.stack(fbanks, dim=0)
+        fbank = (fbank - self.fbank_mean) / (2 * self.fbank_std)
+        return fbank
+
+    def output_size(self) -> int:
+        return self._output_size
+
+    def forward(
+        self,
+        xs_pad: torch.Tensor,
+        ilens: torch.Tensor,
+        prev_states: torch.Tensor = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]:
+        """Wrapper for compatibility with ESPNets' AbsEncoder Interface."""
+
+        # TODO(sbharad2):Explain. If xs is not provided then the operation is costly, because this function tries to create a tensor of size maxlen x maxlen.
+        mask = make_pad_mask(lengths=ilens, xs=xs_pad.unsqueeze(-1).unsqueeze(-1)).to(
+            xs_pad.device
+        )
+        # Adjust shapes to be compatible with BEATs code
+        xs_pad, mask = xs_pad.squeeze(-1).squeeze(-1), mask.squeeze(-1).squeeze(-1)
+        # masks = None
+        last_hidden, _, mask = self.extract_features(
+            xs_pad,
+            mask,
+            max_layer=self.max_layer,
+        )
+        output_lens = (~mask).sum(-1)
+        return last_hidden, output_lens, None
+
+    def extract_features(
+        self,
+        source: torch.Tensor,
+        padding_mask: Optional[torch.Tensor] = None,
+        max_layer: Optional[int] = None,
+    ):
+        """Extract features from raw audio."""
+        with autocast(False):
+            fbank = self.preprocess(source)
+
+            if self.specaug is not None and self.training:
+                # NOTE(sbharad2): We don't provide lengths because espnet implementation does not use it.
+                # /espnet/espnet2/layers/mask_along_axis.py
+                fbank = self.specaug(fbank)[0]
+
+        if padding_mask is not None:
+            padding_mask = self.forward_padding_mask(fbank, padding_mask)
+
+        fbank = fbank.unsqueeze(1).float()
+        features = self.patch_embedding(fbank)
+        features = features.reshape(features.shape[0], features.shape[1], -1)
+        features = features.transpose(1, 2)
+        features = self.layer_norm(features)
+
+        if padding_mask is not None:
+            padding_mask = self.forward_padding_mask(features, padding_mask)
+
+        if self.post_extract_proj is not None:
+            features = self.post_extract_proj(features)
+
+        x = self.dropout_input(features)
+        x, layer_results = self.encoder(x, padding_mask=padding_mask, layer=max_layer)
+        return x, layer_results, padding_mask
+
+
+class TransformerEncoder(nn.Module):
+    """Transformer encoder."""
+
+    def __init__(self, config):
+        super().__init__()
+
+        self.dropout = config.dropout
+        self.embedding_dim = config.encoder_embed_dim
+
+        self.pos_conv = nn.Conv1d(
+            self.embedding_dim,
+            self.embedding_dim,
+            kernel_size=config.conv_pos,
+            padding=config.conv_pos // 2,
+            groups=config.conv_pos_groups,
+        )
+        dropout = 0
+        std = math.sqrt((4 * (1.0 - dropout)) / (config.conv_pos * self.embedding_dim))
+        nn.init.normal_(self.pos_conv.weight, mean=0, std=std)
+        nn.init.constant_(self.pos_conv.bias, 0)
+
+        self.pos_conv = nn.utils.weight_norm(self.pos_conv, name="weight", dim=2)
+        self.pos_conv = nn.Sequential(
+            self.pos_conv, SamePad(config.conv_pos), nn.GELU()
+        )
+
+        if hasattr(config, "relative_position_embedding"):
+            self.relative_position_embedding = config.relative_position_embedding
+            self.num_buckets = config.num_buckets
+            self.max_distance = config.max_distance
+        else:
+            self.relative_position_embedding = False
+            self.num_buckets = 0
+            self.max_distance = 0
+
+        self.layers = nn.ModuleList(
+            [
+                TransformerSentenceEncoderLayer(
+                    embedding_dim=self.embedding_dim,
+                    ffn_embedding_dim=config.encoder_ffn_embed_dim,
+                    num_attention_heads=config.encoder_attention_heads,
+                    dropout=self.dropout,
+                    attention_dropout=config.attention_dropout,
+                    activation_dropout=config.activation_dropout,
+                    activation_fn=config.activation_fn,
+                    layer_norm_first=config.layer_norm_first,
+                    deep_norm=config.deep_norm,
+                    has_relative_attention_bias=self.relative_position_embedding,
+                    num_buckets=self.num_buckets,
+                    max_distance=self.max_distance,
+                    gru_rel_pos=config.gru_rel_pos,
+                    encoder_layers=config.encoder_layers,
+                )
+                for i in range(config.encoder_layers)
+            ]
+        )
+        if self.relative_position_embedding:
+            for i in range(1, config.encoder_layers):
+                del self.layers[i].self_attn.relative_attention_bias
+                self.layers[i].self_attn.relative_attention_bias = self.layers[
+                    0
+                ].self_attn.relative_attention_bias
+
+        self.layer_norm_first = config.layer_norm_first
+        self.layer_norm = LayerNorm(self.embedding_dim)
+        self.layerdrop = config.encoder_layerdrop
+
+        self.apply(init_bert_params)
+
+        if config.deep_norm:
+            deep_norm_beta = math.pow(8 * config.encoder_layers, -1 / 4)
+            for i in range(config.encoder_layers):
+                nn.init.xavier_normal_(self.layers[i].self_attn.k_proj.weight, gain=1)
+                nn.init.xavier_normal_(
+                    self.layers[i].self_attn.v_proj.weight, gain=deep_norm_beta
+                )
+                nn.init.xavier_normal_(self.layers[i].self_attn.q_proj.weight, gain=1)
+                nn.init.xavier_normal_(
+                    self.layers[i].self_attn.out_proj.weight, gain=deep_norm_beta
+                )
+                nn.init.xavier_normal_(self.layers[i].fc1.weight, gain=deep_norm_beta)
+                nn.init.xavier_normal_(self.layers[i].fc2.weight, gain=deep_norm_beta)
+
+        self.layer_wise_gradient_decay_ratio = getattr(
+            config, "layer_wise_gradient_decay_ratio", 1
+        )
+
+    def forward(self, x, padding_mask=None, layer=None):
+        """Forward pass."""
+        x, layer_results = self.extract_features(x, padding_mask, layer)
+
+        if self.layer_norm_first and layer is None:
+            x = self.layer_norm(x)
+
+        return x, layer_results
+
+    def extract_features(self, x, padding_mask=None, tgt_layer=None):
+        """Extract features from the input sequence."""
+
+        if padding_mask is not None:
+            x[padding_mask] = 0
+
+        x_conv = self.pos_conv(x.transpose(1, 2))
+        x_conv = x_conv.transpose(1, 2)
+        x = x + x_conv
+
+        if not self.layer_norm_first:
+            x = self.layer_norm(x)
+
+        x = F.dropout(x, p=self.dropout, training=self.training)
+
+        # B x T x C -> T x B x C
+        x = x.transpose(0, 1)
+
+        layer_results = []
+        z = None
+        if tgt_layer is not None:
+            layer_results.append((x, z))
+        r = None
+        pos_bias = None
+        for i, layer in enumerate(self.layers):
+            if self.layer_wise_gradient_decay_ratio != 1.0:
+                x = GradMultiply.apply((x, self.layer_wise_gradient_decay_ratio))
+            dropout_probability = np.random.random()
+            if not self.training or (dropout_probability > self.layerdrop):
+                x, z, pos_bias = layer(
+                    x,
+                    self_attn_padding_mask=padding_mask,
+                    need_weights=False,
+                    pos_bias=pos_bias,
+                )
+            if tgt_layer is not None:
+                layer_results.append((x, z))
+            if i == tgt_layer:
+                r = x
+                break
+
+        if r is not None:
+            x = r
+
+        # T x B x C -> B x T x C
+        x = x.transpose(0, 1)
+
+        return x, layer_results
+
+
+class TransformerSentenceEncoderLayer(nn.Module):
+    """Transformer encoder layer."""
+
+    def __init__(
+        self,
+        embedding_dim: float = 768,
+        ffn_embedding_dim: float = 3072,
+        num_attention_heads: float = 8,
+        dropout: float = 0.1,
+        attention_dropout: float = 0.1,
+        activation_dropout: float = 0.1,
+        activation_fn: str = "relu",
+        layer_norm_first: bool = False,
+        deep_norm: bool = False,
+        has_relative_attention_bias: bool = False,
+        num_buckets: int = 0,
+        max_distance: int = 0,
+        rescale_init: bool = False,
+        gru_rel_pos: bool = False,
+        encoder_layers: int = 0,
+    ) -> None:
+
+        super().__init__()
+        self.embedding_dim = embedding_dim
+        self.dropout = dropout
+        self.activation_dropout = activation_dropout
+
+        self.activation_name = activation_fn
+        self.activation_fn = get_activation_fn(activation_fn)
+        self.self_attn = MultiheadAttention(
+            self.embedding_dim,
+            num_attention_heads,
+            dropout=attention_dropout,
+            self_attention=True,
+            has_relative_attention_bias=has_relative_attention_bias,
+            num_buckets=num_buckets,
+            max_distance=max_distance,
+            rescale_init=rescale_init,
+            gru_rel_pos=gru_rel_pos,
+        )
+
+        self.dropout1 = nn.Dropout(dropout)
+        self.dropout2 = nn.Dropout(self.activation_dropout)
+        self.dropout3 = nn.Dropout(dropout)
+
+        self.layer_norm_first = layer_norm_first
+
+        self.self_attn_layer_norm = LayerNorm(self.embedding_dim)
+
+        if self.activation_name == "glu":
+            self.fc1 = GLU_Linear(self.embedding_dim, ffn_embedding_dim, "swish")
+        else:
+            self.fc1 = nn.Linear(self.embedding_dim, ffn_embedding_dim)
+        self.fc2 = nn.Linear(ffn_embedding_dim, self.embedding_dim)
+
+        self.final_layer_norm = LayerNorm(self.embedding_dim)
+
+        self.deep_norm = deep_norm
+        if self.deep_norm:
+            self.deep_norm_alpha = math.pow(2 * encoder_layers, 1 / 4)
+        else:
+            self.deep_norm_alpha = 1
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        self_attn_mask: torch.Tensor = None,
+        self_attn_padding_mask: torch.Tensor = None,
+        need_weights: bool = False,
+        pos_bias=None,
+    ):
+        """Forward pass."""
+        residual = x
+
+        if self.layer_norm_first:
+            x = self.self_attn_layer_norm(x)
+            x, attn, pos_bias = self.self_attn(
+                query=x,
+                key=x,
+                value=x,
+                key_padding_mask=self_attn_padding_mask,
+                need_weights=False,
+                attn_mask=self_attn_mask,
+                position_bias=pos_bias,
+            )
+            x = self.dropout1(x)
+            x = residual + x
+
+            residual = x
+            x = self.final_layer_norm(x)
+            if self.activation_name == "glu":
+                x = self.fc1(x)
+            else:
+                x = self.activation_fn(self.fc1(x))
+            x = self.dropout2(x)
+            x = self.fc2(x)
+            x = self.dropout3(x)
+            x = residual + x
+        else:
+            x, attn, pos_bias = self.self_attn(
+                query=x,
+                key=x,
+                value=x,
+                key_padding_mask=self_attn_padding_mask,
+                need_weights=need_weights,
+                attn_mask=self_attn_mask,
+                position_bias=pos_bias,
+            )
+
+            x = self.dropout1(x)
+            x = residual * self.deep_norm_alpha + x
+
+            x = self.self_attn_layer_norm(x)
+
+            residual = x
+            if self.activation_name == "glu":
+                x = self.fc1(x)
+            else:
+                x = self.activation_fn(self.fc1(x))
+            x = self.dropout2(x)
+            x = self.fc2(x)
+            x = self.dropout3(x)
+            x = residual * self.deep_norm_alpha + x
+            x = self.final_layer_norm(x)
+
+        return x, attn, pos_bias
+
+
+class MultiheadAttention(nn.Module):
+    """Multi-headed attention.
+
+    See "Attention Is All You Need" for more details.
+    """
+
+    def __init__(
+        self,
+        embed_dim,
+        num_heads,
+        kdim=None,
+        vdim=None,
+        dropout=0.0,
+        bias=True,
+        add_bias_kv=False,
+        add_zero_attn=False,
+        self_attention=False,
+        encoder_decoder_attention=False,
+        q_noise=0.0,
+        qn_block_size=8,
+        has_relative_attention_bias=False,
+        num_buckets=32,
+        max_distance=128,
+        gru_rel_pos=False,
+        rescale_init=False,
+    ):
+        super().__init__()
+        self.embed_dim = embed_dim
+        self.kdim = kdim if kdim is not None else embed_dim
+        self.vdim = vdim if vdim is not None else embed_dim
+        self.qkv_same_dim = self.kdim == embed_dim and self.vdim == embed_dim
+
+        self.num_heads = num_heads
+        self.dropout_module = nn.Dropout(dropout)
+
+        self.has_relative_attention_bias = has_relative_attention_bias
+        self.num_buckets = num_buckets
+        self.max_distance = max_distance
+        if self.has_relative_attention_bias:
+            self.relative_attention_bias = nn.Embedding(num_buckets, num_heads)
+
+        self.head_dim = embed_dim // num_heads
+        self.q_head_dim = self.head_dim
+        self.k_head_dim = self.head_dim
+        assert (
+            self.head_dim * num_heads == self.embed_dim
+        ), "embed_dim must be divisible by num_heads"
+        self.scaling = self.head_dim**-0.5
+
+        self.self_attention = self_attention
+        self.encoder_decoder_attention = encoder_decoder_attention
+
+        assert not self.self_attention or self.qkv_same_dim, (
+            "Self-attention requires query, key and " "value to be of the same size"
+        )
+
+        k_bias = True
+        if rescale_init:
+            k_bias = False
+
+        k_embed_dim = embed_dim
+        q_embed_dim = embed_dim
+
+        self.k_proj = quant_noise(
+            nn.Linear(self.kdim, k_embed_dim, bias=k_bias), q_noise, qn_block_size
+        )
+        self.v_proj = quant_noise(
+            nn.Linear(self.vdim, embed_dim, bias=bias), q_noise, qn_block_size
+        )
+        self.q_proj = quant_noise(
+            nn.Linear(embed_dim, q_embed_dim, bias=bias), q_noise, qn_block_size
+        )
+
+        self.out_proj = quant_noise(
+            nn.Linear(embed_dim, embed_dim, bias=bias), q_noise, qn_block_size
+        )
+
+        if add_bias_kv:
+            self.bias_k = Parameter(torch.Tensor(1, 1, embed_dim))
+            self.bias_v = Parameter(torch.Tensor(1, 1, embed_dim))
+        else:
+            self.bias_k = self.bias_v = None
+
+        self.add_zero_attn = add_zero_attn
+
+        self.gru_rel_pos = gru_rel_pos
+        if self.gru_rel_pos:
+            self.grep_linear = nn.Linear(self.q_head_dim, 8)
+            self.grep_a = nn.Parameter(torch.ones(1, num_heads, 1, 1))
+
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        """Initiate parameters in the transformer model."""
+        logging.info("Initiate parameters in the MultiheadAttention module.")
+        if self.qkv_same_dim:
+            # Empirically observed the convergence to be much better with
+            # the scaled initialization
+            nn.init.xavier_uniform_(self.k_proj.weight, gain=1 / math.sqrt(2))
+            nn.init.xavier_uniform_(self.v_proj.weight, gain=1 / math.sqrt(2))
+            nn.init.xavier_uniform_(self.q_proj.weight, gain=1 / math.sqrt(2))
+        else:
+            nn.init.xavier_uniform_(self.k_proj.weight)
+            nn.init.xavier_uniform_(self.v_proj.weight)
+            nn.init.xavier_uniform_(self.q_proj.weight)
+
+        nn.init.xavier_uniform_(self.out_proj.weight)
+        if self.out_proj.bias is not None:
+            nn.init.constant_(self.out_proj.bias, 0.0)
+        if self.bias_k is not None:
+            nn.init.xavier_normal_(self.bias_k)
+        if self.bias_v is not None:
+            nn.init.xavier_normal_(self.bias_v)
+        if self.has_relative_attention_bias:
+            nn.init.xavier_normal_(self.relative_attention_bias.weight)
+
+    def _relative_positions_bucket(self, relative_positions, bidirectional=True):
+        num_buckets = self.num_buckets
+        max_distance = self.max_distance
+        relative_buckets = 0
+
+        if bidirectional:
+            num_buckets = num_buckets // 2
+            relative_buckets += (relative_positions > 0).to(torch.long) * num_buckets
+            relative_positions = torch.abs(relative_positions)
+        else:
+            relative_positions = -torch.min(
+                relative_positions, torch.zeros_like(relative_positions)
+            )
+
+        max_exact = num_buckets // 2
+        is_small = relative_positions < max_exact
+
+        relative_postion_if_large = max_exact + (
+            torch.log(relative_positions.float() / max_exact)
+            / math.log(max_distance / max_exact)
+            * (num_buckets - max_exact)
+        ).to(torch.long)
+        relative_postion_if_large = torch.min(
+            relative_postion_if_large,
+            torch.full_like(relative_postion_if_large, num_buckets - 1),
+        )
+
+        relative_buckets += torch.where(
+            is_small, relative_positions, relative_postion_if_large
+        )
+        return relative_buckets
+
+    def compute_bias(self, query_length, key_length):
+        """Compute relative position bias."""
+        context_position = torch.arange(query_length, dtype=torch.long)[:, None]
+        memory_position = torch.arange(key_length, dtype=torch.long)[None, :]
+        relative_position = memory_position - context_position
+        relative_position_bucket = self._relative_positions_bucket(
+            relative_position, bidirectional=True
+        )
+        relative_position_bucket = relative_position_bucket.to(
+            self.relative_attention_bias.weight.device
+        )
+        values = self.relative_attention_bias(relative_position_bucket)
+        values = values.permute([2, 0, 1])
+        return values
+
+    def forward(
+        self,
+        query,
+        key: Optional[torch.Tensor],
+        value: Optional[torch.Tensor],
+        key_padding_mask: Optional[torch.Tensor] = None,
+        incremental_state: Optional[
+            Dict[str, Dict[str, Optional[torch.Tensor]]]
+        ] = None,
+        need_weights: bool = True,
+        static_kv: bool = False,
+        attn_mask: Optional[torch.Tensor] = None,
+        before_softmax: bool = False,
+        need_head_weights: bool = False,
+        position_bias: Optional[torch.Tensor] = None,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]:
+        """Input shape: Time x Batch x Channel
+
+        Args:
+            key_padding_mask (ByteTensor, optional): mask to exclude
+                keys that are pads, of shape `(batch, src_len)`, where
+                padding elements are indicated by 1s.
+            need_weights (bool, optional): return the attention weights,
+                averaged over heads (default: False).
+            attn_mask (ByteTensor, optional): typically used to
+                implement causal attention, where the mask prevents the
+                attention from looking forward in time (default: None).
+            before_softmax (bool, optional): return the raw attention
+                weights and values before the attention softmax.
+            need_head_weights (bool, optional): return the attention
+                weights for each head. Implies *need_weights*. Default:
+                return the average attention weights over all heads.
+        """
+        if need_head_weights:
+            need_weights = True
+
+        is_tpu = query.device.type == "xla"
+
+        tgt_len, bsz, embed_dim = query.size()
+        src_len = tgt_len
+        assert embed_dim == self.embed_dim
+        assert list(query.size()) == [tgt_len, bsz, embed_dim]
+        if key is not None:
+            src_len, key_bsz, _ = key.size()
+            if not torch.jit.is_scripting():
+                assert key_bsz == bsz
+                assert value is not None
+                assert src_len, bsz == value.shape[:2]
+
+        if self.has_relative_attention_bias and position_bias is None:
+            position_bias = self.compute_bias(tgt_len, src_len)
+            position_bias = (
+                position_bias.unsqueeze(0)
+                .repeat(bsz, 1, 1, 1)
+                .view(bsz * self.num_heads, tgt_len, src_len)
+            )
+
+        if incremental_state is not None:
+            saved_state = self._get_input_buffer(incremental_state)
+            if saved_state is not None and "prev_key" in saved_state:
+                # previous time steps are cached - no need to recompute
+                # key and value if they are static
+                if static_kv:
+                    assert self.encoder_decoder_attention and not self.self_attention
+                    key = value = None
+        else:
+            saved_state = None
+
+        if self.self_attention:
+            q = self.q_proj(query)
+            k = self.k_proj(query)
+            v = self.v_proj(query)
+        elif self.encoder_decoder_attention:
+            # encoder-decoder attention
+            q = self.q_proj(query)
+            if key is None:
+                assert value is None
+                k = v = None
+            else:
+                k = self.k_proj(key)
+                v = self.v_proj(key)
+
+        else:
+            assert key is not None and value is not None
+            q = self.q_proj(query)
+            k = self.k_proj(key)
+            v = self.v_proj(value)
+        q *= self.scaling
+        alpha = 32
+        q *= 1 / alpha
+
+        if self.bias_k is not None:
+            assert self.bias_v is not None
+            k = torch.cat([k, self.bias_k.repeat(1, bsz, 1)])
+            v = torch.cat([v, self.bias_v.repeat(1, bsz, 1)])
+            if attn_mask is not None:
+                attn_mask = torch.cat(
+                    [attn_mask, attn_mask.new_zeros(attn_mask.size(0), 1)], dim=1
+                )
+            if key_padding_mask is not None:
+                key_padding_mask = torch.cat(
+                    [
+                        key_padding_mask,
+                        key_padding_mask.new_zeros(key_padding_mask.size(0), 1),
+                    ],
+                    dim=1,
+                )
+
+        q = (
+            q.contiguous()
+            .view(tgt_len, bsz * self.num_heads, self.q_head_dim)
+            .transpose(0, 1)
+        )
+        if k is not None:
+            k = (
+                k.contiguous()
+                .view(-1, bsz * self.num_heads, self.k_head_dim)
+                .transpose(0, 1)
+            )
+        if v is not None:
+            v = (
+                v.contiguous()
+                .view(-1, bsz * self.num_heads, self.head_dim)
+                .transpose(0, 1)
+            )
+
+        if saved_state is not None:
+            # saved states are stored with shape (bsz, num_heads, seq_len, head_dim)
+            if "prev_key" in saved_state:
+                _prev_key = saved_state["prev_key"]
+                assert _prev_key is not None
+                prev_key = _prev_key.view(bsz * self.num_heads, -1, self.head_dim)
+                if static_kv:
+                    k = prev_key
+                else:
+                    assert k is not None
+                    k = torch.cat([prev_key, k], dim=1)
+                src_len = k.size(1)
+            if "prev_value" in saved_state:
+                _prev_value = saved_state["prev_value"]
+                assert _prev_value is not None
+                prev_value = _prev_value.view(bsz * self.num_heads, -1, self.head_dim)
+                if static_kv:
+                    v = prev_value
+                else:
+                    assert v is not None
+                    v = torch.cat([prev_value, v], dim=1)
+            prev_key_padding_mask: Optional[torch.Tensor] = None
+            if "prev_key_padding_mask" in saved_state:
+                prev_key_padding_mask = saved_state["prev_key_padding_mask"]
+            assert k is not None and v is not None
+            key_padding_mask = MultiheadAttention._append_prev_key_padding_mask(
+                key_padding_mask=key_padding_mask,
+                prev_key_padding_mask=prev_key_padding_mask,
+                batch_size=bsz,
+                src_len=k.size(1),
+                static_kv=static_kv,
+            )
+
+            saved_state["prev_key"] = k.view(bsz, self.num_heads, -1, self.head_dim)
+            saved_state["prev_value"] = v.view(bsz, self.num_heads, -1, self.head_dim)
+            saved_state["prev_key_padding_mask"] = key_padding_mask
+            # In this branch incremental_state is never None
+            assert incremental_state is not None
+            incremental_state = self._set_input_buffer(incremental_state, saved_state)
+        assert k is not None
+        assert k.size(1) == src_len
+
+        # This is part of a workaround to get around fork/join parallelism
+        # not supporting Optional types.
+        if key_padding_mask is not None and key_padding_mask.dim() == 0:
+            key_padding_mask = None
+
+        if key_padding_mask is not None:
+            assert key_padding_mask.size(0) == bsz
+            assert key_padding_mask.size(1) == src_len
+
+        if self.add_zero_attn:
+            assert v is not None
+            src_len += 1
+            k = torch.cat([k, k.new_zeros((k.size(0), 1) + k.size()[2:])], dim=1)
+            v = torch.cat([v, v.new_zeros((v.size(0), 1) + v.size()[2:])], dim=1)
+            if attn_mask is not None:
+                attn_mask = torch.cat(
+                    [attn_mask, attn_mask.new_zeros(attn_mask.size(0), 1)], dim=1
+                )
+            if key_padding_mask is not None:
+                key_padding_mask = torch.cat(
+                    [
+                        key_padding_mask,
+                        torch.zeros(key_padding_mask.size(0), 1).type_as(
+                            key_padding_mask
+                        ),
+                    ],
+                    dim=1,
+                )
+
+        attn_weights = torch.bmm(q, k.transpose(1, 2))
+        attn_weights = (
+            attn_weights - attn_weights.max(dim=-1, keepdim=True)[0]
+        ) * alpha
+        attn_weights = self.apply_sparse_mask(attn_weights, tgt_len, src_len, bsz)
+
+        assert list(attn_weights.size()) == [bsz * self.num_heads, tgt_len, src_len]
+
+        if attn_mask is not None:
+            attn_mask = attn_mask.unsqueeze(0)
+            attn_weights += attn_mask
+
+        if key_padding_mask is not None:
+            # don't attend to padding symbols
+            attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len)
+            if not is_tpu:
+                attn_weights = attn_weights.masked_fill(
+                    key_padding_mask.unsqueeze(1).unsqueeze(2).to(torch.bool),
+                    float("-inf"),
+                )
+            else:
+                attn_weights = attn_weights.transpose(0, 2)
+                attn_weights = attn_weights.masked_fill(key_padding_mask, float("-inf"))
+                attn_weights = attn_weights.transpose(0, 2)
+            attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len)
+
+        if before_softmax:
+            return attn_weights, v, position_bias
+
+        if position_bias is not None:
+            attn_mask_rel_pos = position_bias
+            if self.gru_rel_pos == 1:
+                query_layer = (
+                    q.view(bsz, self.num_heads, tgt_len, self.q_head_dim)
+                    * alpha
+                    / self.scaling
+                )
+                _B, _H, _L, __ = query_layer.size()
+                gate_a, gate_b = torch.sigmoid(
+                    self.grep_linear(query_layer)
+                    .view(_B, _H, _L, 2, 4)
+                    .sum(-1, keepdim=False)
+                ).chunk(2, dim=-1)
+                gate_a_1 = gate_a * (gate_b * self.grep_a - 1.0) + 2.0
+                attn_mask_rel_pos = (
+                    gate_a_1.view(bsz * self.num_heads, tgt_len, 1) * position_bias
+                )
+
+            attn_mask_rel_pos = attn_mask_rel_pos.view(attn_weights.size())
+
+            attn_weights = attn_weights + attn_mask_rel_pos
+
+        attn_weights_float = F.softmax(attn_weights, dim=-1)
+        attn_weights = attn_weights_float.type_as(attn_weights)
+        attn_probs = self.dropout_module(attn_weights)
+
+        assert v is not None
+        attn = torch.bmm(attn_probs, v)
+        assert list(attn.size()) == [bsz * self.num_heads, tgt_len, self.head_dim]
+        attn = attn.transpose(0, 1).contiguous().view(tgt_len, bsz, embed_dim)
+        attn = self.out_proj(attn)
+        attn_weights: Optional[torch.Tensor] = None
+        if need_weights:
+            attn_weights = attn_weights_float.view(
+                bsz, self.num_heads, tgt_len, src_len
+            ).transpose(1, 0)
+            if not need_head_weights:
+                # average attention weights over heads
+                attn_weights = attn_weights.mean(dim=0)
+
+        return attn, attn_weights, position_bias
+
+    @staticmethod
+    def _append_prev_key_padding_mask(
+        key_padding_mask: Optional[torch.Tensor],
+        prev_key_padding_mask: Optional[torch.Tensor],
+        batch_size: int,
+        src_len: int,
+        static_kv: bool,
+    ) -> Optional[torch.Tensor]:
+        # saved key padding masks have shape (bsz, seq_len)
+        if prev_key_padding_mask is not None and static_kv:
+            new_key_padding_mask = prev_key_padding_mask
+        elif prev_key_padding_mask is not None and key_padding_mask is not None:
+            new_key_padding_mask = torch.cat(
+                [prev_key_padding_mask.float(), key_padding_mask.float()], dim=1
+            )
+        # During incremental decoding, as the padding token enters and
+        # leaves the frame, there will be a time when prev or current
+        # is None
+        elif prev_key_padding_mask is not None:
+            if src_len > prev_key_padding_mask.size(1):
+                filler = torch.zeros(
+                    (batch_size, src_len - prev_key_padding_mask.size(1)),
+                    device=prev_key_padding_mask.device,
+                )
+                new_key_padding_mask = torch.cat(
+                    [prev_key_padding_mask.float(), filler.float()], dim=1
+                )
+            else:
+                new_key_padding_mask = prev_key_padding_mask.float()
+        elif key_padding_mask is not None:
+            if src_len > key_padding_mask.size(1):
+                filler = torch.zeros(
+                    (batch_size, src_len - key_padding_mask.size(1)),
+                    device=key_padding_mask.device,
+                )
+                new_key_padding_mask = torch.cat(
+                    [filler.float(), key_padding_mask.float()], dim=1
+                )
+            else:
+                new_key_padding_mask = key_padding_mask.float()
+        else:
+            new_key_padding_mask = prev_key_padding_mask
+        return new_key_padding_mask
+
+    def _get_input_buffer(
+        self, incremental_state: Optional[Dict[str, Dict[str, Optional[torch.Tensor]]]]
+    ) -> Dict[str, Optional[torch.Tensor]]:
+        result = self.get_incremental_state(incremental_state, "attn_state")
+        if result is not None:
+            return result
+        else:
+            empty_result: Dict[str, Optional[torch.Tensor]] = {}
+            return empty_result
+
+    def _set_input_buffer(
+        self,
+        incremental_state: Dict[str, Dict[str, Optional[torch.Tensor]]],
+        buffer: Dict[str, Optional[torch.Tensor]],
+    ):
+        return self.set_incremental_state(incremental_state, "attn_state", buffer)
+
+    def apply_sparse_mask(self, attn_weights, tgt_len: int, src_len: int, bsz: int):
+        """No op"""
+        return attn_weights
+
+
+def init_bert_params(module):
+    """
+    Initialize the weights specific to the BERT Model.
+    This overrides the default initializations depending on the specified arguments.
+        1. If normal_init_linear_weights is set then weights of linear
+           layer will be initialized using the normal distribution and
+           bais will be set to the specified value.
+        2. If normal_init_embed_weights is set then weights of embedding
+           layer will be initialized using the normal distribution.
+        3. If normal_init_proj_weights is set then weights of
+           in_project_weight for MultiHeadAttention initialized using
+           the normal distribution (to be validated).
+    """
+
+    def normal_(data):
+        # with FSDP, module params will be on CUDA, so we cast them back to CPU
+        # so that the RNG is consistent with and without FSDP
+        data.copy_(data.cpu().normal_(mean=0.0, std=0.02).to(data.device))
+
+    if isinstance(module, nn.Linear):
+        logging.info("Intializing Linear Layer")
+        normal_(module.weight.data)
+        if module.bias is not None:
+            module.bias.data.zero_()
+    if isinstance(module, nn.Embedding):
+        logging.info("Intializing Embedding Layer")
+        normal_(module.weight.data)
+        if module.padding_idx is not None:
+            module.weight.data[module.padding_idx].zero_()
+    if isinstance(module, MultiheadAttention):
+        logging.info("Intializing Multihead Attention")
+        normal_(module.q_proj.weight.data)
+        normal_(module.k_proj.weight.data)
+        normal_(module.v_proj.weight.data)
+
+
+class GradMultiply(torch.autograd.Function):
+    """A gradient modification function that scales the gradient by a fixed scalar"""
+
+    @staticmethod
+    def forward(ctx, i):
+        """Forward pass"""
+        x, scale = i
+        ctx.scale = scale
+        res = x.new(x)
+        return res
+
+    @staticmethod
+    def backward(ctx, grad):
+        """Backward pass"""
+        return grad * ctx.scale
+
+
+class SamePad(nn.Module):
+    """Change input tensor shape according to the kernel size and type of LM"""
+
+    def __init__(self, kernel_size, causal=False):
+        super().__init__()
+        if causal:
+            self.remove = kernel_size - 1
+        else:
+            self.remove = 1 if kernel_size % 2 == 0 else 0
+
+    def forward(self, x):
+        """Forward pass"""
+        if self.remove > 0:
+            x = x[:, :, : -self.remove]
+        return x
+
+
+class Swish(nn.Module):
+    """Swish activation function"""
+
+    def __init__(self):
+        super(Swish, self).__init__()
+        self.act = torch.nn.Sigmoid()
+
+    def forward(self, x):
+        """Forward pass"""
+        return x * self.act(x)
+
+
+class GLU_Linear(nn.Module):
+    """GLU Linear layer"""
+
+    def __init__(self, input_dim, output_dim, glu_type="sigmoid", bias_in_glu=True):
+        super(GLU_Linear, self).__init__()
+
+        self.glu_type = glu_type
+        self.output_dim = output_dim
+
+        if glu_type == "sigmoid":
+            self.glu_act = torch.nn.Sigmoid()
+        elif glu_type == "swish":
+            self.glu_act = Swish()
+        elif glu_type == "relu":
+            self.glu_act = torch.nn.ReLU()
+        elif glu_type == "gelu":
+            self.glu_act = torch.nn.GELU()
+
+        if bias_in_glu:
+            self.linear = nn.Linear(input_dim, output_dim * 2, True)
+        else:
+            self.linear = nn.Linear(input_dim, output_dim * 2, False)
+
+    def forward(self, x):
+        """Forward pass"""
+        # to be consistent with GLU_Linear, we assume the input always has the #channel (#dim) in the last dimension of the tensor, so need to switch the dimension first for 1D-Conv case
+        x = self.linear(x)
+
+        if self.glu_type == "bilinear":
+            x = (
+                x[:, :, 0 : self.output_dim]
+                * x[:, :, self.output_dim : self.output_dim * 2]
+            )
+        else:
+            x = x[:, :, 0 : self.output_dim] * self.glu_act(
+                x[:, :, self.output_dim : self.output_dim * 2]
+            )
+
+        return x
+
+
+def gelu_accurate(x):
+    if not hasattr(gelu_accurate, "_a"):
+        gelu_accurate._a = math.sqrt(2 / math.pi)
+    return (
+        0.5 * x * (1 + torch.tanh(gelu_accurate._a * (x + 0.044715 * torch.pow(x, 3))))
+    )
+
+
+def gelu(x: torch.Tensor) -> torch.Tensor:
+    return F.gelu(x.float()).type_as(x)
+
+
+def get_activation_fn(activation: str):
+    """Returns the activation function corresponding to `activation`"""
+
+    if activation == "relu":
+        return F.relu
+    elif activation == "gelu":
+        return gelu
+    elif activation == "gelu_fast":
+        warnings.warn("--activation-fn=gelu_fast has been renamed to gelu_accurate")
+        return gelu_accurate
+    elif activation == "gelu_accurate":
+        return gelu_accurate
+    elif activation == "tanh":
+        return torch.tanh
+    elif activation == "linear":
+        return lambda x: x
+    elif activation == "glu":
+        return lambda x: x
+    else:
+        raise RuntimeError(f"--activation-fn {activation} not supported")
+
+
+def quant_noise(module, p, block_size):
+    """
+    Wraps modules and applies quantization noise to the weights for
+    subsequent quantization with Iterative Product Quantization as
+    described in "Training with Quantization Noise for Extreme Model Compression"
+
+    Args:
+        - module: nn.Module
+        - p: amount of Quantization Noise
+        - block_size: size of the blocks for subsequent quantization with iPQ
+
+    Remarks:
+        - Module weights must have the right sizes wrt the block size
+        - Only Linear, Embedding and Conv2d modules are supported for the moment
+        - For more detail on how to quantize by blocks with convolutional weights,
+          see "And the Bit Goes Down: Revisiting the Quantization of Neural Networks"
+        - We implement the simplest form of noise here as stated in the paper
+          which consists in randomly dropping blocks
+    """
+
+    # if no quantization noise, don't register hook
+    if p <= 0:
+        return module
+
+    # supported modules
+    assert isinstance(module, (nn.Linear, nn.Embedding, nn.Conv2d))
+
+    # test whether module.weight has the right sizes wrt block_size
+    is_conv = module.weight.ndim == 4
+
+    # 2D matrix
+    if not is_conv:
+        assert (
+            module.weight.size(1) % block_size == 0
+        ), "Input features must be a multiple of block sizes"
+
+    # 4D matrix
+    else:
+        # 1x1 convolutions
+        if module.kernel_size == (1, 1):
+            assert (
+                module.in_channels % block_size == 0
+            ), "Input channels must be a multiple of block sizes"
+        # regular convolutions
+        else:
+            k = module.kernel_size[0] * module.kernel_size[1]
+            assert k % block_size == 0, "Kernel size must be a multiple of block size"
+
+    def _forward_pre_hook(mod, input):
+        # no noise for evaluation
+        if mod.training:
+            if not is_conv:
+                # gather weight and sizes
+                weight = mod.weight
+                in_features = weight.size(1)
+                out_features = weight.size(0)
+
+                # split weight matrix into blocks and randomly drop selected blocks
+                mask = torch.zeros(
+                    in_features // block_size * out_features, device=weight.device
+                )
+                mask.bernoulli_(p)
+                mask = mask.repeat_interleave(block_size, -1).view(-1, in_features)
+
+            else:
+                # gather weight and sizes
+                weight = mod.weight
+                in_channels = mod.in_channels
+                out_channels = mod.out_channels
+
+                # split weight matrix into blocks and randomly drop selected blocks
+                if mod.kernel_size == (1, 1):
+                    mask = torch.zeros(
+                        int(in_channels // block_size * out_channels),
+                        device=weight.device,
+                    )
+                    mask.bernoulli_(p)
+                    mask = mask.repeat_interleave(block_size, -1).view(-1, in_channels)
+                else:
+                    mask = torch.zeros(
+                        weight.size(0), weight.size(1), device=weight.device
+                    )
+                    mask.bernoulli_(p)
+                    mask = (
+                        mask.unsqueeze(2)
+                        .unsqueeze(3)
+                        .repeat(1, 1, mod.kernel_size[0], mod.kernel_size[1])
+                    )
+
+            # scale weights and apply mask
+            mask = mask.to(
+                torch.bool
+            )  # x.bool() is not currently supported in TorchScript
+            s = 1 / (1 - p)
+            mod.weight.data = s * weight.masked_fill(mask, 0)
+
+    module.register_forward_pre_hook(_forward_pre_hook)
+    return module
diff --git a/espnet2/tasks/asr.py b/espnet2/tasks/asr.py
index 5c737da0dfe..fbb351325d2 100644
--- a/espnet2/tasks/asr.py
+++ b/espnet2/tasks/asr.py
@@ -24,6 +24,7 @@
 )
 from espnet2.asr.decoder.whisper_decoder import OpenAIWhisperDecoder
 from espnet2.asr.encoder.abs_encoder import AbsEncoder
+from espnet2.asr.encoder.beats_encoder import BeatsEncoder
 from espnet2.asr.encoder.avhubert_encoder import FairseqAVHubertEncoder
 from espnet2.asr.encoder.branchformer_encoder import BranchformerEncoder
 from espnet2.asr.encoder.conformer_encoder import ConformerEncoder
@@ -96,7 +97,7 @@
         s3prl=S3prlFrontend,
         fused=FusedFrontends,
         whisper=WhisperFrontend,
-    ),
+    ),  # If setting this to none, please make sure to provide input_size in the config.
     type_check=AbsFrontend,
     default="default",
 )
@@ -159,6 +160,7 @@
         e_branchformer=EBranchformerEncoder,
         avhubert=FairseqAVHubertEncoder,
         multiconv_conformer=MultiConvConformerEncoder,
+        beats=BeatsEncoder,
     ),
     type_check=AbsEncoder,
     default="rnn",

From fd6c73ecb2b0782932c78e2e18e9a63eaa66cb1b Mon Sep 17 00:00:00 2001
From: Shikhar Bharadwaj <sbharad2@andrew.cmu.edu>
Date: Mon, 30 Sep 2024 12:44:25 -0400
Subject: [PATCH 02/47] BEATs encoder with BART decoder for Clotho AAC task

---
 egs2/clotho_v2/asr1/conf/dcase.yaml  | 41 ++++++++++++-------------
 egs2/clotho_v2/asr1/conf/slurm.conf  |  2 +-
 egs2/clotho_v2/asr1/local/score.sh   | 10 ++-----
 egs2/clotho_v2/asr1/run.sh           | 41 ++++++++++++-------------
 espnet2/asr/encoder/beats_encoder.py | 45 ++++++++++++++++++++--------
 espnet2/tasks/abs_task.py            |  8 ++---
 6 files changed, 81 insertions(+), 66 deletions(-)

diff --git a/egs2/clotho_v2/asr1/conf/dcase.yaml b/egs2/clotho_v2/asr1/conf/dcase.yaml
index ee5c7944bed..8b572f51264 100644
--- a/egs2/clotho_v2/asr1/conf/dcase.yaml
+++ b/egs2/clotho_v2/asr1/conf/dcase.yaml
@@ -1,7 +1,7 @@
 batch_type: folded
-batch_size: 80
+batch_size: 64
 accum_grad: 1    # gradient accumulation steps
-max_epoch: 20
+max_epoch: 100
 patience: none
 best_model_criterion:  # criterion to save best models
 -   - valid
@@ -12,11 +12,11 @@ use_amp: true    # whether to use automatic mixed precision
 num_att_plot: 0  # do not save attention plots to save time in the demo
 num_workers: 2   # number of workers in dataloader
 
-input_size: 1
-# beats code takes care of generating mel spectrogram and specaug
+
+# BEATs implementation takes care of generating mel spectrogram, normalization and specaug
 frontend: none
-normalize: none
-init: xavier_normal # For the decoder
+input_size: 1 # important to set input_size to 1 if frontend is none
+normalize: none # BEATs code does global mean and variance normalization
 
 encoder: beats
 encoder_conf:
@@ -35,26 +35,27 @@ encoder_conf:
         - 0
         - 0.12
         num_time_mask: 5
-
-decoder: transformer
+decoder: hugging_face_transformers
 decoder_conf:
-    attention_heads: 4
-    linear_units: 1024
-    num_blocks: 6
-    dropout_rate: 0.1
-    positional_dropout_rate: 0.1
-    self_attention_dropout_rate: 0.1
-    src_attention_dropout_rate: 0.1
-    # use_flash_attn: false
+    model_name_or_path: facebook/bart-base
+
+# The initilization is only for the linear layer between the BEATs encoder and the BART decoder. 
+# BEATs and BART weights override this for encoder and decoder respectively.
+init: kaiming_normal
 
+token_type: hugging_face
+
+# Loss, optimizer, scheduler
 model_conf:
-    ctc_weight: 0.3  # joint CTC/attention training
+    ctc_weight: 0.0  # No CTC, only attention branch
     lsm_weight: 0.1  # label smoothing weight
-    length_normalized_loss: false
+    length_normalized_loss: true
 
 optim: adam
 optim_conf:
-    lr: 0.001
+    lr: 0.00002
+    weight_decay: 0.001
+
 scheduler: warmuplr  # linearly increase and exponentially decrease
 scheduler_conf:
-    warmup_steps: 500
\ No newline at end of file
+    warmup_steps: 100
\ No newline at end of file
diff --git a/egs2/clotho_v2/asr1/conf/slurm.conf b/egs2/clotho_v2/asr1/conf/slurm.conf
index 3a2e19f78df..ae9e7f74459 100644
--- a/egs2/clotho_v2/asr1/conf/slurm.conf
+++ b/egs2/clotho_v2/asr1/conf/slurm.conf
@@ -11,4 +11,4 @@ option num_threads=1 --cpus-per-task 1
 option num_nodes=* --nodes $0
 default gpu=0
 option gpu=0 -p cpu
-option gpu=* -p general --gres=gpu:$0 -c $0
\ No newline at end of file
+option gpu=* -p swl_general --gres=gpu:$0 -c $0
\ No newline at end of file
diff --git a/egs2/clotho_v2/asr1/local/score.sh b/egs2/clotho_v2/asr1/local/score.sh
index 20b81278462..08f3bad9b8c 100755
--- a/egs2/clotho_v2/asr1/local/score.sh
+++ b/egs2/clotho_v2/asr1/local/score.sh
@@ -1,9 +1,4 @@
 #!/usr/bin/env bash
-# Copyright 2021 Carnegie Mellon University (Author : Roshan Sharma)
-
-# TODO(sbharad2): Change for your own use case.
-## begin configuration section.
-# end configuration section.
 
 
 [ -f ./path.sh ] && . ./path.sh
@@ -15,8 +10,9 @@ if [ $# -lt 1 ]; then
   exit 1;
 fi
 
-# pip install aac-metrics
-# aac-metrics-download
+pip install aac-metrics
+aac-metrics-download
+
 asr_expdir=$1
 splits=(evaluation validation)
 
diff --git a/egs2/clotho_v2/asr1/run.sh b/egs2/clotho_v2/asr1/run.sh
index a5fc61f626d..c7488a8aff2 100755
--- a/egs2/clotho_v2/asr1/run.sh
+++ b/egs2/clotho_v2/asr1/run.sh
@@ -1,31 +1,32 @@
-# !/usr/bin/env bash
-# Set bash to 'debug' mode, it will exit on :
-# -e 'error', -u 'undefined variable', -o ... 'error in pipeline', -x 'print commands',
-set -e
-set -u
-set -o pipefail
+#!/usr/bin/env bash
+
+set -euo pipefail
 
 timestamp=$(date "+%Y%m%d.%H%M%S")
-    # --dumpdir "dump.${timestamp}" \
-    # --expdir "exp.${timestamp}" \
+
+wandb_init_args="--use_wandb true --wandb_project DCASE_AAC --wandb_model_log_interval 0"
+wandb_sweep_args="$@"
+
+timestamp=20240928.224746
 
 ./asr.sh \
+    --asr_tag ${timestamp} \
     --feats_normalize uttmvn \
-    --stage 11 \
-    --stop_stage 11 \
-    --asr_config conf/dcase.yaml \
-    --local_score_opts exp/asr_branchformer_raw_en_word/inference_beam_size10_ctc_weight0.3_asr_model_valid.acc.ave \
+    --stage 12 \
+    --stop_stage 12 \
     --ngpu 1 \
     --gpu_inference true \
-    --nj 1 \
+    --nj 10 \
     --inference_nj 1 \
-    --use_lm false \
-    --lang en \
-    --token_type word \
     --max_wav_duration 30 \
-    --inference_args "--beam_size 10 --ctc_weight 0.3" \
-    --train_set  development \
+    --token_type hugging_face \
+    --use_lm false \
+    --hugging_face_model_name_or_path "facebook/bart-base" \
+    --inference_args "--beam_size 10 --ctc_weight 0.0 --hugging_face_decoder True" \
+    --train_set development \
     --valid_set validation \
     --test_sets "validation evaluation" \
-    --bpe_train_text "data/development/text" \
-    --lm_train_text "data/development/text" "$@"
\ No newline at end of file
+    --asr_config conf/dcase.yaml \
+    --inference_asr_model 68epoch.pth \
+    --asr_args "${wandb_init_args} ${wandb_sweep_args}"
+    # --local_score_opts exp/asr_branchformer_raw_en_word/inference_beam_size10_ctc_weight0.3_asr_model_valid.acc.ave \
\ No newline at end of file
diff --git a/espnet2/asr/encoder/beats_encoder.py b/espnet2/asr/encoder/beats_encoder.py
index 7729fc6b562..6f1543fa55b 100644
--- a/espnet2/asr/encoder/beats_encoder.py
+++ b/espnet2/asr/encoder/beats_encoder.py
@@ -5,6 +5,9 @@
 # Licensed under The MIT License [see LICENSE for details]
 # Based on fairseq code bases
 # https://github.com/pytorch/fairseq
+
+# This code is adapted from the original BEATs implementation and
+#  can be used to pre-train/and or fine-tune BEATs model.
 # --------------------------------------------------------
 
 from typing import Optional, Dict, Tuple
@@ -123,19 +126,27 @@ def __init__(
         self.fbank_mean = 15.41663
         self.fbank_std = 6.55582
         self.max_layer = max_layer
-
+        self.beats_ckpt_path = beats_ckpt_path
+
+        # Four cases for loading Beats config:
+        # 1. No checkpoint and no config: Default config
+        # 2. Checkpoint and no user-provided config: Load config from checkpoint
+        # 3. Checkpoint and user-provided config: Merge the two, but override with user-provided config
+        # 4. No checkpoint and user-provided config: Use user-provided config
+        config = BEATsConfig()  # Default config
         if beats_ckpt_path and beats_config:
             logging.warning(
-                "Both pretrained checkpoint and config are provided. Config will override."
+                "Both pretrained checkpoint and config are provided."
+                " We will override ckpt config with user-provided config."
             )
-        config = BEATsConfig()
         self.loaded_state_dict_ = None
         if beats_ckpt_path is not None:
             self.loaded_state_dict_ = torch.load(beats_ckpt_path)
-            config = BEATsConfig(self.loaded_state_dict_["cfg"])
             logging.info(f"Loaded BEATs pretrained config from {beats_ckpt_path}.")
+            config = BEATsConfig(self.loaded_state_dict_["cfg"])
         if beats_config is not None:
             config.update(beats_config)
+            logging.info("Overriding BEATs config with user-provided config.")
 
         self.specaug = None
         if specaug_config is not None:
@@ -162,11 +173,16 @@ def __init__(
         self.encoder = TransformerEncoder(config)
         self.layer_norm = LayerNorm(self.embed)
 
-        self.espnet_initialization_fn()
-        # TODO(sbharad2): add label predictor for fine-tuned model
-
     def espnet_initialization_fn(self):
-        logging.info("BEATs Intialization function called.")
+        """Initialization function for BEATs.
+
+        This must be called last in the initialization procedure.
+        The initialization occurs in three steps:
+        1. ESPNet initializes all modules.
+        2. This function initializes BEATs encoder overriding 1.
+        3. Optionally, if we have the pretrained checkpoint, we load the weights from the checkpoint overriding 2 and 1.
+        """
+        logging.info("BEATs Initialization function called.")
         if self.post_extract_proj:
             torch.nn.init.xavier_normal_(self.post_extract_proj.weight)
             if self.post_extract_proj.bias is not None:
@@ -183,8 +199,9 @@ def espnet_initialization_fn(self):
             )
             # strict=False to ignore Weights in the predictor
             logging.info(
-                f"Loaded BEATs pretrained model. Following keys were missing: {load_info.missing_keys}. "
-                f"Follwing keys were not found: {load_info.unexpected_keys}."
+                f"Loaded BEATs pretrained model. Following keys were missing in your custom model: {load_info.missing_keys}. "
+                f"Follwing keys could not be loaded from the pretrained checkpoint: {load_info.unexpected_keys}."
+                "It is expected to have 'predictor' here if you are fine-tuning with only the BEATs backbone."
             )
 
     def forward_padding_mask(
@@ -231,7 +248,9 @@ def forward(
     ) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]:
         """Wrapper for compatibility with ESPNets' AbsEncoder Interface."""
 
-        # TODO(sbharad2):Explain. If xs is not provided then the operation is costly, because this function tries to create a tensor of size maxlen x maxlen.
+        # NOTE(shikhar): If xs is not provided then the operation is costly, because this
+        # function tries to create a tensor of size maxlen x maxlen. Therfore, we unsqueeze
+        # and then squeeze tensors.
         mask = make_pad_mask(lengths=ilens, xs=xs_pad.unsqueeze(-1).unsqueeze(-1)).to(
             xs_pad.device
         )
@@ -257,8 +276,8 @@ def extract_features(
             fbank = self.preprocess(source)
 
             if self.specaug is not None and self.training:
-                # NOTE(sbharad2): We don't provide lengths because espnet implementation does not use it.
-                # /espnet/espnet2/layers/mask_along_axis.py
+                # NOTE(shikhar): We don't provide lengths because espnet implementation
+                # does not use it --> /espnet/espnet2/layers/mask_along_axis.py
                 fbank = self.specaug(fbank)[0]
 
         if padding_mask is not None:
diff --git a/espnet2/tasks/abs_task.py b/espnet2/tasks/abs_task.py
index 3f84fc6d1e1..a575d3c1440 100644
--- a/espnet2/tasks/abs_task.py
+++ b/espnet2/tasks/abs_task.py
@@ -1518,10 +1518,8 @@ def main_worker(cls, args: argparse.Namespace):
                     else:
                         project = args.wandb_project
 
-                    if args.wandb_name is None:
-                        name = str(Path(".").resolve()).replace("/", "_")
-                    else:
-                        name = args.wandb_name
+                    # Wandb server generates a random name, if args.wandb_name is None
+                    name = args.wandb_name
 
                     wandb.init(
                         entity=args.wandb_entity,
@@ -1534,7 +1532,7 @@ def main_worker(cls, args: argparse.Namespace):
                     wandb.config.update(args)
                 else:
                     # wandb also supports grouping for distributed training,
-                    # but we only logs aggregated data,
+                    # but we only log aggregated data,
                     # so it's enough to perform on rank0 node.
                     args.use_wandb = False
 

From dfa7aeee904b1dc64cf6048cd20af54acd0ce21e Mon Sep 17 00:00:00 2001
From: Shikhar Bharadwaj <sbharad2@andrew.cmu.edu>
Date: Mon, 30 Sep 2024 13:14:22 -0400
Subject: [PATCH 03/47] recipe downloads data now

---
 egs2/TEMPLATE/asr1/db.sh          |  2 +-
 egs2/clotho_v2/asr1/local/data.sh | 90 ++++++++++++++++++-------------
 2 files changed, 54 insertions(+), 38 deletions(-)

diff --git a/egs2/TEMPLATE/asr1/db.sh b/egs2/TEMPLATE/asr1/db.sh
index e94e95f2733..157471e78a8 100755
--- a/egs2/TEMPLATE/asr1/db.sh
+++ b/egs2/TEMPLATE/asr1/db.sh
@@ -219,7 +219,7 @@ SPRING_INX=downloads
 VOXCELEB=
 KSPONSPEECH=
 HIFITTS=downloads
-CLOTHO_V2=/data/user_data/sbharad2/espnet/egs2/clotho_v2/asr1/local/CLOTHO_v2.1 # TODO(sbharad2): Downloads
+CLOTHO_V2=downloads
 
 # For only CMU TIR environment
 if [[ "$(hostname)" == tir* ]]; then
diff --git a/egs2/clotho_v2/asr1/local/data.sh b/egs2/clotho_v2/asr1/local/data.sh
index 69714c59a60..3d85cdb25e3 100755
--- a/egs2/clotho_v2/asr1/local/data.sh
+++ b/egs2/clotho_v2/asr1/local/data.sh
@@ -11,9 +11,6 @@ log() {
 }
 SECONDS=0
 
-stage=1
-stop_stage=100
-
 log "$0 $*"
 . utils/parse_options.sh
 
@@ -31,45 +28,64 @@ if [ -z "${CLOTHO_V2}" ]; then
     exit 1
 fi
 
-SPLITS=(development validation evaluation)
-N_REF=5
-
-if [ ${stage} -le 1 ] && [ ${stop_stage} -ge 1 ]; then
-    log "stage 1: Data preparation"
-    for split_name in ${SPLITS[@]}; do
-        mkdir -p "data/${split_name}"
-    done
-
-    if [ ! -d ${CLOTHO_V2} ]; then
-        echo Cannot find CLOTHO_V2 root! Exiting...
+## DOWNLOAD DATA
+CLOTHO_V2_ROOT_DIR="$(pwd)/local"
+if [ ! -e "${CLOTHO_V2_ROOT_DIR}/download_done" ]; then
+    echo "Installing aac-datasets."
+    if ! pip install aac-datasets; then
+        echo "Error: Installing aac-datasets failed."
         exit 1
     fi
-    
-    # Prepare data in the Kaldi format, including three files:
-    # text, wav.scp, utt2spk
-    echo "$(which python)"
-    python3 local/data_prep.py ${CLOTHO_V2} ${N_REF}
-
-    for split_name in ${SPLITS[@]}; do
-        for f in wav.scp utt2spk; do
-            sort data/${split_name}/${f} -o data/${split_name}/${f}
-        done
-        # Sort all text files
-        if [ -f data/${split_name}/text ]; then
-            sort data/${split_name}/text -o data/${split_name}/text
+    echo "Downlaoding clotho into ${CLOTHO_V2_ROOT_DIR}."
+    mkdir -p "${CLOTHO_V2_ROOT_DIR}"
+    for split in val eval; do
+        echo "Downloading ${split} split."
+        if ! aac-datasets-download --root "${CLOTHO_V2_ROOT_DIR}" clotho --subsets "${split}"; then
+            echo "Error: Downloading Clotho dataset failed."
+            exit 1
         fi
-        for i in $(seq ${N_REF}); do
-            if [ -f data/${split_name}/text_spk${i} ]; then
-                sort data/${split_name}/text_spk${i} -o data/${split_name}/text_spk${i}
-            fi
-        done
-        echo "Running spk2utt for ${split_name}"
-        utils/utt2spk_to_spk2utt.pl data/${split_name}/utt2spk > "data/${split_name}/spk2utt"
     done
+    touch "${CLOTHO_V2_ROOT_DIR}/download_done"
+else
+    echo "Clotho dataset is already downloaded. ${CLOTHO_V2_ROOT_DIR}/download_done exists."
 fi
 
-log "Successfully finished. [elapsed=${SECONDS}s]"
+SPLITS=(development validation evaluation) # map dev, val, eval to development, validation, evaluation
+N_REF=5
+CLOTHO_V2_ROOT_DIR="${CLOTHO_V2_ROOT_DIR}/CLOTHO_v2.1"
+
+
+## PREPARE DATA
+log "stage 1: Data preparation"
+for split_name in ${SPLITS[@]}; do
+    mkdir -p "data/${split_name}"
+done
 
+if [ ! -d ${CLOTHO_V2_ROOT_DIR} ]; then
+    echo Cannot find CLOTHO_V2_ROOT_DIR root! Exiting...
+    exit 1
+fi
+
+# Prepare data in the Kaldi format, including three files:
+# text, wav.scp, utt2spk
+echo "$(which python)"
+python3 local/data_prep.py ${CLOTHO_V2_ROOT_DIR} ${N_REF}
+
+for split_name in ${SPLITS[@]}; do
+    for f in wav.scp utt2spk; do
+        sort data/${split_name}/${f} -o data/${split_name}/${f}
+    done
+    # Sort all text files
+    if [ -f data/${split_name}/text ]; then
+        sort data/${split_name}/text -o data/${split_name}/text
+    fi
+    for i in $(seq ${N_REF}); do
+        if [ -f data/${split_name}/text_spk${i} ]; then
+            sort data/${split_name}/text_spk${i} -o data/${split_name}/text_spk${i}
+        fi
+    done
+    echo "Running spk2utt for ${split_name}"
+    utils/utt2spk_to_spk2utt.pl data/${split_name}/utt2spk > "data/${split_name}/spk2utt"
+done
 
-# local/data.sh ${post_process_local_data_opts} --asr_data_dir "${data_feats}/${train_set}"
-# TODO(sbharad2): Copy files for own case
\ No newline at end of file
+log "Successfully finished. [elapsed=${SECONDS}s]"
\ No newline at end of file

From 7800c2c9049703369118de6a14c2930895bc0be3 Mon Sep 17 00:00:00 2001
From: Shikhar Bharadwaj <sbharad2@andrew.cmu.edu>
Date: Mon, 30 Sep 2024 13:22:53 -0400
Subject: [PATCH 04/47] add readme

---
 egs2/README.md | 1 +
 1 file changed, 1 insertion(+)

diff --git a/egs2/README.md b/egs2/README.md
index e6461313648..99e7192d1ba 100644
--- a/egs2/README.md
+++ b/egs2/README.md
@@ -30,6 +30,7 @@ See: https://espnet.github.io/espnet/espnet2_tutorial.html#recipes-using-espnet2
 | bur_openslr80           | Burmese ASR training dataset                                                                                                     | ASR                     | BUR                   | https://openslr.org/80/                                                                                      |              |
 | catslu               	  | CATSLU-MAPS                                                                                                                      | SLU                     | CMN           	       | https://sites.google.com/view/catslu/home                                                                    |              |
 | catslu_entity        	  | CATSLU                                                                                                                           | SLU/Entity Classifi.    | CMN           	       | https://sites.google.com/view/catslu/home                                                                    |              |
+| clotho_v2                  | Clotho v2.1 dataset for audio captioning | ASR   | ENG                   | https://zenodo.org/records/4783391
 | chime1                  | The 1st CHiME Speech Separation and Recognition Challenge                                                                        | ASR/Multichannel ASR    | ENG                   | https://spandh.dcs.shef.ac.uk/chime_challenge/chime2011/                                                      |              |
 | chime2                  | The 2nd CHiME Speech Separation and Recognition Challenge                                                                        | ASR/Multichannel ASR    | ENG                   | https://spandh.dcs.shef.ac.uk/chime_challenge/chime2013/                                                      |              |
 | chime4                  | The 4th CHiME Speech Separation and Recognition Challenge                                                                        | ASR/Multichannel ASR    | ENG                   | http://spandh.dcs.shef.ac.uk/chime_challenge/chime2016/                                                      |              |

From d9ac98e011d30bc9d838f78d919d6e21b383d30b Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Mon, 30 Sep 2024 17:39:55 +0000
Subject: [PATCH 05/47] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 egs2/clotho_v2/asr1/conf/dcase.yaml                      | 4 ++--
 egs2/clotho_v2/asr1/conf/slurm.conf                      | 2 +-
 .../asr1/local/.ipynb_checkpoints/data-checkpoint.sh     | 4 ++--
 .../local/.ipynb_checkpoints/data_prep-checkpoint.py     | 4 ++--
 egs2/clotho_v2/asr1/local/data.sh                        | 2 +-
 egs2/clotho_v2/asr1/local/data_prep.py                   | 4 ++--
 egs2/clotho_v2/asr1/local/evaluation.py                  | 1 +
 egs2/clotho_v2/asr1/local/score.sh                       | 2 +-
 egs2/clotho_v2/asr1/run.sh                               | 2 +-
 espnet2/asr/encoder/beats_encoder.py                     | 9 ++++-----
 espnet2/tasks/asr.py                                     | 2 +-
 11 files changed, 18 insertions(+), 18 deletions(-)

diff --git a/egs2/clotho_v2/asr1/conf/dcase.yaml b/egs2/clotho_v2/asr1/conf/dcase.yaml
index 8b572f51264..987096f6ddf 100644
--- a/egs2/clotho_v2/asr1/conf/dcase.yaml
+++ b/egs2/clotho_v2/asr1/conf/dcase.yaml
@@ -39,7 +39,7 @@ decoder: hugging_face_transformers
 decoder_conf:
     model_name_or_path: facebook/bart-base
 
-# The initilization is only for the linear layer between the BEATs encoder and the BART decoder. 
+# The initilization is only for the linear layer between the BEATs encoder and the BART decoder.
 # BEATs and BART weights override this for encoder and decoder respectively.
 init: kaiming_normal
 
@@ -58,4 +58,4 @@ optim_conf:
 
 scheduler: warmuplr  # linearly increase and exponentially decrease
 scheduler_conf:
-    warmup_steps: 100
\ No newline at end of file
+    warmup_steps: 100
diff --git a/egs2/clotho_v2/asr1/conf/slurm.conf b/egs2/clotho_v2/asr1/conf/slurm.conf
index ae9e7f74459..cb27726bedb 100644
--- a/egs2/clotho_v2/asr1/conf/slurm.conf
+++ b/egs2/clotho_v2/asr1/conf/slurm.conf
@@ -11,4 +11,4 @@ option num_threads=1 --cpus-per-task 1
 option num_nodes=* --nodes $0
 default gpu=0
 option gpu=0 -p cpu
-option gpu=* -p swl_general --gres=gpu:$0 -c $0
\ No newline at end of file
+option gpu=* -p swl_general --gres=gpu:$0 -c $0
diff --git a/egs2/clotho_v2/asr1/local/.ipynb_checkpoints/data-checkpoint.sh b/egs2/clotho_v2/asr1/local/.ipynb_checkpoints/data-checkpoint.sh
index a44a32ef0d7..e984e589eaf 100755
--- a/egs2/clotho_v2/asr1/local/.ipynb_checkpoints/data-checkpoint.sh
+++ b/egs2/clotho_v2/asr1/local/.ipynb_checkpoints/data-checkpoint.sh
@@ -43,7 +43,7 @@ if [ ${stage} -le 1 ] && [ ${stop_stage} -ge 1 ]; then
         echo Cannot find CLOTHO_V2 root! Exiting...
         exit 1
     fi
-    
+
     # Prepare data in the Kaldi format, including three files:
     # text, wav.scp, utt2spk
     echo "$(which python)"
@@ -58,4 +58,4 @@ if [ ${stage} -le 1 ] && [ ${stop_stage} -ge 1 ]; then
     done
 fi
 
-log "Successfully finished. [elapsed=${SECONDS}s]"
\ No newline at end of file
+log "Successfully finished. [elapsed=${SECONDS}s]"
diff --git a/egs2/clotho_v2/asr1/local/.ipynb_checkpoints/data_prep-checkpoint.py b/egs2/clotho_v2/asr1/local/.ipynb_checkpoints/data_prep-checkpoint.py
index fb6e5810301..52c7ba98290 100644
--- a/egs2/clotho_v2/asr1/local/.ipynb_checkpoints/data_prep-checkpoint.py
+++ b/egs2/clotho_v2/asr1/local/.ipynb_checkpoints/data_prep-checkpoint.py
@@ -1,10 +1,10 @@
 """Prepares data for Clotho_v2 dataset."""
 
-import os
 import glob
+import os
 import sys
-import pandas as pd
 
+import pandas as pd
 
 if __name__ == "__main__":
     ROOT_DATA_DIR = sys.argv[1]
diff --git a/egs2/clotho_v2/asr1/local/data.sh b/egs2/clotho_v2/asr1/local/data.sh
index 3d85cdb25e3..a9b5c0de3ea 100755
--- a/egs2/clotho_v2/asr1/local/data.sh
+++ b/egs2/clotho_v2/asr1/local/data.sh
@@ -88,4 +88,4 @@ for split_name in ${SPLITS[@]}; do
     utils/utt2spk_to_spk2utt.pl data/${split_name}/utt2spk > "data/${split_name}/spk2utt"
 done
 
-log "Successfully finished. [elapsed=${SECONDS}s]"
\ No newline at end of file
+log "Successfully finished. [elapsed=${SECONDS}s]"
diff --git a/egs2/clotho_v2/asr1/local/data_prep.py b/egs2/clotho_v2/asr1/local/data_prep.py
index 3a9e5abde3a..a515eb5b6b7 100644
--- a/egs2/clotho_v2/asr1/local/data_prep.py
+++ b/egs2/clotho_v2/asr1/local/data_prep.py
@@ -1,10 +1,10 @@
 """Prepares data for Clotho_v2 dataset."""
 
-import os
 import glob
+import os
 import sys
-import pandas as pd
 
+import pandas as pd
 
 if __name__ == "__main__":
     ROOT_DATA_DIR = sys.argv[1]
diff --git a/egs2/clotho_v2/asr1/local/evaluation.py b/egs2/clotho_v2/asr1/local/evaluation.py
index 04f2062febd..158fe6e9d5d 100644
--- a/egs2/clotho_v2/asr1/local/evaluation.py
+++ b/egs2/clotho_v2/asr1/local/evaluation.py
@@ -1,6 +1,7 @@
 """Wrapper around aac_metrics (https://github.com/Labbeti/aac-metrics) to evaluate the predictions."""
 
 import argparse
+
 from aac_metrics import Evaluate
 
 parser = argparse.ArgumentParser()
diff --git a/egs2/clotho_v2/asr1/local/score.sh b/egs2/clotho_v2/asr1/local/score.sh
index 08f3bad9b8c..06b2d1751c8 100755
--- a/egs2/clotho_v2/asr1/local/score.sh
+++ b/egs2/clotho_v2/asr1/local/score.sh
@@ -22,4 +22,4 @@ for split in ${splits[@]}; do
 		echo "Decode file: ${decode_file} for split:${split}"
 		python local/evaluation.py --decode_file=${decode_file} --split=${split}
 	done
-done
\ No newline at end of file
+done
diff --git a/egs2/clotho_v2/asr1/run.sh b/egs2/clotho_v2/asr1/run.sh
index c7488a8aff2..57be575e4c3 100755
--- a/egs2/clotho_v2/asr1/run.sh
+++ b/egs2/clotho_v2/asr1/run.sh
@@ -29,4 +29,4 @@ timestamp=20240928.224746
     --asr_config conf/dcase.yaml \
     --inference_asr_model 68epoch.pth \
     --asr_args "${wandb_init_args} ${wandb_sweep_args}"
-    # --local_score_opts exp/asr_branchformer_raw_en_word/inference_beam_size10_ctc_weight0.3_asr_model_valid.acc.ave \
\ No newline at end of file
+    # --local_score_opts exp/asr_branchformer_raw_en_word/inference_beam_size10_ctc_weight0.3_asr_model_valid.acc.ave \
diff --git a/espnet2/asr/encoder/beats_encoder.py b/espnet2/asr/encoder/beats_encoder.py
index 6f1543fa55b..d58ee0e0174 100644
--- a/espnet2/asr/encoder/beats_encoder.py
+++ b/espnet2/asr/encoder/beats_encoder.py
@@ -10,24 +10,23 @@
 #  can be used to pre-train/and or fine-tune BEATs model.
 # --------------------------------------------------------
 
-from typing import Optional, Dict, Tuple
 import logging
 import math
 import warnings
+from typing import Dict, Optional, Tuple
+
 import numpy as np
 import torch
 import torch.nn as nn
-from torch.nn import LayerNorm, Parameter
 import torch.nn.functional as F
 import torchaudio.compliance.kaldi as ta_kaldi
+from packaging.version import parse as V
+from torch.nn import LayerNorm, Parameter
 
 from espnet2.asr.encoder.abs_encoder import AbsEncoder
 from espnet2.asr.specaug.specaug import SpecAug
-
 from espnet.nets.pytorch_backend.nets_utils import make_pad_mask
 
-from packaging.version import parse as V
-
 if V(torch.__version__) >= V("1.6.0"):
     from torch.cuda.amp import autocast
 else:
diff --git a/espnet2/tasks/asr.py b/espnet2/tasks/asr.py
index fbb351325d2..a9b9bd1eb7d 100644
--- a/espnet2/tasks/asr.py
+++ b/espnet2/tasks/asr.py
@@ -24,8 +24,8 @@
 )
 from espnet2.asr.decoder.whisper_decoder import OpenAIWhisperDecoder
 from espnet2.asr.encoder.abs_encoder import AbsEncoder
-from espnet2.asr.encoder.beats_encoder import BeatsEncoder
 from espnet2.asr.encoder.avhubert_encoder import FairseqAVHubertEncoder
+from espnet2.asr.encoder.beats_encoder import BeatsEncoder
 from espnet2.asr.encoder.branchformer_encoder import BranchformerEncoder
 from espnet2.asr.encoder.conformer_encoder import ConformerEncoder
 from espnet2.asr.encoder.contextual_block_conformer_encoder import (

From 1b4b322ec2505a347781fe580f0a7c0cbe0614cf Mon Sep 17 00:00:00 2001
From: Shikhar Bharadwaj <sbharad2@andrew.cmu.edu>
Date: Sun, 13 Oct 2024 09:42:53 -0400
Subject: [PATCH 06/47] fix bug in pad, add option to use all layers, use last
 layer as adapter

---
 espnet2/asr/encoder/beats_encoder.py | 34 +++++++++++++++++++++++++---
 1 file changed, 31 insertions(+), 3 deletions(-)

diff --git a/espnet2/asr/encoder/beats_encoder.py b/espnet2/asr/encoder/beats_encoder.py
index d58ee0e0174..67415a9fb65 100644
--- a/espnet2/asr/encoder/beats_encoder.py
+++ b/espnet2/asr/encoder/beats_encoder.py
@@ -117,6 +117,7 @@ def __init__(
         input_size: int,
         beats_ckpt_path: str = None,
         max_layer: int = None,
+        use_all_layers: bool = False,
         beats_config: Optional[BEATsConfig] = None,
         specaug_config: Optional[Dict] = None,
     ) -> None:
@@ -171,6 +172,10 @@ def __init__(
         assert not config.deep_norm or not config.layer_norm_first
         self.encoder = TransformerEncoder(config)
         self.layer_norm = LayerNorm(self.embed)
+        self.use_all_layers = use_all_layers
+        if self.use_all_layers:
+            assert self.max_layer is not None, "max_layer must be provided."
+            self.layer_weights = nn.Parameter(torch.ones(self.max_layer))
 
     def espnet_initialization_fn(self):
         """Initialization function for BEATs.
@@ -193,6 +198,20 @@ def espnet_initialization_fn(self):
         # BEATs has different initialization from ESPNet for other modules, so override.
         self.encoder.apply(init_bert_params)
         if self.loaded_state_dict_ is not None:
+            # TODO(shikhar): Hardcoded
+            all_keys = list(self.loaded_state_dict_["model"].items())
+            for key, _ in all_keys:
+                if (
+                    key.startswith("encoder.layers.11")
+                    # or key.startswith("encoder.layers.10")
+                    # or key.startswith("encoder.layers.9")
+                    # or key.startswith("encoder.layers.8")
+                    # or key.startswith("encoder.layers.7")
+                    # or key.startswith("encoder.layers.6")
+                ):
+                    self.loaded_state_dict_["model"].pop(key)
+                    logging.info(f"Removed key: {key}")
+
             load_info = self.load_state_dict(
                 self.loaded_state_dict_["model"], strict=False
             )
@@ -250,9 +269,9 @@ def forward(
         # NOTE(shikhar): If xs is not provided then the operation is costly, because this
         # function tries to create a tensor of size maxlen x maxlen. Therfore, we unsqueeze
         # and then squeeze tensors.
-        mask = make_pad_mask(lengths=ilens, xs=xs_pad.unsqueeze(-1).unsqueeze(-1)).to(
-            xs_pad.device
-        )
+        mask = make_pad_mask(
+            lengths=ilens, xs=xs_pad.unsqueeze(-1).unsqueeze(-1), length_dim=1
+        ).to(xs_pad.device)
         # Adjust shapes to be compatible with BEATs code
         xs_pad, mask = xs_pad.squeeze(-1).squeeze(-1), mask.squeeze(-1).squeeze(-1)
         # masks = None
@@ -296,6 +315,15 @@ def extract_features(
 
         x = self.dropout_input(features)
         x, layer_results = self.encoder(x, padding_mask=padding_mask, layer=max_layer)
+
+        if self.use_all_layers:
+            for i, (x_i, _) in enumerate(layer_results):
+                if i > self.max_layer:
+                    break
+                if i == 0:
+                    x = self.layer_weights[i] * x_i
+                x += self.layer_weights[i] * x_i
+
         return x, layer_results, padding_mask
 
 

From aa8910e107440c14a9e22e35e252c562e636552e Mon Sep 17 00:00:00 2001
From: Shikhar Bharadwaj <sbharadwaj@dt-login03.delta.ncsa.illinois.edu>
Date: Sun, 13 Oct 2024 10:22:11 -0500
Subject: [PATCH 07/47] fix downloading for train set, change pre-trained model
 path for delta

---
 egs2/clotho_v2/asr1/conf/dcase.yaml | 2 +-
 egs2/clotho_v2/asr1/local/data.sh   | 2 +-
 egs2/clotho_v2/asr1/run.sh          | 2 +-
 3 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/egs2/clotho_v2/asr1/conf/dcase.yaml b/egs2/clotho_v2/asr1/conf/dcase.yaml
index 987096f6ddf..dccb6e7ed6c 100644
--- a/egs2/clotho_v2/asr1/conf/dcase.yaml
+++ b/egs2/clotho_v2/asr1/conf/dcase.yaml
@@ -20,7 +20,7 @@ normalize: none # BEATs code does global mean and variance normalization
 
 encoder: beats
 encoder_conf:
-    beats_ckpt_path: /home/sbharad2/dcase2024_cp/dcase24task6/pretrained_weights/BEATs_iter3_plus_AS2M_finetuned_on_AS2M_cpt1.pt
+    beats_ckpt_path: /u/sbharadwaj/BEATs_models/BEATs_iter3_plus_AS2M.pt
     beats_config:
         input_patch_size: 16
         embed_dim: 512
diff --git a/egs2/clotho_v2/asr1/local/data.sh b/egs2/clotho_v2/asr1/local/data.sh
index a9b5c0de3ea..db5df6682ca 100755
--- a/egs2/clotho_v2/asr1/local/data.sh
+++ b/egs2/clotho_v2/asr1/local/data.sh
@@ -38,7 +38,7 @@ if [ ! -e "${CLOTHO_V2_ROOT_DIR}/download_done" ]; then
     fi
     echo "Downlaoding clotho into ${CLOTHO_V2_ROOT_DIR}."
     mkdir -p "${CLOTHO_V2_ROOT_DIR}"
-    for split in val eval; do
+    for split in val eval dev; do
         echo "Downloading ${split} split."
         if ! aac-datasets-download --root "${CLOTHO_V2_ROOT_DIR}" clotho --subsets "${split}"; then
             echo "Error: Downloading Clotho dataset failed."
diff --git a/egs2/clotho_v2/asr1/run.sh b/egs2/clotho_v2/asr1/run.sh
index 57be575e4c3..579124dd774 100755
--- a/egs2/clotho_v2/asr1/run.sh
+++ b/egs2/clotho_v2/asr1/run.sh
@@ -12,7 +12,7 @@ timestamp=20240928.224746
 ./asr.sh \
     --asr_tag ${timestamp} \
     --feats_normalize uttmvn \
-    --stage 12 \
+    --stage 1 \
     --stop_stage 12 \
     --ngpu 1 \
     --gpu_inference true \

From 01e54fa5cdb6b133fa8b52613c59f89654e2d5a6 Mon Sep 17 00:00:00 2001
From: Shikhar Bharadwaj <sbharadwaj@dt-login03.delta.ncsa.illinois.edu>
Date: Tue, 5 Nov 2024 09:02:37 -0600
Subject: [PATCH 08/47] fix beats layernorm initialization bug

---
 .gitignore                                    |   3 +-
 .../conf/{dcase.yaml => beats_bart_ft.yaml}   |  38 +-
 egs2/clotho_v2/asr1/conf/beats_bart_pt.yaml   |  79 +++
 egs2/clotho_v2/asr1/conf/branchformer.yaml    |  65 ---
 egs2/clotho_v2/asr1/conf/slurm.conf           |   6 +-
 egs2/clotho_v2/asr1/local/data.sh             |  65 ++-
 .../asr1/local/data_prep_audiocaps.py         |  54 +++
 .../{data_prep.py => data_prep_clotho.py}     |  16 +-
 .../local/data_prep_clotho_chatgpt_mixup.py   | 449 ++++++++++++++++++
 egs2/clotho_v2/asr1/local/evaluation.py       |  19 +-
 egs2/clotho_v2/asr1/local/score.sh            |   5 +-
 egs2/clotho_v2/asr1/run.sh                    |  32 --
 egs2/clotho_v2/asr1/run_ft.sh                 |  33 ++
 egs2/clotho_v2/asr1/run_pt.sh                 |  70 +++
 espnet2/asr/encoder/beats_encoder.py          |  70 +--
 15 files changed, 833 insertions(+), 171 deletions(-)
 rename egs2/clotho_v2/asr1/conf/{dcase.yaml => beats_bart_ft.yaml} (73%)
 create mode 100644 egs2/clotho_v2/asr1/conf/beats_bart_pt.yaml
 delete mode 100644 egs2/clotho_v2/asr1/conf/branchformer.yaml
 create mode 100644 egs2/clotho_v2/asr1/local/data_prep_audiocaps.py
 rename egs2/clotho_v2/asr1/local/{data_prep.py => data_prep_clotho.py} (82%)
 create mode 100644 egs2/clotho_v2/asr1/local/data_prep_clotho_chatgpt_mixup.py
 delete mode 100755 egs2/clotho_v2/asr1/run.sh
 create mode 100755 egs2/clotho_v2/asr1/run_ft.sh
 create mode 100755 egs2/clotho_v2/asr1/run_pt.sh

diff --git a/.gitignore b/.gitignore
index ed37f59e70c..9a1fb57bfa9 100644
--- a/.gitignore
+++ b/.gitignore
@@ -102,5 +102,4 @@ tools/miniconda
 tools/ffmpeg-*-static
 tools/ffmpeg-release
 tools/ffmpeg-release-*.tar.xz
-tools/Miniconda3-latest-*.sh
-egs2/clotho_v2/asr1/local/CLOTHO_v2.1/*
+tools/Miniconda3-latest-*.sh
\ No newline at end of file
diff --git a/egs2/clotho_v2/asr1/conf/dcase.yaml b/egs2/clotho_v2/asr1/conf/beats_bart_ft.yaml
similarity index 73%
rename from egs2/clotho_v2/asr1/conf/dcase.yaml
rename to egs2/clotho_v2/asr1/conf/beats_bart_ft.yaml
index dccb6e7ed6c..dd4e1fcdbcd 100644
--- a/egs2/clotho_v2/asr1/conf/dcase.yaml
+++ b/egs2/clotho_v2/asr1/conf/beats_bart_ft.yaml
@@ -1,7 +1,9 @@
 batch_type: folded
+unused_parameters: true
 batch_size: 64
-accum_grad: 1    # gradient accumulation steps
-max_epoch: 100
+accum_grad: 2
+max_epoch: 50
+grad_clip: 2
 patience: none
 best_model_criterion:  # criterion to save best models
 -   - valid
@@ -18,12 +20,16 @@ frontend: none
 input_size: 1 # important to set input_size to 1 if frontend is none
 normalize: none # BEATs code does global mean and variance normalization
 
+freeze_param: [
+    "encoder.encoder",
+    "encoder.layer_norm",
+    "encoder.patch_embedding",
+    "encoder.post_extract_proj",
+]
+
 encoder: beats
 encoder_conf:
-    beats_ckpt_path: /u/sbharadwaj/BEATs_models/BEATs_iter3_plus_AS2M.pt
-    beats_config:
-        input_patch_size: 16
-        embed_dim: 512
+    beats_ckpt_path: /u/sbharadwaj/BEATs_models/BEATs_iter3_plus_AS2M.pt # _finetuned_on_AS2M_cpt1
     specaug_config:
         apply_freq_mask: true
         freq_mask_width_range:
@@ -35,13 +41,17 @@ encoder_conf:
         - 0
         - 0.12
         num_time_mask: 5
+    adapter_layers: 2
+    max_layer: 12
+    use_weighted_representation: true
+    
 decoder: hugging_face_transformers
 decoder_conf:
     model_name_or_path: facebook/bart-base
 
 # The initilization is only for the linear layer between the BEATs encoder and the BART decoder.
 # BEATs and BART weights override this for encoder and decoder respectively.
-init: kaiming_normal
+init: kaiming_uniform
 
 token_type: hugging_face
 
@@ -50,12 +60,18 @@ model_conf:
     ctc_weight: 0.0  # No CTC, only attention branch
     lsm_weight: 0.1  # label smoothing weight
     length_normalized_loss: true
+    # BART Dictionary customizations
+    ignore_id: 1
+    sym_blank: "<pad>"
+    sym_sos: "<s>"
+    sym_eos: "</s>"
+    sym_space: "Ġ"
 
-optim: adam
+optim: adamw
 optim_conf:
-    lr: 0.00002
+    lr: 0.00005 # 5e-5
     weight_decay: 0.001
 
-scheduler: warmuplr  # linearly increase and exponentially decrease
+scheduler: warmuplr
 scheduler_conf:
-    warmup_steps: 100
+    warmup_steps: 4000
diff --git a/egs2/clotho_v2/asr1/conf/beats_bart_pt.yaml b/egs2/clotho_v2/asr1/conf/beats_bart_pt.yaml
new file mode 100644
index 00000000000..74a65947c81
--- /dev/null
+++ b/egs2/clotho_v2/asr1/conf/beats_bart_pt.yaml
@@ -0,0 +1,79 @@
+batch_type: folded
+unused_parameters: true
+batch_size: 64
+accum_grad: 8
+max_epoch: 30
+grad_clip: 2
+patience: none
+best_model_criterion:  # criterion to save best models
+-   - valid
+    - acc
+    - max
+keep_nbest_models: 5  # save nbest models and average these checkpoints
+use_amp: true    # whether to use automatic mixed precision
+num_att_plot: 0  # do not save attention plots to save time in the demo
+num_workers: 2   # number of workers in dataloader
+
+
+# BEATs implementation takes care of generating mel spectrogram, normalization and specaug
+frontend: none
+input_size: 1 # important to set input_size to 1 if frontend is none
+normalize: none # BEATs code does global mean and variance normalization
+
+
+freeze_param: [
+    "encoder.encoder",
+    "encoder.layer_norm",
+    "encoder.patch_embedding",
+    "encoder.post_extract_proj",
+]
+
+
+encoder: beats
+encoder_conf:
+    beats_ckpt_path: /u/sbharadwaj/BEATs_models/BEATs_iter3_plus_AS2M_finetuned_on_AS2M_cpt1.pt # _finetuned_on_AS2M_cpt1
+    specaug_config:
+        apply_freq_mask: true
+        freq_mask_width_range:
+        - 0
+        - 64
+        num_freq_mask: 2
+        apply_time_mask: true
+        time_mask_width_ratio_range:
+        - 0
+        - 0.12
+        num_time_mask: 5
+    adapter_layers: 2
+    max_layer: 12
+    use_weighted_representation: true
+    
+decoder: hugging_face_transformers
+decoder_conf:
+    model_name_or_path: facebook/bart-base
+
+# The initilization is only for the linear layer between the BEATs encoder and the BART decoder.
+# BEATs and BART weights override this for encoder and decoder respectively.
+init: kaiming_uniform
+
+token_type: hugging_face
+
+# Loss, optimizer, scheduler
+model_conf:
+    ctc_weight: 0.0  # No CTC, only attention branch
+    lsm_weight: 0.1  # label smoothing weight
+    length_normalized_loss: true
+    # BART Dictionary customizations
+    ignore_id: 1
+    sym_blank: "<pad>"
+    sym_sos: "<s>"
+    sym_eos: "</s>"
+    sym_space: "Ġ"
+
+optim: adamw
+optim_conf:
+    lr: 0.0002 # 2e-4
+    weight_decay: 0.001
+
+scheduler: warmuplr
+scheduler_conf:
+    warmup_steps: 8000
diff --git a/egs2/clotho_v2/asr1/conf/branchformer.yaml b/egs2/clotho_v2/asr1/conf/branchformer.yaml
deleted file mode 100644
index b092371029a..00000000000
--- a/egs2/clotho_v2/asr1/conf/branchformer.yaml
+++ /dev/null
@@ -1,65 +0,0 @@
-batch_type: folded
-batch_size: 64
-accum_grad: 2    # gradient accumulation steps
-max_epoch: 100
-patience: none
-init: xavier_uniform
-best_model_criterion:  # criterion to save best models
--   - valid
-    - acc
-    - max
-keep_nbest_models: 10  # save nbest models and average these checkpoints
-use_amp: true    # whether to use automatic mixed precision
-num_att_plot: 0  # do not save attention plots to save time in the demo
-num_workers: 2   # number of workers in dataloader
-
-encoder: transformer
-encoder_conf:
-    output_size: 256
-    attention_heads: 4
-    linear_units: 1024
-    num_blocks: 12
-    dropout_rate: 0.1
-    positional_dropout_rate: 0.1
-    attention_dropout_rate: 0.1
-    input_layer: conv2d
-    normalize_before: true
-    # use_flash_attn: false
-
-decoder: transformer
-decoder_conf:
-    attention_heads: 4
-    linear_units: 1024
-    num_blocks: 6
-    dropout_rate: 0.1
-    positional_dropout_rate: 0.1
-    self_attention_dropout_rate: 0.1
-    src_attention_dropout_rate: 0.1
-    # use_flash_attn: false
-
-model_conf:
-    ctc_weight: 0.3  # joint CTC/attention training
-    lsm_weight: 0.1  # label smoothing weight
-    length_normalized_loss: false
-
-optim: adam
-optim_conf:
-    lr: 0.001
-scheduler: warmuplr  # linearly increase and exponentially decrease
-scheduler_conf:
-    warmup_steps: 800
-specaug: specaug
-specaug_conf:
-    apply_time_warp: true
-    time_warp_window: 5
-    time_warp_mode: bicubic
-    apply_freq_mask: true
-    freq_mask_width_range:
-    - 0
-    - 30
-    num_freq_mask: 2
-    apply_time_mask: true
-    time_mask_width_range:
-    - 0
-    - 40
-    num_time_mask: 2
diff --git a/egs2/clotho_v2/asr1/conf/slurm.conf b/egs2/clotho_v2/asr1/conf/slurm.conf
index cb27726bedb..9f397a1b55d 100644
--- a/egs2/clotho_v2/asr1/conf/slurm.conf
+++ b/egs2/clotho_v2/asr1/conf/slurm.conf
@@ -1,14 +1,16 @@
 # Default configuration
 command sbatch --export=PATH
 option name=* --job-name $0
+default time=2-00:00:00
 option time=* --time $0
 default mem=10G
 option mem=* --mem-per-cpu $0
 option mem=10G --mem-per-cpu 10G
 option mem=0
+default num_threads=4
 option num_threads=* --cpus-per-task $0
 option num_threads=1 --cpus-per-task 1
 option num_nodes=* --nodes $0
 default gpu=0
-option gpu=0 -p cpu
-option gpu=* -p swl_general --gres=gpu:$0 -c $0
+option gpu=0 -p cpu -A bbjs-delta-cpu
+option gpu=* -p gpuA40x4 -A bbjs-delta-gpu --gres=gpu:$0 -c $0
diff --git a/egs2/clotho_v2/asr1/local/data.sh b/egs2/clotho_v2/asr1/local/data.sh
index db5df6682ca..aa005c9059e 100755
--- a/egs2/clotho_v2/asr1/local/data.sh
+++ b/egs2/clotho_v2/asr1/local/data.sh
@@ -28,10 +28,11 @@ if [ -z "${CLOTHO_V2}" ]; then
     exit 1
 fi
 
+log "stage 1: Data preparation"
 ## DOWNLOAD DATA
-CLOTHO_V2_ROOT_DIR="$(pwd)/local"
+CLOTHO_V2_ROOT_DIR="$(pwd)/downloads"
 if [ ! -e "${CLOTHO_V2_ROOT_DIR}/download_done" ]; then
-    echo "Installing aac-datasets."
+    log "stage 1: Data preparation - Installing aac-datasets"
     if ! pip install aac-datasets; then
         echo "Error: Installing aac-datasets failed."
         exit 1
@@ -50,27 +51,65 @@ else
     echo "Clotho dataset is already downloaded. ${CLOTHO_V2_ROOT_DIR}/download_done exists."
 fi
 
-SPLITS=(development validation evaluation) # map dev, val, eval to development, validation, evaluation
-N_REF=5
 CLOTHO_V2_ROOT_DIR="${CLOTHO_V2_ROOT_DIR}/CLOTHO_v2.1"
 
-
 ## PREPARE DATA
-log "stage 1: Data preparation"
-for split_name in ${SPLITS[@]}; do
-    mkdir -p "data/${split_name}"
-done
-
 if [ ! -d ${CLOTHO_V2_ROOT_DIR} ]; then
-    echo Cannot find CLOTHO_V2_ROOT_DIR root! Exiting...
+    echo Cannot find ${CLOTHO_V2_ROOT_DIR} directory! Exiting...
+    exit 1
+fi
+if [ -z "${AUDIOCAPS}" ]; then
+    log "Fill the value of 'AUDIOCAPS' of db.sh"
+    exit 1
+fi
+if [ -z "${CLOTHO_CHATGPT_MIXUP}" ]; then
+    log "Fill the value of 'CLOTHO_CHATGPT_MIXUP' of db.sh"
     exit 1
 fi
 
+AUDIOCAPS_ROOT_DIR=${AUDIOCAPS}
+if [ ! -d ${AUDIOCAPS_ROOT_DIR} ]; then
+    echo Cannot find ${AUDIOCAPS_ROOT_DIR} directory! Exiting...
+    exit 1
+fi
+CLOTHO_CHATGPT_MIXUP_ROOT_DIR=${CLOTHO_CHATGPT_MIXUP}
+if [ ! -d ${CLOTHO_CHATGPT_MIXUP} ]; then
+    echo Cannot find ${CLOTHO_CHATGPT_MIXUP} directory! Exiting...
+    exit 1
+fi
+
+SPLITS=(development_clotho validation evaluation) # map dev, val, eval to development, validation, evaluation
+N_REF=5
+for split_name in ${SPLITS[@]}; do
+    mkdir -p "data/${split_name}"
+done
+
+echo "$(which python)"
 # Prepare data in the Kaldi format, including three files:
 # text, wav.scp, utt2spk
-echo "$(which python)"
-python3 local/data_prep.py ${CLOTHO_V2_ROOT_DIR} ${N_REF}
+##########
+python3 local/data_prep_clotho.py ${CLOTHO_V2_ROOT_DIR} ${N_REF}
+##########
+mkdir -p data/development_audiocaps
+python3 local/data_prep_audiocaps.py ${AUDIOCAPS_ROOT_DIR}
+SPLITS+=(development_audiocaps)
+##########
+mkdir -p data/development_clotho_chatgpt_mixup
+python3 local/data_prep_clotho_chatgpt_mixup.py ${CLOTHO_CHATGPT_MIXUP_ROOT_DIR} ${CLOTHO_V2_ROOT_DIR}
+SPLITS+=(development_clotho_chatgpt_mixup)
+##########
+
+# CONCATENATE MIXUP AND AUDIOCAPS
+mkdir -p data/pretrain
+for f in wav.scp utt2spk text; do
+    cat data/development_audiocaps/${f} data/development_clotho_chatgpt_mixup/${f} > data/pretrain/${f}
+done
+SPLITS+=(pretrain)
+
+# SIX directories: 4 training sets - development_clotho, development_audiocaps, development_clotho_chatgpt_mixup, development_pretrain
+# and 2 for val, eval - validation, evaluation
 
+# SORT ALL
 for split_name in ${SPLITS[@]}; do
     for f in wav.scp utt2spk; do
         sort data/${split_name}/${f} -o data/${split_name}/${f}
diff --git a/egs2/clotho_v2/asr1/local/data_prep_audiocaps.py b/egs2/clotho_v2/asr1/local/data_prep_audiocaps.py
new file mode 100644
index 00000000000..93b3938ec06
--- /dev/null
+++ b/egs2/clotho_v2/asr1/local/data_prep_audiocaps.py
@@ -0,0 +1,54 @@
+"""Prepares data for AudioCaps dataset."""
+
+import os
+import sys
+import json
+from scipy.io import wavfile
+
+if __name__ == "__main__":
+    ROOT_DATA_DIR = sys.argv[1]
+    TRAIN_FILES = "train_split.json" # this is a json with key "audiocap_ids" containing list of folder_ids 
+    DATA_DIR = "data" # this has folders each for a folder_id, containing audio.wav and metadata.json files
+    
+    LOCAL_DATA_DIR = "data/development_audiocaps"
+
+    all_train_ids = None
+    with open(os.path.join(ROOT_DATA_DIR,TRAIN_FILES)) as f:
+        train_meta = json.load(f)
+        all_train_ids = train_meta['audiocap_ids']
+    assert all_train_ids is not None, f'Could not read {ROOT_DATA_DIR/TRAIN_FILES}'
+
+    output_wav_scp_path= os.path.join(LOCAL_DATA_DIR, 'wav.scp')
+    output_utt2spk_path= os.path.join(LOCAL_DATA_DIR, 'utt2spk')
+    output_txt_path = os.path.join(LOCAL_DATA_DIR, 'text')
+    missing_audio = []
+    N_PROCESSED = 0
+    N_ZERO_LENGTH_CAPTIONS = 0
+    N_ZERO_LENGTH_AUDIO =0 
+    N_VERY_LONG_AUDIO = 0
+    AUDIO_LENGTH_THRESHOLD = 30 # seconds
+    with open(output_txt_path, 'w') as text_f, open(output_wav_scp_path,'w') as wav_scp_f, open(output_utt2spk_path,'w') as utt2spk_f:
+        for train_id in all_train_ids:
+            uttid = f'audiocaps_train_{train_id}'
+            audio_path = os.path.join(ROOT_DATA_DIR, DATA_DIR, train_id, "audio.wav")
+            samplerate, data = wavfile.read(audio_path) 
+            audio_length = len(data)/samplerate
+            if audio_length==0:
+                N_ZERO_LENGTH_AUDIO+=1
+                continue
+            if audio_length > AUDIO_LENGTH_THRESHOLD:
+                N_VERY_LONG_AUDIO+=1
+                continue
+            with open(os.path.join(ROOT_DATA_DIR, DATA_DIR, train_id, "metadata.json")) as f:
+                meta = json.load(f)
+                text = meta['caption'].strip()
+                if len(text)==0:
+                    N_ZERO_LENGTH_CAPTIONS += 1
+                    continue
+                print(f"{uttid} {audio_path}", file=wav_scp_f)
+                print(f"{uttid} {text}", file=text_f)
+                print(f"{uttid} dummy", file=utt2spk_f)
+                N_PROCESSED += 1
+                if N_PROCESSED % 1000 == 0:
+                    print(f"Processed {N_PROCESSED} audio files.")
+    print(f"Processed {N_PROCESSED} audio files. Skipped {N_VERY_LONG_AUDIO} long audio, {N_ZERO_LENGTH_AUDIO} empty audio and {N_ZERO_LENGTH_CAPTIONS} empty caption items.")
\ No newline at end of file
diff --git a/egs2/clotho_v2/asr1/local/data_prep.py b/egs2/clotho_v2/asr1/local/data_prep_clotho.py
similarity index 82%
rename from egs2/clotho_v2/asr1/local/data_prep.py
rename to egs2/clotho_v2/asr1/local/data_prep_clotho.py
index a515eb5b6b7..d9e8b92d033 100644
--- a/egs2/clotho_v2/asr1/local/data_prep.py
+++ b/egs2/clotho_v2/asr1/local/data_prep_clotho.py
@@ -23,17 +23,21 @@
         )
         captions_df = pd.read_csv(all_captions_path)
         N_PROCESSED = 0
+        if data_split=='development':
+            dir_write_name = 'development_clotho'
+        else:
+            dir_write_name = data_split # for validation and evaluation
         with open(
-            os.path.join("data", data_split, "wav.scp"), "w", encoding="utf-8"
+            os.path.join("data", dir_write_name, "wav.scp"), "w", encoding="utf-8"
         ) as wav_scp_f, open(
-            os.path.join("data", data_split, "utt2spk"), "w", encoding="utf-8"
+            os.path.join("data", dir_write_name, "utt2spk"), "w", encoding="utf-8"
         ) as utt2spk_f:
             text_f = []
             try:
                 if data_split == "evaluation":
                     text_f = [
                         open(
-                            os.path.join("data", data_split, f"text_spk{n_ref}"),
+                            os.path.join("data", dir_write_name, f"text_spk{n_ref}"),
                             "w",
                             encoding="utf-8",
                         )
@@ -42,13 +46,13 @@
                 else:
                     text_f = [
                         open(
-                            os.path.join("data", data_split, "text"),
+                            os.path.join("data", dir_write_name, "text"),
                             "w",
                             encoding="utf-8",
                         )
                     ]
                 for uttid, row in captions_df.iterrows():
-                    uttid = f"{data_split}_{uttid}"
+                    uttid = f"{data_split}_clotho_{uttid}"
                     file_name = row["file_name"]
                     audio_path = os.path.join(
                         ROOT_DATA_DIR, AUDIO_DIR, data_split, f"{file_name}"
@@ -64,8 +68,6 @@
                             text_i = row[f"caption_{i}"].strip()
                             print(f"{uttid} {text_i}", file=text_f[i - 1])
                     else:
-                        # TODO(sbharad2): Add options for concatenating data from multiple reference columns,
-                        # choosing particular ref column, randomly choosing col for each example.
                         text = row["caption_1"].strip()
                         print(f"{uttid} {text}", file=text_f[0])
                     N_PROCESSED += 1
diff --git a/egs2/clotho_v2/asr1/local/data_prep_clotho_chatgpt_mixup.py b/egs2/clotho_v2/asr1/local/data_prep_clotho_chatgpt_mixup.py
new file mode 100644
index 00000000000..d331cde0d46
--- /dev/null
+++ b/egs2/clotho_v2/asr1/local/data_prep_clotho_chatgpt_mixup.py
@@ -0,0 +1,449 @@
+import sys
+import os
+import json
+import random
+
+from collections import defaultdict
+from typing import Optional, List
+from string import punctuation
+from multiprocessing import Pool
+
+import librosa
+import soundfile as sf
+import pandas as pd
+import numpy as np
+from torch.utils.data import Dataset
+from transformers import AutoTokenizer
+
+
+SAMPLE_RATE = 16000
+strip_punct_table = str.maketrans("", "", punctuation)
+
+
+class ClothoMixupDataset(Dataset):
+    def __init__(
+        self,
+        audio_dir: str,
+        tokenizer_dir: str,
+        chatgpt_augmented_data_path: str,
+        chatgpt_rejected_data_path: str,
+        split_name: Optional[str] = None,
+        bypass_resample: Optional[bool] = False,
+        lemma_file: Optional[str] = None,
+        n_lemmas: Optional[int] = 2609,
+        caption_embed_dir: Optional[str] = None,
+        clap_embed_dir: Optional[str] = None,
+        do_speed_augment: Optional[bool] = False,
+        do_pitch_augment: Optional[bool] = False,
+        speed_augment_range: Optional[List[float]] = [0.8, 1.2],
+        pitch_augment_range: Optional[List[int]] = [-3, 4],
+        max_audio_len: Optional[int] = 480000,
+        do_mixup_normalize: Optional[bool] = False,
+        sample_rate: Optional[int] = SAMPLE_RATE,
+        audio_mixup_base_dir: Optional[str] = "local/clotho_audio_mixup",
+    ):
+        super().__init__()
+        self.audio_mixup_base_dir = audio_mixup_base_dir
+        self.sample_rate = sample_rate
+        self.bypass_resample = bypass_resample
+        self.audio_dir = audio_dir
+        self.split_name = split_name if split_name is not None \
+                                     else os.path.basename(audio_dir)
+
+        (
+            self.idx_to_sample,
+            self.idx_to_audio,
+            self.captions
+        ) = \
+            self.get_captions(
+                chatgpt_augmented_data_path,
+                chatgpt_rejected_data_path,
+            )
+
+        self.tokenizer = self.get_tokenizer(tokenizer_dir)
+
+        # lemma labels for encoder loss
+        if lemma_file is not None:
+            self.lemma_data = json.load(
+                open(lemma_file)
+            )
+            self.n_lemmas = n_lemmas
+        else:
+            self.lemma_data = None
+
+        # text embedding for encoder loss
+        self.caption_embed_dir = caption_embed_dir
+        self.clap_embed_dir = clap_embed_dir
+
+        # parameters for audio speed/pitch augmentation
+        self.do_speed_augment = do_speed_augment
+        self.do_pitch_augment = do_pitch_augment
+        self.speed_augment_range = speed_augment_range
+        self.pitch_augment_range = list(range(*pitch_augment_range))
+        if self.do_pitch_augment:
+            print(
+                "[info] perform pitch shifting:",
+                self.pitch_augment_range,
+                "half notes"
+            )
+        self.max_audio_len = max_audio_len
+
+        self.do_mixup_normalize = do_mixup_normalize
+
+
+    def tokenize(self, text):
+        text_id = self.tokenizer.encode(
+                                    text,
+                                    add_special_tokens=False
+                                )
+        text_id += [self.tokenizer.eos_token_id]
+        text_id = np.array(text_id, dtype=int)
+
+        return text_id
+
+
+    def get_captions(self, caption_json, rejected_json):
+        rejected_df = json.load(open(rejected_json))
+        caption_df = json.load(open(caption_json))["dataset"]
+
+        all_rejects = set()
+        seen_combs = set()
+        for rej in rejected_df:
+            all_rejects.add(tuple(rej["idx"]))
+
+        idx_to_sample = []
+        idx_to_audio = []
+        captions = defaultdict(list)
+
+        for i in range(len(caption_df)):
+            # if i >= 100:
+            #     break
+            samp_a, samp_b = caption_df[i]["selected_pair"]
+            audio_a, audio_b = caption_df[i]["audio_files"]
+
+            assert (samp_a, samp_b) not in seen_combs
+            seen_combs.add((samp_a, samp_b))
+
+            for j in range(len(caption_df[i]["chatgpt_mixups"])):
+                if (i, j) in all_rejects:
+                    continue
+
+                idx_to_sample.append((samp_a, samp_b))
+                idx_to_audio.append((audio_a, audio_b))
+
+                captions[(samp_a, samp_b)].append(
+                    caption_df[i]["chatgpt_mixups"][j]
+                        .strip()
+                        .lower()
+                        .translate(strip_punct_table)
+                )
+
+        return idx_to_sample, idx_to_audio, captions
+
+    def get_tokenizer(self, tokenizer_dir):
+        return AutoTokenizer.from_pretrained(tokenizer_dir, use_fast=True)
+
+    def load_audio(self, audio_file):
+        if not self.bypass_resample:
+            wav, sr = librosa.load(
+                audio_file,
+                sr=self.sample_rate,
+            )
+            assert sr == self.sample_rate
+        else:
+            wav, sr = sf.read(audio_file)
+            assert sr == 44100
+
+        return wav, sr
+
+
+    def mixup_audio(self, audio_a, audio_b, sample_name):
+        mixed_audio_path = os.path.join(self.audio_mixup_base_dir , f"{sample_name}.wav")
+        if not self.do_mixup_normalize:
+            energy_a = np.mean(
+                librosa.feature.rms(y=audio_a)[0] ** 2
+            )
+            energy_b = np.mean(
+                librosa.feature.rms(y=audio_b)[0] ** 2
+            )
+
+            mix_db_ratio = np.random.uniform(-5, 5)
+            mix_scale = np.sqrt(
+                energy_a /
+                (10 ** (mix_db_ratio / 10) * energy_b)
+            )
+        else:
+            import pyloudnorm as pyln
+            mix_scale = 1
+            base_db = np.random.uniform(-12, -5)
+            mix_db_ratio = np.random.uniform(-5, 5)
+            audio_a = pyln.normalize.peak(
+                audio_a,
+                base_db
+            )
+            audio_b = pyln.normalize.peak(
+                audio_b,
+                base_db + mix_db_ratio
+            )
+
+        # sf.write(
+        #     os.path.join("tmp", f"audio_{base_db:.2f}_{mix_db_ratio:.2f}_A_prenorm.wav"),
+        #     audio_a,
+        #     SAMPLE_RATE
+        # )
+        # sf.write(
+        #     os.path.join("tmp", f"audio_{base_db:.2f}_{mix_db_ratio:.2f}_B_prenorm.wav"),
+        #     audio_b,
+        #     SAMPLE_RATE
+        # )
+
+        # sf.write(
+        #     os.path.join("tmp", f"audio_{base_db:.2f}_{mix_db_ratio:.2f}_A.wav"),
+        #     audio_a,
+        #     SAMPLE_RATE
+        # )
+        # sf.write(
+        #     os.path.join("tmp", f"audio_{base_db:.2f}_{mix_db_ratio:.2f}_B.wav"),
+        #     audio_b,
+        #     SAMPLE_RATE
+        # )
+
+        if len(audio_a) > len(audio_b):
+            st_idx = random.choice(
+                range(len(audio_a) - len(audio_b) + 1)
+            )
+            audio_a[st_idx : st_idx + len(audio_b)] += \
+                mix_scale * audio_b
+            mixed_audio = audio_a
+        else:
+            st_idx = random.choice(
+                range(len(audio_b) - len(audio_a) + 1)
+            )
+            audio_b[st_idx : st_idx + len(audio_a)] += \
+                mix_scale * audio_a
+            mixed_audio = audio_b
+
+        try:
+            sf.write(
+                mixed_audio_path,
+                # os.path.join("tmp", f"{base_db:.2f}_{mix_db_ratio:.2f}_mixed.wav"),
+                mixed_audio,
+                SAMPLE_RATE
+            )
+        except:
+            print(f"Can not write audio {mixed_audio_path}")
+            mixed_audio=None
+            mixed_audio_path=None
+
+        return mixed_audio, mixed_audio_path
+
+
+    def augment_audio(self, audio):
+        if self.do_pitch_augment and random.random() < 0.5:
+            audio = librosa.effects.pitch_shift(
+                audio,
+                sr=self.sample_rate,
+                n_steps=random.choice(self.pitch_augment_range)
+            )
+
+        if self.do_speed_augment and random.random() < 0.5:
+            if audio.shape[0] / self.speed_augment_range[0] > self.max_audio_len:
+                _min_speed = audio.shape[0] / self.max_audio_len
+            else:
+                _min_speed = self.speed_augment_range[0]
+
+            assert audio.shape[0] / _min_speed <= self.max_audio_len + 1
+
+            audio = librosa.effects.time_stretch(
+                audio,
+                rate=np.random.uniform(
+                    _min_speed,
+                    self.speed_augment_range[1]
+                )
+            )
+
+        return audio
+
+
+    def read_lemma_labels(self, sample_name, sample_idx):
+        lemmas = self.lemma_data[sample_name][sample_idx]["caption_kw_labels"]
+        labels = np.zeros((self.n_lemmas,))
+
+        for l in lemmas:
+            labels[l] = 1
+
+        return labels
+
+
+    def read_caption_embed(self, sample_pair, sample_idx):
+        sample_embed_dir = os.path.join(
+            self.caption_embed_dir,
+            str(sample_pair[0]) + "_" + str(sample_pair[1])
+        )
+        embed_choices = [
+            f for f in os.listdir(sample_embed_dir) \
+            if f"chatgpt{sample_idx + 1:>02d}" in f
+        ]
+        # print(len(embed_choices))
+
+        return np.load(
+            os.path.join(
+                sample_embed_dir,
+                random.choice(embed_choices)
+            )
+        )
+
+
+    def read_clap_embed(self, sample_pair, sample_idx):
+        sample_embed_dir = os.path.join(
+            self.clap_embed_dir,
+            str(sample_pair[0]) + "_" + str(sample_pair[1])
+        )
+        embed_choices = [
+            f for f in os.listdir(sample_embed_dir) \
+            if f"chatgpt{sample_idx + 1:>02d}" in f
+        ]
+        # print(len(embed_choices))
+
+        return np.load(
+            os.path.join(
+                sample_embed_dir,
+                random.choice(embed_choices)
+            )
+        )
+
+
+    def __len__(self):
+        return len(self.idx_to_sample)
+        # return 10
+
+
+    def __getitem__(self, index):
+        samp_a, samp_b = self.idx_to_sample[index]
+        audio_a, audio_b = self.idx_to_audio[index]
+        audios_to_mix = []
+        try:
+            for audfile in [audio_a, audio_b]:
+                audio, sr = self.load_audio(
+                    os.path.join(
+                        self.audio_dir,
+                        audfile
+                    )
+                )
+
+                audios_to_mix.append(audio)
+        except OSError as e:
+            print('Skipping', samp_a, samp_b, 'because' ,e)
+            return None
+
+        audio,audio_path = self.mixup_audio(
+                        audios_to_mix[0],
+                        audios_to_mix[1],
+                        sample_name=str(samp_a) + "_" + str(samp_b),
+                    )
+
+        if (self.do_pitch_augment or self.do_speed_augment) and \
+        random.random() < 0.5:
+            audio = self.augment_audio(audio)
+
+        caption_idx = random.choice(
+                        range(len(self.captions[(samp_a, samp_b)]))
+                    )
+
+        caption = self.captions[(samp_a, samp_b)][caption_idx]
+        # caption = self.tokenize(caption)
+
+        bundle = {
+            "sample_name": str(samp_a) + "_" + str(samp_b),
+            "sr": sr,
+            "audio_path": audio_path,
+            "labels": caption,
+            # "encoder_input": audio,
+            # "attention_mask": np.full_like(audio, False, dtype=bool)
+        }
+
+        if self.lemma_data is not None:
+            encoder_labels = self.read_lemma_labels(
+                                    (samp_a, samp_b),
+                                    caption_idx
+                                )
+            bundle["encoder_labels"] = encoder_labels
+        elif self.caption_embed_dir is not None:
+            encoder_embed = self.read_caption_embed(
+                                    (samp_a, samp_b),
+                                    caption_idx
+                                )
+
+            if self.clap_embed_dir is not None:
+                clap_embed = self.read_clap_embed(
+                    (samp_a, samp_b),
+                    caption_idx,
+                )
+                encoder_embed = np.concatenate(
+                    [encoder_embed, clap_embed],
+                    axis=-1
+                )
+
+            bundle["encoder_labels"] = encoder_embed
+        else:
+            bundle["encoder_labels"] = np.array([0])
+
+        return bundle
+
+
+# create directory paths for storing clotho mixup audio
+audio_mixup_write_dir = "downloads/CLOTHO_v2.1_audio_mixup/development/"
+os.makedirs(os.path.dirname(audio_mixup_write_dir), exist_ok=True)
+
+mixup_captions_root_dir=sys.argv[1]
+clotho_audio_base_dir=sys.argv[2]
+# Create mixup dataset
+dset = ClothoMixupDataset(
+    audio_dir=os.path.join(clotho_audio_base_dir, "clotho_audio_files/development"),
+    tokenizer_dir="facebook/bart-base",
+    chatgpt_augmented_data_path=os.path.join(mixup_captions_root_dir,"clotho_development_chatgpt_mixups.json"),
+    chatgpt_rejected_data_path=os.path.join(mixup_captions_root_dir,"clotho_development_chatgpt_mixups_err.json"),
+    # caption_embed_dir="sent_embedding/datasets/instructor-xl/clotho/development_chatgpt_mixup",
+    # clap_embed_dir="sent_embedding/datasets/clap/clotho/development_chatgpt_mixup",
+    audio_mixup_base_dir=audio_mixup_write_dir
+)
+
+data_base="data/development_clotho_chatgpt_mixup/"
+output_txt_path=os.path.join(data_base, "text")
+output_wav_scp_path=os.path.join(data_base, "wav.scp")
+output_utt2spk_path=os.path.join(data_base, "utt2spk")
+
+os.makedirs(os.path.dirname(output_txt_path), exist_ok=True)
+import time
+start = time.time()
+def _get_ith_item(i):
+    data_i=dset[i]
+    if data_i is None:
+        return None, None, None
+    uttid=f"clotho_chatmixup_{data_i['sample_name']}"
+    return uttid, data_i['audio_path'], data_i['labels']
+
+results=None
+with Pool() as pool: # Use all cores
+    results = pool.map(_get_ith_item, range(len(dset)))
+
+assert results is not None
+
+N_ERROR=0
+N_PROCESSED=0
+with open(output_txt_path, 'w') as text_f, open(output_wav_scp_path,'w') as wav_scp_f, open(output_utt2spk_path,'w') as utt2spk_f:
+    for i in range(len(dset)):
+        uttid, audio_path, text = results[i]
+        if uttid is None:
+            N_ERROR+=1
+            continue
+        print(f"{uttid} {audio_path}", file=wav_scp_f)
+        print(f"{uttid} {text}", file=text_f)
+        print(f"{uttid} dummy", file=utt2spk_f)
+        N_PROCESSED+=1
+        if (i+1)%1000==0:
+            print(f"Processed {i+1} mixup samples.")
+
+print(f'Number of files skipped {N_ERROR}')
+print(f'Number of files processed {N_PROCESSED}')
+#TODO(shikhar): change location of downloaded data to downloads, check gitignore for more info.
\ No newline at end of file
diff --git a/egs2/clotho_v2/asr1/local/evaluation.py b/egs2/clotho_v2/asr1/local/evaluation.py
index 158fe6e9d5d..88d0d63c22e 100644
--- a/egs2/clotho_v2/asr1/local/evaluation.py
+++ b/egs2/clotho_v2/asr1/local/evaluation.py
@@ -74,12 +74,12 @@ def _evaluate(decode_fname: str, split: str):
                 references[key] = [value]
 
     candidates, mult_references = _convert_to_evaluation_format(predictions, references)
-    candidates = candidates[:10]
-    mult_references = mult_references[:10]  # TODO(sbharad2): Only for testing
-    evaluate = Evaluate(metrics=["spider", "fense", "vocab"])
+    
+    evaluate = Evaluate(metrics=["spider", "fense", "meteor", "rouge_l", "spice", "spider_fl", "cider_d"])
     corpus_scores, _ = evaluate(candidates, mult_references)
 
-    header_str = f" Split: {split} Evaluating over {len(candidates)} predictions. "
+
+    header_str = f" Split: {split} Evaluation over {len(candidates)} predictions. "
     header_str_l = len(header_str)
     print("=" * header_str_l)
     print(header_str)
@@ -88,5 +88,16 @@ def _evaluate(decode_fname: str, split: str):
         print(f" {metric:<20}: {value} ")
     print("=" * header_str_l)
 
+    # Write results in a file
+    results_file = decode_fname+'.result'
+    with open(results_file, 'w') as result_f:
+        header_str = f" Split: {split} Evaluation over {len(candidates)} predictions. "
+        header_str_l = len(header_str)
+        print("=" * header_str_l, file=result_f)
+        print(header_str, file=result_f)
+        print("=" * header_str_l, file=result_f)
+        for metric, value in corpus_scores.items():
+            print(f" {metric:<20}: {value} ", file=result_f)
+        print("=" * header_str_l, file=result_f)
 
 _evaluate(args.decode_file, args.split)
diff --git a/egs2/clotho_v2/asr1/local/score.sh b/egs2/clotho_v2/asr1/local/score.sh
index 06b2d1751c8..147604a16ff 100755
--- a/egs2/clotho_v2/asr1/local/score.sh
+++ b/egs2/clotho_v2/asr1/local/score.sh
@@ -1,5 +1,5 @@
 #!/usr/bin/env bash
-
+set -e
 
 [ -f ./path.sh ] && . ./path.sh
 . parse_options.sh || exit 1;
@@ -14,9 +14,8 @@ pip install aac-metrics
 aac-metrics-download
 
 asr_expdir=$1
-splits=(evaluation validation)
+splits=(evaluation)
 
-set -ex
 for split in ${splits[@]}; do
 	for decode_file in $((ls -d ${asr_expdir}/*/*/* && ls -d ${asr_expdir}/*/*) | grep "${split}/text"); do
 		echo "Decode file: ${decode_file} for split:${split}"
diff --git a/egs2/clotho_v2/asr1/run.sh b/egs2/clotho_v2/asr1/run.sh
deleted file mode 100755
index 579124dd774..00000000000
--- a/egs2/clotho_v2/asr1/run.sh
+++ /dev/null
@@ -1,32 +0,0 @@
-#!/usr/bin/env bash
-
-set -euo pipefail
-
-timestamp=$(date "+%Y%m%d.%H%M%S")
-
-wandb_init_args="--use_wandb true --wandb_project DCASE_AAC --wandb_model_log_interval 0"
-wandb_sweep_args="$@"
-
-timestamp=20240928.224746
-
-./asr.sh \
-    --asr_tag ${timestamp} \
-    --feats_normalize uttmvn \
-    --stage 1 \
-    --stop_stage 12 \
-    --ngpu 1 \
-    --gpu_inference true \
-    --nj 10 \
-    --inference_nj 1 \
-    --max_wav_duration 30 \
-    --token_type hugging_face \
-    --use_lm false \
-    --hugging_face_model_name_or_path "facebook/bart-base" \
-    --inference_args "--beam_size 10 --ctc_weight 0.0 --hugging_face_decoder True" \
-    --train_set development \
-    --valid_set validation \
-    --test_sets "validation evaluation" \
-    --asr_config conf/dcase.yaml \
-    --inference_asr_model 68epoch.pth \
-    --asr_args "${wandb_init_args} ${wandb_sweep_args}"
-    # --local_score_opts exp/asr_branchformer_raw_en_word/inference_beam_size10_ctc_weight0.3_asr_model_valid.acc.ave \
diff --git a/egs2/clotho_v2/asr1/run_ft.sh b/egs2/clotho_v2/asr1/run_ft.sh
new file mode 100755
index 00000000000..a012172175e
--- /dev/null
+++ b/egs2/clotho_v2/asr1/run_ft.sh
@@ -0,0 +1,33 @@
+#!/usr/bin/env bash
+set -euo pipefail
+# Runs fine-tuning
+
+# wandb_init_args="--use_wandb true --wandb_project DCASE_AAC --wandb_model_log_interval 0"
+wandb_init_args=""
+other_args="$@"
+
+pre_trained_model_path=exp/asr_pt.initfix.bigbatch512.lr2e-4.weighted12layers.20241103.145125/valid.acc.ave_5best.pth
+
+./asr.sh \
+    --asr_tag ft_lr5e-5.initfix.bigbatch512.lr2e-4.weighted12layers.20241103.145125 \
+    --feats_normalize uttmvn \
+    --stage 11 \
+    --stop_stage 13 \
+    --asr_stats_dir exp/asr_stats_finetune \
+    --ngpu 2 \
+    --gpu_inference true \
+    --nj 20 \
+    --inference_nj 1 \
+    --max_wav_duration 30 \
+    --token_type hugging_face \
+    --use_lm false \
+    --hugging_face_model_name_or_path "facebook/bart-base" \
+    --inference_args "--beam_size 10 --ctc_weight 0.0 --hugging_face_decoder True" \
+    --train_set development_clotho \
+    --valid_set validation \
+    --test_sets "evaluation" \
+    --asr_config conf/beats_bart_ft.yaml \
+    --pretrained_model ${pre_trained_model_path} \
+    --inference_asr_model valid.acc.ave_5best.pth \
+    --asr_args "${wandb_init_args} ${other_args}" \
+    --local_score_opts exp/asr_ft.initfix.bigbatch512.lr2e-4.weighted12layers.20241103.145125/inference_beam_size10_ctc_weight0.0_hugging_face_decoderTrue_asr_model_valid.acc.ave_5best
diff --git a/egs2/clotho_v2/asr1/run_pt.sh b/egs2/clotho_v2/asr1/run_pt.sh
new file mode 100755
index 00000000000..f64472baefa
--- /dev/null
+++ b/egs2/clotho_v2/asr1/run_pt.sh
@@ -0,0 +1,70 @@
+#!/usr/bin/env bash
+# Runs pre-training
+set -euo pipefail
+
+timestamp=$(date "+%Y%m%d.%H%M%S")
+
+# wandb_init_args="--use_wandb true --wandb_project DCASE_AAC --wandb_model_log_interval 0"
+wandb_init_args=""
+other_args="$@"
+
+./asr.sh \
+    --asr_tag pt.initfix.bigbatch512.lr2e-4.weighted12layers.${timestamp} \
+    --asr_speech_fold_length 1600 \
+    --feats_normalize uttmvn \
+    --stage 11 \
+    --stop_stage 13 \
+    --ngpu 2 \
+    --gpu_inference true \
+    --nj 20 \
+    --inference_nj 1 \
+    --max_wav_duration 30 \
+    --token_type hugging_face \
+    --use_lm false \
+    --hugging_face_model_name_or_path "facebook/bart-base" \
+    --inference_args "--beam_size 10 --ctc_weight 0.0 --hugging_face_decoder True" \
+    --train_set pretrain \
+    --valid_set validation \
+    --test_sets "validation evaluation" \
+    --asr_config conf/beats_bart_pt.yaml \
+    --inference_asr_model valid.acc.best.pth \
+    --asr_args "${wandb_init_args} ${other_args}" \
+    --local_score_opts "exp/asr_pt.${timestamp}/inference_beam_size10_ctc_weight0.0_hugging_face_decoderTrue_asr_model_valid.acc.best"
+
+
+
+# # Run inference on all checkpoints
+
+# wandb_sweep_args="$@"
+
+# name_tag=ft.initfix.bigbatch512.lr2e-4.weighted12layers.20241103.145125
+
+# all_ckpts=$(find exp/asr_${name_tag} -name "*.pth"  -printf "%f\n" | awk '{print substr($0, 1, length($0) - 4)}')
+
+# # for ckpt in valid.acc.best; do
+# for ckpt in $all_ckpts ; do
+#     if [[ $ckpt == "checkpoint" || $ckpt == "latest"  || $ckpt == "valid.acc.best" ]]; then
+#         continue
+#     fi
+#     ./asr.sh \
+#         --asr_tag $name_tag \
+#         --feats_normalize uttmvn \
+#         --stage 12 \
+#         --stop_stage 13 \
+#         --ngpu 1 \
+#         --gpu_inference true \
+#         --nj 20 \
+#         --inference_nj 1 \
+#         --max_wav_duration 30 \
+#         --token_type hugging_face \
+#         --use_lm false \
+#         --hugging_face_model_name_or_path "facebook/bart-base" \
+#         --inference_args "--beam_size 10 --ctc_weight 0.0 --hugging_face_decoder True" \
+#         --train_set pretrain \
+#         --valid_set validation \
+#         --test_sets "evaluation" \
+#         --asr_config conf/dcase.yaml \
+#         --inference_asr_model ${ckpt}.pth \
+#         --asr_args "${wandb_sweep_args}" \
+#         --local_score_opts "exp/asr_${name_tag}/inference_beam_size10_ctc_weight0.0_hugging_face_decoderTrue_asr_model_${ckpt}"
+# done
diff --git a/espnet2/asr/encoder/beats_encoder.py b/espnet2/asr/encoder/beats_encoder.py
index 67415a9fb65..a6c60cc6e7b 100644
--- a/espnet2/asr/encoder/beats_encoder.py
+++ b/espnet2/asr/encoder/beats_encoder.py
@@ -14,6 +14,7 @@
 import math
 import warnings
 from typing import Dict, Optional, Tuple
+from copy import deepcopy
 
 import numpy as np
 import torch
@@ -102,10 +103,11 @@ class BeatsEncoder(AbsEncoder):
 
     (https://arxiv.org/abs/2212.09058)
     Args:
-        input_size: Input size for xs_pad in `forward`.
         beats_ckpt_path: Path to a pretrained BEATs checkpoint. If `beats_config` is provided and it does
             not match the config in the checkpoint, code might throw an error.
-        max_layer: Use all layers for encoding if None. Otherwise use upto `max_layer`.
+        adapter_layers: Number of transformer-based adapter layers after encoder.
+        max_layer: Propagate input through all layers for encoding if None. Otherwise use upto `max_layer`.
+        use_weighted_representation: Use weighted representations from max_layer if True. Weights are randomly initialized.
         beats_config: `BEATsConfig` object. If provided, we will try to override the config in the
             checkpoint. This can be used to change dropouts etc for fine-tuning the model while starting
             from a pretrained checkpoint.
@@ -116,8 +118,9 @@ def __init__(
         self,
         input_size: int,
         beats_ckpt_path: str = None,
+        adapter_layers: int = 0,
         max_layer: int = None,
-        use_all_layers: bool = False,
+        use_weighted_representation: bool = False,
         beats_config: Optional[BEATsConfig] = None,
         specaug_config: Optional[Dict] = None,
     ) -> None:
@@ -170,14 +173,24 @@ def __init__(
         )
         self.dropout_input = nn.Dropout(config.dropout_input)
         assert not config.deep_norm or not config.layer_norm_first
+        # Add adapter layer to config.
         self.encoder = TransformerEncoder(config)
         self.layer_norm = LayerNorm(self.embed)
-        self.use_all_layers = use_all_layers
-        if self.use_all_layers:
-            assert self.max_layer is not None, "max_layer must be provided."
-            self.layer_weights = nn.Parameter(torch.ones(self.max_layer))
-
-    def espnet_initialization_fn(self):
+        self.use_weighted_representation = use_weighted_representation
+        if self.use_weighted_representation:
+            if self.max_layer is None:
+                logging.warning(f"max_layer must be provided when using weighted representations. Set to {config.encoder_layers}.")
+                self.max_layer = config.encoder_layers
+            self.layer_weights = nn.Parameter(torch.ones(self.max_layer)/self.max_layer, requires_grad=True) # Uniform init
+        
+        self.transformer_adapter = None
+        if adapter_layers !=0:
+            adapter_config = deepcopy(config)
+            adapter_config.encoder_layers=adapter_layers
+            self.transformer_adapter = TransformerEncoder(adapter_config)
+
+
+    def reload_pretrained_parameters(self):
         """Initialization function for BEATs.
 
         This must be called last in the initialization procedure.
@@ -198,19 +211,6 @@ def espnet_initialization_fn(self):
         # BEATs has different initialization from ESPNet for other modules, so override.
         self.encoder.apply(init_bert_params)
         if self.loaded_state_dict_ is not None:
-            # TODO(shikhar): Hardcoded
-            all_keys = list(self.loaded_state_dict_["model"].items())
-            for key, _ in all_keys:
-                if (
-                    key.startswith("encoder.layers.11")
-                    # or key.startswith("encoder.layers.10")
-                    # or key.startswith("encoder.layers.9")
-                    # or key.startswith("encoder.layers.8")
-                    # or key.startswith("encoder.layers.7")
-                    # or key.startswith("encoder.layers.6")
-                ):
-                    self.loaded_state_dict_["model"].pop(key)
-                    logging.info(f"Removed key: {key}")
 
             load_info = self.load_state_dict(
                 self.loaded_state_dict_["model"], strict=False
@@ -219,7 +219,7 @@ def espnet_initialization_fn(self):
             logging.info(
                 f"Loaded BEATs pretrained model. Following keys were missing in your custom model: {load_info.missing_keys}. "
                 f"Follwing keys could not be loaded from the pretrained checkpoint: {load_info.unexpected_keys}."
-                "It is expected to have 'predictor' here if you are fine-tuning with only the BEATs backbone."
+                "It is expected to have 'predictor' listed above if you are fine-tuning with only the BEATs backbone."
             )
 
     def forward_padding_mask(
@@ -275,7 +275,7 @@ def forward(
         # Adjust shapes to be compatible with BEATs code
         xs_pad, mask = xs_pad.squeeze(-1).squeeze(-1), mask.squeeze(-1).squeeze(-1)
         # masks = None
-        last_hidden, _, mask = self.extract_features(
+        last_hidden, mask = self.extract_features(
             xs_pad,
             mask,
             max_layer=self.max_layer,
@@ -313,18 +313,24 @@ def extract_features(
         if self.post_extract_proj is not None:
             features = self.post_extract_proj(features)
 
-        x = self.dropout_input(features)
-        x, layer_results = self.encoder(x, padding_mask=padding_mask, layer=max_layer)
+        features = self.dropout_input(features)
+        features, layer_results = self.encoder(features, padding_mask=padding_mask, layer=max_layer)
 
-        if self.use_all_layers:
+        if self.use_weighted_representation:
             for i, (x_i, _) in enumerate(layer_results):
-                if i > self.max_layer:
+                if i >= max_layer:
                     break
                 if i == 0:
-                    x = self.layer_weights[i] * x_i
-                x += self.layer_weights[i] * x_i
-
-        return x, layer_results, padding_mask
+                    features = self.layer_weights[i] * x_i
+                else:
+                    features += self.layer_weights[i] * x_i
+            # T x B x C -> B x T x C
+            features = features.transpose(0, 1)
+        
+        if self.transformer_adapter:
+            features, _ = self.transformer_adapter(features, padding_mask = padding_mask, layer = None)
+
+        return features, padding_mask
 
 
 class TransformerEncoder(nn.Module):

From 0425cfb218d3c29005d3f53a68cd088de5802c55 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Tue, 5 Nov 2024 15:08:52 +0000
Subject: [PATCH 09/47] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 egs2/clotho_v2/asr1/conf/beats_bart_ft.yaml   |   2 +-
 egs2/clotho_v2/asr1/conf/beats_bart_pt.yaml   |   2 +-
 .../asr1/local/data_prep_audiocaps.py         |  53 ++--
 egs2/clotho_v2/asr1/local/data_prep_clotho.py |   6 +-
 .../local/data_prep_clotho_chatgpt_mixup.py   | 239 +++++++-----------
 egs2/clotho_v2/asr1/local/evaluation.py       |  20 +-
 espnet2/asr/encoder/beats_encoder.py          |  27 +-
 7 files changed, 156 insertions(+), 193 deletions(-)

diff --git a/egs2/clotho_v2/asr1/conf/beats_bart_ft.yaml b/egs2/clotho_v2/asr1/conf/beats_bart_ft.yaml
index dd4e1fcdbcd..d95bb1d0939 100644
--- a/egs2/clotho_v2/asr1/conf/beats_bart_ft.yaml
+++ b/egs2/clotho_v2/asr1/conf/beats_bart_ft.yaml
@@ -44,7 +44,7 @@ encoder_conf:
     adapter_layers: 2
     max_layer: 12
     use_weighted_representation: true
-    
+
 decoder: hugging_face_transformers
 decoder_conf:
     model_name_or_path: facebook/bart-base
diff --git a/egs2/clotho_v2/asr1/conf/beats_bart_pt.yaml b/egs2/clotho_v2/asr1/conf/beats_bart_pt.yaml
index 74a65947c81..86862c04de3 100644
--- a/egs2/clotho_v2/asr1/conf/beats_bart_pt.yaml
+++ b/egs2/clotho_v2/asr1/conf/beats_bart_pt.yaml
@@ -46,7 +46,7 @@ encoder_conf:
     adapter_layers: 2
     max_layer: 12
     use_weighted_representation: true
-    
+
 decoder: hugging_face_transformers
 decoder_conf:
     model_name_or_path: facebook/bart-base
diff --git a/egs2/clotho_v2/asr1/local/data_prep_audiocaps.py b/egs2/clotho_v2/asr1/local/data_prep_audiocaps.py
index 93b3938ec06..822aa1d4d64 100644
--- a/egs2/clotho_v2/asr1/local/data_prep_audiocaps.py
+++ b/egs2/clotho_v2/asr1/local/data_prep_audiocaps.py
@@ -1,48 +1,53 @@
 """Prepares data for AudioCaps dataset."""
 
+import json
 import os
 import sys
-import json
+
 from scipy.io import wavfile
 
 if __name__ == "__main__":
     ROOT_DATA_DIR = sys.argv[1]
-    TRAIN_FILES = "train_split.json" # this is a json with key "audiocap_ids" containing list of folder_ids 
-    DATA_DIR = "data" # this has folders each for a folder_id, containing audio.wav and metadata.json files
-    
+    TRAIN_FILES = "train_split.json"  # this is a json with key "audiocap_ids" containing list of folder_ids
+    DATA_DIR = "data"  # this has folders each for a folder_id, containing audio.wav and metadata.json files
+
     LOCAL_DATA_DIR = "data/development_audiocaps"
 
     all_train_ids = None
-    with open(os.path.join(ROOT_DATA_DIR,TRAIN_FILES)) as f:
+    with open(os.path.join(ROOT_DATA_DIR, TRAIN_FILES)) as f:
         train_meta = json.load(f)
-        all_train_ids = train_meta['audiocap_ids']
-    assert all_train_ids is not None, f'Could not read {ROOT_DATA_DIR/TRAIN_FILES}'
+        all_train_ids = train_meta["audiocap_ids"]
+    assert all_train_ids is not None, f"Could not read {ROOT_DATA_DIR/TRAIN_FILES}"
 
-    output_wav_scp_path= os.path.join(LOCAL_DATA_DIR, 'wav.scp')
-    output_utt2spk_path= os.path.join(LOCAL_DATA_DIR, 'utt2spk')
-    output_txt_path = os.path.join(LOCAL_DATA_DIR, 'text')
+    output_wav_scp_path = os.path.join(LOCAL_DATA_DIR, "wav.scp")
+    output_utt2spk_path = os.path.join(LOCAL_DATA_DIR, "utt2spk")
+    output_txt_path = os.path.join(LOCAL_DATA_DIR, "text")
     missing_audio = []
     N_PROCESSED = 0
     N_ZERO_LENGTH_CAPTIONS = 0
-    N_ZERO_LENGTH_AUDIO =0 
+    N_ZERO_LENGTH_AUDIO = 0
     N_VERY_LONG_AUDIO = 0
-    AUDIO_LENGTH_THRESHOLD = 30 # seconds
-    with open(output_txt_path, 'w') as text_f, open(output_wav_scp_path,'w') as wav_scp_f, open(output_utt2spk_path,'w') as utt2spk_f:
+    AUDIO_LENGTH_THRESHOLD = 30  # seconds
+    with open(output_txt_path, "w") as text_f, open(
+        output_wav_scp_path, "w"
+    ) as wav_scp_f, open(output_utt2spk_path, "w") as utt2spk_f:
         for train_id in all_train_ids:
-            uttid = f'audiocaps_train_{train_id}'
+            uttid = f"audiocaps_train_{train_id}"
             audio_path = os.path.join(ROOT_DATA_DIR, DATA_DIR, train_id, "audio.wav")
-            samplerate, data = wavfile.read(audio_path) 
-            audio_length = len(data)/samplerate
-            if audio_length==0:
-                N_ZERO_LENGTH_AUDIO+=1
+            samplerate, data = wavfile.read(audio_path)
+            audio_length = len(data) / samplerate
+            if audio_length == 0:
+                N_ZERO_LENGTH_AUDIO += 1
                 continue
             if audio_length > AUDIO_LENGTH_THRESHOLD:
-                N_VERY_LONG_AUDIO+=1
+                N_VERY_LONG_AUDIO += 1
                 continue
-            with open(os.path.join(ROOT_DATA_DIR, DATA_DIR, train_id, "metadata.json")) as f:
+            with open(
+                os.path.join(ROOT_DATA_DIR, DATA_DIR, train_id, "metadata.json")
+            ) as f:
                 meta = json.load(f)
-                text = meta['caption'].strip()
-                if len(text)==0:
+                text = meta["caption"].strip()
+                if len(text) == 0:
                     N_ZERO_LENGTH_CAPTIONS += 1
                     continue
                 print(f"{uttid} {audio_path}", file=wav_scp_f)
@@ -51,4 +56,6 @@
                 N_PROCESSED += 1
                 if N_PROCESSED % 1000 == 0:
                     print(f"Processed {N_PROCESSED} audio files.")
-    print(f"Processed {N_PROCESSED} audio files. Skipped {N_VERY_LONG_AUDIO} long audio, {N_ZERO_LENGTH_AUDIO} empty audio and {N_ZERO_LENGTH_CAPTIONS} empty caption items.")
\ No newline at end of file
+    print(
+        f"Processed {N_PROCESSED} audio files. Skipped {N_VERY_LONG_AUDIO} long audio, {N_ZERO_LENGTH_AUDIO} empty audio and {N_ZERO_LENGTH_CAPTIONS} empty caption items."
+    )
diff --git a/egs2/clotho_v2/asr1/local/data_prep_clotho.py b/egs2/clotho_v2/asr1/local/data_prep_clotho.py
index d9e8b92d033..d3d29391802 100644
--- a/egs2/clotho_v2/asr1/local/data_prep_clotho.py
+++ b/egs2/clotho_v2/asr1/local/data_prep_clotho.py
@@ -23,10 +23,10 @@
         )
         captions_df = pd.read_csv(all_captions_path)
         N_PROCESSED = 0
-        if data_split=='development':
-            dir_write_name = 'development_clotho'
+        if data_split == "development":
+            dir_write_name = "development_clotho"
         else:
-            dir_write_name = data_split # for validation and evaluation
+            dir_write_name = data_split  # for validation and evaluation
         with open(
             os.path.join("data", dir_write_name, "wav.scp"), "w", encoding="utf-8"
         ) as wav_scp_f, open(
diff --git a/egs2/clotho_v2/asr1/local/data_prep_clotho_chatgpt_mixup.py b/egs2/clotho_v2/asr1/local/data_prep_clotho_chatgpt_mixup.py
index d331cde0d46..1be84bfda95 100644
--- a/egs2/clotho_v2/asr1/local/data_prep_clotho_chatgpt_mixup.py
+++ b/egs2/clotho_v2/asr1/local/data_prep_clotho_chatgpt_mixup.py
@@ -1,21 +1,19 @@
-import sys
-import os
 import json
+import os
 import random
-
+import sys
 from collections import defaultdict
-from typing import Optional, List
-from string import punctuation
 from multiprocessing import Pool
+from string import punctuation
+from typing import List, Optional
 
 import librosa
-import soundfile as sf
-import pandas as pd
 import numpy as np
+import pandas as pd
+import soundfile as sf
 from torch.utils.data import Dataset
 from transformers import AutoTokenizer
 
-
 SAMPLE_RATE = 16000
 strip_punct_table = str.maketrans("", "", punctuation)
 
@@ -47,26 +45,20 @@ def __init__(
         self.sample_rate = sample_rate
         self.bypass_resample = bypass_resample
         self.audio_dir = audio_dir
-        self.split_name = split_name if split_name is not None \
-                                     else os.path.basename(audio_dir)
-
-        (
-            self.idx_to_sample,
-            self.idx_to_audio,
-            self.captions
-        ) = \
-            self.get_captions(
-                chatgpt_augmented_data_path,
-                chatgpt_rejected_data_path,
-            )
+        self.split_name = (
+            split_name if split_name is not None else os.path.basename(audio_dir)
+        )
+
+        (self.idx_to_sample, self.idx_to_audio, self.captions) = self.get_captions(
+            chatgpt_augmented_data_path,
+            chatgpt_rejected_data_path,
+        )
 
         self.tokenizer = self.get_tokenizer(tokenizer_dir)
 
         # lemma labels for encoder loss
         if lemma_file is not None:
-            self.lemma_data = json.load(
-                open(lemma_file)
-            )
+            self.lemma_data = json.load(open(lemma_file))
             self.n_lemmas = n_lemmas
         else:
             self.lemma_data = None
@@ -82,26 +74,19 @@ def __init__(
         self.pitch_augment_range = list(range(*pitch_augment_range))
         if self.do_pitch_augment:
             print(
-                "[info] perform pitch shifting:",
-                self.pitch_augment_range,
-                "half notes"
+                "[info] perform pitch shifting:", self.pitch_augment_range, "half notes"
             )
         self.max_audio_len = max_audio_len
 
         self.do_mixup_normalize = do_mixup_normalize
 
-
     def tokenize(self, text):
-        text_id = self.tokenizer.encode(
-                                    text,
-                                    add_special_tokens=False
-                                )
+        text_id = self.tokenizer.encode(text, add_special_tokens=False)
         text_id += [self.tokenizer.eos_token_id]
         text_id = np.array(text_id, dtype=int)
 
         return text_id
 
-
     def get_captions(self, caption_json, rejected_json):
         rejected_df = json.load(open(rejected_json))
         caption_df = json.load(open(caption_json))["dataset"]
@@ -133,9 +118,9 @@ def get_captions(self, caption_json, rejected_json):
 
                 captions[(samp_a, samp_b)].append(
                     caption_df[i]["chatgpt_mixups"][j]
-                        .strip()
-                        .lower()
-                        .translate(strip_punct_table)
+                    .strip()
+                    .lower()
+                    .translate(strip_punct_table)
                 )
 
         return idx_to_sample, idx_to_audio, captions
@@ -156,35 +141,22 @@ def load_audio(self, audio_file):
 
         return wav, sr
 
-
     def mixup_audio(self, audio_a, audio_b, sample_name):
-        mixed_audio_path = os.path.join(self.audio_mixup_base_dir , f"{sample_name}.wav")
+        mixed_audio_path = os.path.join(self.audio_mixup_base_dir, f"{sample_name}.wav")
         if not self.do_mixup_normalize:
-            energy_a = np.mean(
-                librosa.feature.rms(y=audio_a)[0] ** 2
-            )
-            energy_b = np.mean(
-                librosa.feature.rms(y=audio_b)[0] ** 2
-            )
+            energy_a = np.mean(librosa.feature.rms(y=audio_a)[0] ** 2)
+            energy_b = np.mean(librosa.feature.rms(y=audio_b)[0] ** 2)
 
             mix_db_ratio = np.random.uniform(-5, 5)
-            mix_scale = np.sqrt(
-                energy_a /
-                (10 ** (mix_db_ratio / 10) * energy_b)
-            )
+            mix_scale = np.sqrt(energy_a / (10 ** (mix_db_ratio / 10) * energy_b))
         else:
             import pyloudnorm as pyln
+
             mix_scale = 1
             base_db = np.random.uniform(-12, -5)
             mix_db_ratio = np.random.uniform(-5, 5)
-            audio_a = pyln.normalize.peak(
-                audio_a,
-                base_db
-            )
-            audio_b = pyln.normalize.peak(
-                audio_b,
-                base_db + mix_db_ratio
-            )
+            audio_a = pyln.normalize.peak(audio_a, base_db)
+            audio_b = pyln.normalize.peak(audio_b, base_db + mix_db_ratio)
 
         # sf.write(
         #     os.path.join("tmp", f"audio_{base_db:.2f}_{mix_db_ratio:.2f}_A_prenorm.wav"),
@@ -209,18 +181,12 @@ def mixup_audio(self, audio_a, audio_b, sample_name):
         # )
 
         if len(audio_a) > len(audio_b):
-            st_idx = random.choice(
-                range(len(audio_a) - len(audio_b) + 1)
-            )
-            audio_a[st_idx : st_idx + len(audio_b)] += \
-                mix_scale * audio_b
+            st_idx = random.choice(range(len(audio_a) - len(audio_b) + 1))
+            audio_a[st_idx : st_idx + len(audio_b)] += mix_scale * audio_b
             mixed_audio = audio_a
         else:
-            st_idx = random.choice(
-                range(len(audio_b) - len(audio_a) + 1)
-            )
-            audio_b[st_idx : st_idx + len(audio_a)] += \
-                mix_scale * audio_a
+            st_idx = random.choice(range(len(audio_b) - len(audio_a) + 1))
+            audio_b[st_idx : st_idx + len(audio_a)] += mix_scale * audio_a
             mixed_audio = audio_b
 
         try:
@@ -228,22 +194,21 @@ def mixup_audio(self, audio_a, audio_b, sample_name):
                 mixed_audio_path,
                 # os.path.join("tmp", f"{base_db:.2f}_{mix_db_ratio:.2f}_mixed.wav"),
                 mixed_audio,
-                SAMPLE_RATE
+                SAMPLE_RATE,
             )
         except:
             print(f"Can not write audio {mixed_audio_path}")
-            mixed_audio=None
-            mixed_audio_path=None
+            mixed_audio = None
+            mixed_audio_path = None
 
         return mixed_audio, mixed_audio_path
 
-
     def augment_audio(self, audio):
         if self.do_pitch_augment and random.random() < 0.5:
             audio = librosa.effects.pitch_shift(
                 audio,
                 sr=self.sample_rate,
-                n_steps=random.choice(self.pitch_augment_range)
+                n_steps=random.choice(self.pitch_augment_range),
             )
 
         if self.do_speed_augment and random.random() < 0.5:
@@ -255,16 +220,11 @@ def augment_audio(self, audio):
             assert audio.shape[0] / _min_speed <= self.max_audio_len + 1
 
             audio = librosa.effects.time_stretch(
-                audio,
-                rate=np.random.uniform(
-                    _min_speed,
-                    self.speed_augment_range[1]
-                )
+                audio, rate=np.random.uniform(_min_speed, self.speed_augment_range[1])
             )
 
         return audio
 
-
     def read_lemma_labels(self, sample_name, sample_idx):
         lemmas = self.lemma_data[sample_name][sample_idx]["caption_kw_labels"]
         labels = np.zeros((self.n_lemmas,))
@@ -274,81 +234,59 @@ def read_lemma_labels(self, sample_name, sample_idx):
 
         return labels
 
-
     def read_caption_embed(self, sample_pair, sample_idx):
         sample_embed_dir = os.path.join(
-            self.caption_embed_dir,
-            str(sample_pair[0]) + "_" + str(sample_pair[1])
+            self.caption_embed_dir, str(sample_pair[0]) + "_" + str(sample_pair[1])
         )
         embed_choices = [
-            f for f in os.listdir(sample_embed_dir) \
+            f
+            for f in os.listdir(sample_embed_dir)
             if f"chatgpt{sample_idx + 1:>02d}" in f
         ]
         # print(len(embed_choices))
 
-        return np.load(
-            os.path.join(
-                sample_embed_dir,
-                random.choice(embed_choices)
-            )
-        )
-
+        return np.load(os.path.join(sample_embed_dir, random.choice(embed_choices)))
 
     def read_clap_embed(self, sample_pair, sample_idx):
         sample_embed_dir = os.path.join(
-            self.clap_embed_dir,
-            str(sample_pair[0]) + "_" + str(sample_pair[1])
+            self.clap_embed_dir, str(sample_pair[0]) + "_" + str(sample_pair[1])
         )
         embed_choices = [
-            f for f in os.listdir(sample_embed_dir) \
+            f
+            for f in os.listdir(sample_embed_dir)
             if f"chatgpt{sample_idx + 1:>02d}" in f
         ]
         # print(len(embed_choices))
 
-        return np.load(
-            os.path.join(
-                sample_embed_dir,
-                random.choice(embed_choices)
-            )
-        )
-
+        return np.load(os.path.join(sample_embed_dir, random.choice(embed_choices)))
 
     def __len__(self):
         return len(self.idx_to_sample)
         # return 10
 
-
     def __getitem__(self, index):
         samp_a, samp_b = self.idx_to_sample[index]
         audio_a, audio_b = self.idx_to_audio[index]
         audios_to_mix = []
         try:
             for audfile in [audio_a, audio_b]:
-                audio, sr = self.load_audio(
-                    os.path.join(
-                        self.audio_dir,
-                        audfile
-                    )
-                )
+                audio, sr = self.load_audio(os.path.join(self.audio_dir, audfile))
 
                 audios_to_mix.append(audio)
         except OSError as e:
-            print('Skipping', samp_a, samp_b, 'because' ,e)
+            print("Skipping", samp_a, samp_b, "because", e)
             return None
 
-        audio,audio_path = self.mixup_audio(
-                        audios_to_mix[0],
-                        audios_to_mix[1],
-                        sample_name=str(samp_a) + "_" + str(samp_b),
-                    )
+        audio, audio_path = self.mixup_audio(
+            audios_to_mix[0],
+            audios_to_mix[1],
+            sample_name=str(samp_a) + "_" + str(samp_b),
+        )
 
-        if (self.do_pitch_augment or self.do_speed_augment) and \
-        random.random() < 0.5:
+        if (self.do_pitch_augment or self.do_speed_augment) and random.random() < 0.5:
             audio = self.augment_audio(audio)
 
-        caption_idx = random.choice(
-                        range(len(self.captions[(samp_a, samp_b)]))
-                    )
+        caption_idx = random.choice(range(len(self.captions[(samp_a, samp_b)])))
 
         caption = self.captions[(samp_a, samp_b)][caption_idx]
         # caption = self.tokenize(caption)
@@ -363,26 +301,17 @@ def __getitem__(self, index):
         }
 
         if self.lemma_data is not None:
-            encoder_labels = self.read_lemma_labels(
-                                    (samp_a, samp_b),
-                                    caption_idx
-                                )
+            encoder_labels = self.read_lemma_labels((samp_a, samp_b), caption_idx)
             bundle["encoder_labels"] = encoder_labels
         elif self.caption_embed_dir is not None:
-            encoder_embed = self.read_caption_embed(
-                                    (samp_a, samp_b),
-                                    caption_idx
-                                )
+            encoder_embed = self.read_caption_embed((samp_a, samp_b), caption_idx)
 
             if self.clap_embed_dir is not None:
                 clap_embed = self.read_clap_embed(
                     (samp_a, samp_b),
                     caption_idx,
                 )
-                encoder_embed = np.concatenate(
-                    [encoder_embed, clap_embed],
-                    axis=-1
-                )
+                encoder_embed = np.concatenate([encoder_embed, clap_embed], axis=-1)
 
             bundle["encoder_labels"] = encoder_embed
         else:
@@ -395,55 +324,65 @@ def __getitem__(self, index):
 audio_mixup_write_dir = "downloads/CLOTHO_v2.1_audio_mixup/development/"
 os.makedirs(os.path.dirname(audio_mixup_write_dir), exist_ok=True)
 
-mixup_captions_root_dir=sys.argv[1]
-clotho_audio_base_dir=sys.argv[2]
+mixup_captions_root_dir = sys.argv[1]
+clotho_audio_base_dir = sys.argv[2]
 # Create mixup dataset
 dset = ClothoMixupDataset(
     audio_dir=os.path.join(clotho_audio_base_dir, "clotho_audio_files/development"),
     tokenizer_dir="facebook/bart-base",
-    chatgpt_augmented_data_path=os.path.join(mixup_captions_root_dir,"clotho_development_chatgpt_mixups.json"),
-    chatgpt_rejected_data_path=os.path.join(mixup_captions_root_dir,"clotho_development_chatgpt_mixups_err.json"),
+    chatgpt_augmented_data_path=os.path.join(
+        mixup_captions_root_dir, "clotho_development_chatgpt_mixups.json"
+    ),
+    chatgpt_rejected_data_path=os.path.join(
+        mixup_captions_root_dir, "clotho_development_chatgpt_mixups_err.json"
+    ),
     # caption_embed_dir="sent_embedding/datasets/instructor-xl/clotho/development_chatgpt_mixup",
     # clap_embed_dir="sent_embedding/datasets/clap/clotho/development_chatgpt_mixup",
-    audio_mixup_base_dir=audio_mixup_write_dir
+    audio_mixup_base_dir=audio_mixup_write_dir,
 )
 
-data_base="data/development_clotho_chatgpt_mixup/"
-output_txt_path=os.path.join(data_base, "text")
-output_wav_scp_path=os.path.join(data_base, "wav.scp")
-output_utt2spk_path=os.path.join(data_base, "utt2spk")
+data_base = "data/development_clotho_chatgpt_mixup/"
+output_txt_path = os.path.join(data_base, "text")
+output_wav_scp_path = os.path.join(data_base, "wav.scp")
+output_utt2spk_path = os.path.join(data_base, "utt2spk")
 
 os.makedirs(os.path.dirname(output_txt_path), exist_ok=True)
 import time
+
 start = time.time()
+
+
 def _get_ith_item(i):
-    data_i=dset[i]
+    data_i = dset[i]
     if data_i is None:
         return None, None, None
-    uttid=f"clotho_chatmixup_{data_i['sample_name']}"
-    return uttid, data_i['audio_path'], data_i['labels']
+    uttid = f"clotho_chatmixup_{data_i['sample_name']}"
+    return uttid, data_i["audio_path"], data_i["labels"]
+
 
-results=None
-with Pool() as pool: # Use all cores
+results = None
+with Pool() as pool:  # Use all cores
     results = pool.map(_get_ith_item, range(len(dset)))
 
 assert results is not None
 
-N_ERROR=0
-N_PROCESSED=0
-with open(output_txt_path, 'w') as text_f, open(output_wav_scp_path,'w') as wav_scp_f, open(output_utt2spk_path,'w') as utt2spk_f:
+N_ERROR = 0
+N_PROCESSED = 0
+with open(output_txt_path, "w") as text_f, open(
+    output_wav_scp_path, "w"
+) as wav_scp_f, open(output_utt2spk_path, "w") as utt2spk_f:
     for i in range(len(dset)):
         uttid, audio_path, text = results[i]
         if uttid is None:
-            N_ERROR+=1
+            N_ERROR += 1
             continue
         print(f"{uttid} {audio_path}", file=wav_scp_f)
         print(f"{uttid} {text}", file=text_f)
         print(f"{uttid} dummy", file=utt2spk_f)
-        N_PROCESSED+=1
-        if (i+1)%1000==0:
+        N_PROCESSED += 1
+        if (i + 1) % 1000 == 0:
             print(f"Processed {i+1} mixup samples.")
 
-print(f'Number of files skipped {N_ERROR}')
-print(f'Number of files processed {N_PROCESSED}')
-#TODO(shikhar): change location of downloaded data to downloads, check gitignore for more info.
\ No newline at end of file
+print(f"Number of files skipped {N_ERROR}")
+print(f"Number of files processed {N_PROCESSED}")
+# TODO(shikhar): change location of downloaded data to downloads, check gitignore for more info.
diff --git a/egs2/clotho_v2/asr1/local/evaluation.py b/egs2/clotho_v2/asr1/local/evaluation.py
index 88d0d63c22e..257f6f7f105 100644
--- a/egs2/clotho_v2/asr1/local/evaluation.py
+++ b/egs2/clotho_v2/asr1/local/evaluation.py
@@ -74,10 +74,19 @@ def _evaluate(decode_fname: str, split: str):
                 references[key] = [value]
 
     candidates, mult_references = _convert_to_evaluation_format(predictions, references)
-    
-    evaluate = Evaluate(metrics=["spider", "fense", "meteor", "rouge_l", "spice", "spider_fl", "cider_d"])
-    corpus_scores, _ = evaluate(candidates, mult_references)
 
+    evaluate = Evaluate(
+        metrics=[
+            "spider",
+            "fense",
+            "meteor",
+            "rouge_l",
+            "spice",
+            "spider_fl",
+            "cider_d",
+        ]
+    )
+    corpus_scores, _ = evaluate(candidates, mult_references)
 
     header_str = f" Split: {split} Evaluation over {len(candidates)} predictions. "
     header_str_l = len(header_str)
@@ -89,8 +98,8 @@ def _evaluate(decode_fname: str, split: str):
     print("=" * header_str_l)
 
     # Write results in a file
-    results_file = decode_fname+'.result'
-    with open(results_file, 'w') as result_f:
+    results_file = decode_fname + ".result"
+    with open(results_file, "w") as result_f:
         header_str = f" Split: {split} Evaluation over {len(candidates)} predictions. "
         header_str_l = len(header_str)
         print("=" * header_str_l, file=result_f)
@@ -100,4 +109,5 @@ def _evaluate(decode_fname: str, split: str):
             print(f" {metric:<20}: {value} ", file=result_f)
         print("=" * header_str_l, file=result_f)
 
+
 _evaluate(args.decode_file, args.split)
diff --git a/espnet2/asr/encoder/beats_encoder.py b/espnet2/asr/encoder/beats_encoder.py
index a6c60cc6e7b..b71945072bc 100644
--- a/espnet2/asr/encoder/beats_encoder.py
+++ b/espnet2/asr/encoder/beats_encoder.py
@@ -13,8 +13,8 @@
 import logging
 import math
 import warnings
-from typing import Dict, Optional, Tuple
 from copy import deepcopy
+from typing import Dict, Optional, Tuple
 
 import numpy as np
 import torch
@@ -179,17 +179,20 @@ def __init__(
         self.use_weighted_representation = use_weighted_representation
         if self.use_weighted_representation:
             if self.max_layer is None:
-                logging.warning(f"max_layer must be provided when using weighted representations. Set to {config.encoder_layers}.")
+                logging.warning(
+                    f"max_layer must be provided when using weighted representations. Set to {config.encoder_layers}."
+                )
                 self.max_layer = config.encoder_layers
-            self.layer_weights = nn.Parameter(torch.ones(self.max_layer)/self.max_layer, requires_grad=True) # Uniform init
-        
+            self.layer_weights = nn.Parameter(
+                torch.ones(self.max_layer) / self.max_layer, requires_grad=True
+            )  # Uniform init
+
         self.transformer_adapter = None
-        if adapter_layers !=0:
+        if adapter_layers != 0:
             adapter_config = deepcopy(config)
-            adapter_config.encoder_layers=adapter_layers
+            adapter_config.encoder_layers = adapter_layers
             self.transformer_adapter = TransformerEncoder(adapter_config)
 
-
     def reload_pretrained_parameters(self):
         """Initialization function for BEATs.
 
@@ -314,7 +317,9 @@ def extract_features(
             features = self.post_extract_proj(features)
 
         features = self.dropout_input(features)
-        features, layer_results = self.encoder(features, padding_mask=padding_mask, layer=max_layer)
+        features, layer_results = self.encoder(
+            features, padding_mask=padding_mask, layer=max_layer
+        )
 
         if self.use_weighted_representation:
             for i, (x_i, _) in enumerate(layer_results):
@@ -326,9 +331,11 @@ def extract_features(
                     features += self.layer_weights[i] * x_i
             # T x B x C -> B x T x C
             features = features.transpose(0, 1)
-        
+
         if self.transformer_adapter:
-            features, _ = self.transformer_adapter(features, padding_mask = padding_mask, layer = None)
+            features, _ = self.transformer_adapter(
+                features, padding_mask=padding_mask, layer=None
+            )
 
         return features, padding_mask
 

From 64a00c4efca618777ba4564a3755e56ad5e878ba Mon Sep 17 00:00:00 2001
From: Shikhar Bharadwaj <sbharadwaj@dt-login03.delta.ncsa.illinois.edu>
Date: Tue, 5 Nov 2024 09:17:04 -0600
Subject: [PATCH 10/47] add local paths to audiocaps and clotho_chatgpt_mixup

---
 egs2/TEMPLATE/asr1/db.sh | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/egs2/TEMPLATE/asr1/db.sh b/egs2/TEMPLATE/asr1/db.sh
index 157471e78a8..a89a8900048 100755
--- a/egs2/TEMPLATE/asr1/db.sh
+++ b/egs2/TEMPLATE/asr1/db.sh
@@ -221,6 +221,11 @@ KSPONSPEECH=
 HIFITTS=downloads
 CLOTHO_V2=downloads
 
+if [[ "$(hostname -d)" == delta.ncsa.illinois.edu ]]; then
+    AUDIOCAPS=/old_scratch/bbjs/shared/corpora/audiocaps/
+    CLOTHO_CHATGPT_MIXUP=/old_scratch/bbjs/slseanwu/dcase23_aac/beats_baseline/chatgpt/
+fi
+
 # For only CMU TIR environment
 if [[ "$(hostname)" == tir* ]]; then
     BABEL_101=/projects/tir5/data/speech_corpora/babel/IARPA_BABEL_BP_101/

From ac484495e6e3cbf58bd3d1175c607c3f05bf6898 Mon Sep 17 00:00:00 2001
From: Shikhar Bharadwaj <sbharadwaj@dt-login03.delta.ncsa.illinois.edu>
Date: Tue, 5 Nov 2024 09:24:48 -0600
Subject: [PATCH 11/47] remove ipynb ckpts

---
 .../Untitled-checkpoint.ipynb                 |  6 --
 .../Checker-checkpoint.ipynb                  | 91 -------------------
 .../.ipynb_checkpoints/data-checkpoint.sh     | 61 -------------
 .../data_prep-checkpoint.py                   | 49 ----------
 4 files changed, 207 deletions(-)
 delete mode 100644 egs2/clotho_v2/.ipynb_checkpoints/Untitled-checkpoint.ipynb
 delete mode 100644 egs2/clotho_v2/asr1/local/.ipynb_checkpoints/Checker-checkpoint.ipynb
 delete mode 100755 egs2/clotho_v2/asr1/local/.ipynb_checkpoints/data-checkpoint.sh
 delete mode 100644 egs2/clotho_v2/asr1/local/.ipynb_checkpoints/data_prep-checkpoint.py

diff --git a/egs2/clotho_v2/.ipynb_checkpoints/Untitled-checkpoint.ipynb b/egs2/clotho_v2/.ipynb_checkpoints/Untitled-checkpoint.ipynb
deleted file mode 100644
index 363fcab7ed6..00000000000
--- a/egs2/clotho_v2/.ipynb_checkpoints/Untitled-checkpoint.ipynb
+++ /dev/null
@@ -1,6 +0,0 @@
-{
- "cells": [],
- "metadata": {},
- "nbformat": 4,
- "nbformat_minor": 5
-}
diff --git a/egs2/clotho_v2/asr1/local/.ipynb_checkpoints/Checker-checkpoint.ipynb b/egs2/clotho_v2/asr1/local/.ipynb_checkpoints/Checker-checkpoint.ipynb
deleted file mode 100644
index 231fa5bf132..00000000000
--- a/egs2/clotho_v2/asr1/local/.ipynb_checkpoints/Checker-checkpoint.ipynb
+++ /dev/null
@@ -1,91 +0,0 @@
-{
- "cells": [
-  {
-   "cell_type": "code",
-   "execution_count": 14,
-   "id": "1b6f0ae7-c574-4da9-8ed2-01513728be7e",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import pandas as pd\n",
-    "import os\n",
-    "from scipy.io import wavfile\n",
-    "from collections import defaultdict as dd"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 18,
-   "id": "fe8ed4fc-4561-4e71-aa1d-8f41e61c7633",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "base_path = \"/data/user_data/sbharad2/espnet/egs2/clotho_v2/asr1/local/CLOTHO_v2.1/\"\n",
-    "splits = ['development', 'evaluation','test','validation']\n",
-    "\n",
-    "def get_audio_duration(fname):\n",
-    "    samplerate, data = wavfile.read(fname)\n",
-    "    duration_s = len(data)/samplerate\n",
-    "    return duration_s\n",
-    "\n",
-    "# Audio length\n",
-    "audio_s = dd(lambda : [])\n",
-    "for split in splits:\n",
-    "    for fname in os.listdir(os.path.join(base_path,'clotho_audio_files',split)):\n",
-    "        ffname = os.path.join(base_path,'clotho_audio_files',split,fname)\n",
-    "        audio_s[split].append(get_audio_duration(ffname))"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 23,
-   "id": "59ea5e97-0f9d-4931-b28c-699513e38318",
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "development 30.0 15.0\n",
-      "evaluation 30.0 15.003174603174603\n",
-      "test 30.0 15.000090702947846\n",
-      "validation 29.95390022675737 15.05\n"
-     ]
-    }
-   ],
-   "source": [
-    "for split in splits:\n",
-    "    print(split,max(audio_s[split]),min(audio_s[split]))"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "b3884f12-1a4a-41ef-b4c9-e7c99af27086",
-   "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.8.19"
-  }
- },
- "nbformat": 4,
- "nbformat_minor": 5
-}
diff --git a/egs2/clotho_v2/asr1/local/.ipynb_checkpoints/data-checkpoint.sh b/egs2/clotho_v2/asr1/local/.ipynb_checkpoints/data-checkpoint.sh
deleted file mode 100755
index e984e589eaf..00000000000
--- a/egs2/clotho_v2/asr1/local/.ipynb_checkpoints/data-checkpoint.sh
+++ /dev/null
@@ -1,61 +0,0 @@
-#!/usr/bin/env bash
-# Set bash to 'debug' mode, it will exit on :
-# -e 'error', -u 'undefined variable', -o ... 'error in pipeline', -x 'print commands',
-set -e
-set -u
-set -o pipefail
-
-log() {
-    local fname=${BASH_SOURCE[1]##*/}
-    echo -e "$(date '+%Y-%m-%dT%H:%M:%S') (${fname}:${BASH_LINENO[0]}:${FUNCNAME[1]}) $*"
-}
-SECONDS=0
-
-stage=1
-stop_stage=100
-
-log "$0 $*"
-. utils/parse_options.sh
-
-. ./db.sh
-. ./path.sh
-. ./cmd.sh
-
-if [ $# -ne 0 ]; then
-    log "Error: No positional arguments are required."
-    exit 2
-fi
-
-if [ -z "${CLOTHO_V2}" ]; then
-    log "Fill the value of 'CLOTHO_V2' of db.sh"
-    exit 1
-fi
-
-SPLITS=(development validation evaluation)
-
-if [ ${stage} -le 1 ] && [ ${stop_stage} -ge 1 ]; then
-    log "stage 1: Data preparation"
-    for split_name in ${SPLITS[@]}; do
-        mkdir -p "data/${split_name}"
-    done
-
-    if [ ! -d ${CLOTHO_V2} ]; then
-        echo Cannot find CLOTHO_V2 root! Exiting...
-        exit 1
-    fi
-
-    # Prepare data in the Kaldi format, including three files:
-    # text, wav.scp, utt2spk
-    echo "$(which python)"
-    python3 local/data_prep.py ${CLOTHO_V2}
-
-    for split_name in ${SPLITS[@]}; do
-        for f in text wav.scp utt2spk; do
-            sort data/${split_name}/${f} -o data/${split_name}/${f}
-        done
-        echo "Running spk2utt"
-        utils/utt2spk_to_spk2utt.pl data/${split_name}/utt2spk > "data/${split_name}/spk2utt"
-    done
-fi
-
-log "Successfully finished. [elapsed=${SECONDS}s]"
diff --git a/egs2/clotho_v2/asr1/local/.ipynb_checkpoints/data_prep-checkpoint.py b/egs2/clotho_v2/asr1/local/.ipynb_checkpoints/data_prep-checkpoint.py
deleted file mode 100644
index 52c7ba98290..00000000000
--- a/egs2/clotho_v2/asr1/local/.ipynb_checkpoints/data_prep-checkpoint.py
+++ /dev/null
@@ -1,49 +0,0 @@
-"""Prepares data for Clotho_v2 dataset."""
-
-import glob
-import os
-import sys
-
-import pandas as pd
-
-if __name__ == "__main__":
-    ROOT_DATA_DIR = sys.argv[1]
-    AUDIO_DIR = "clotho_audio_files"
-    CSV_DIR = "clotho_csv_files"
-
-    SPLIT_NAMES = ["development", "validation", "evaluation"]
-
-    for data_split in SPLIT_NAMES:
-        missing_audio = []
-        audio_regex = os.path.join(ROOT_DATA_DIR, AUDIO_DIR, data_split, "*.wav")
-        all_audio_list = glob.glob(audio_regex)
-        all_captions_path = os.path.join(
-            ROOT_DATA_DIR, CSV_DIR, f"clotho_captions_{data_split}.csv"
-        )
-        captions_df = pd.read_csv(all_captions_path)
-        N_PROCESSED = 0
-        with open(
-            os.path.join("data", data_split, "text"), "w", encoding="utf-8"
-        ) as text_f, open(
-            os.path.join("data", data_split, "wav.scp"), "w", encoding="utf-8"
-        ) as wav_scp_f, open(
-            os.path.join("data", data_split, "utt2spk"), "w", encoding="utf-8"
-        ) as utt2spk_f:
-            for uttid, row in captions_df.iterrows():
-                uttid = f"{data_split}_uttid"
-                text = row["caption_1"].strip()
-                file_name = row["file_name"]
-                audio_path = os.path.join(
-                    ROOT_DATA_DIR, AUDIO_DIR, data_split, f"{file_name}"
-                )
-                if audio_path not in all_audio_list:
-                    # Missing audio file
-                    missing_audio += [audio_path]
-                    continue
-                print(f"{uttid} dummy", file=utt2spk_f)
-                print(f"{uttid} {audio_path}", file=wav_scp_f)
-                print(f"{uttid} {text}", file=text_f)
-                N_PROCESSED += 1
-        print(
-            f"For split {data_split}: Processed {N_PROCESSED} audio and their captions. {len(missing_audio)} audio files were missing."
-        )

From 2303395cfdc77bccd0161f9e01994509c00344d8 Mon Sep 17 00:00:00 2001
From: Shikhar Bharadwaj <sbharad2@andrew.cmu.edu>
Date: Tue, 12 Nov 2024 15:31:37 -0500
Subject: [PATCH 12/47] changes for running on babel environment

---
 egs2/TEMPLATE/asr1/db.sh                      |  7 +++
 egs2/clotho_v2/asr1/cmd.sh                    |  2 +-
 egs2/clotho_v2/asr1/conf/beats_bart_ft.yaml   |  2 +-
 egs2/clotho_v2/asr1/conf/beats_bart_pt.yaml   |  2 +-
 egs2/clotho_v2/asr1/conf/slurm.conf           |  4 +-
 egs2/clotho_v2/asr1/local/data.sh             | 19 ++++---
 .../local/data_prep_clotho_chatgpt_mixup.py   |  7 +--
 egs2/clotho_v2/asr1/run_ft.sh                 | 27 +++++++---
 egs2/clotho_v2/asr1/run_inference.sh          | 45 ++++++++++++++++
 egs2/clotho_v2/asr1/run_pt.sh                 | 51 ++++---------------
 10 files changed, 104 insertions(+), 62 deletions(-)
 create mode 100755 egs2/clotho_v2/asr1/run_inference.sh

diff --git a/egs2/TEMPLATE/asr1/db.sh b/egs2/TEMPLATE/asr1/db.sh
index fdc658e7eb2..1073f873c12 100755
--- a/egs2/TEMPLATE/asr1/db.sh
+++ b/egs2/TEMPLATE/asr1/db.sh
@@ -222,11 +222,18 @@ KSPONSPEECH=
 HIFITTS=downloads
 CLOTHO_V2=downloads
 
+# For only Delta environment
 if [[ "$(hostname -d)" == delta.ncsa.illinois.edu ]]; then
     AUDIOCAPS=/old_scratch/bbjs/shared/corpora/audiocaps/
     CLOTHO_CHATGPT_MIXUP=/old_scratch/bbjs/slseanwu/dcase23_aac/beats_baseline/chatgpt/
 fi
 
+# For only CMU Babel environment on compute nodes
+if [[ "$(hostname)" == babel* ]]; then
+    AUDIOCAPS=/compute/babel-13-33/sbharad2/audiocaps/
+    CLOTHO_CHATGPT_MIXUP=/compute/babel-13-33/sbharad2/chatgpt/
+fi
+
 # For only CMU TIR environment
 if [[ "$(hostname)" == tir* ]]; then
     BABEL_101=/projects/tir5/data/speech_corpora/babel/IARPA_BABEL_BP_101/
diff --git a/egs2/clotho_v2/asr1/cmd.sh b/egs2/clotho_v2/asr1/cmd.sh
index 2aae6919fef..257c6c739df 100644
--- a/egs2/clotho_v2/asr1/cmd.sh
+++ b/egs2/clotho_v2/asr1/cmd.sh
@@ -28,7 +28,7 @@
 
 
 # Select the backend used by run.sh from "local", "stdout", "sge", "slurm", or "ssh"
-cmd_backend='local'
+cmd_backend='slurm'
 
 # Local machine, without any Job scheduling system
 if [ "${cmd_backend}" = local ]; then
diff --git a/egs2/clotho_v2/asr1/conf/beats_bart_ft.yaml b/egs2/clotho_v2/asr1/conf/beats_bart_ft.yaml
index d95bb1d0939..7e01b01f665 100644
--- a/egs2/clotho_v2/asr1/conf/beats_bart_ft.yaml
+++ b/egs2/clotho_v2/asr1/conf/beats_bart_ft.yaml
@@ -29,7 +29,7 @@ freeze_param: [
 
 encoder: beats
 encoder_conf:
-    beats_ckpt_path: /u/sbharadwaj/BEATs_models/BEATs_iter3_plus_AS2M.pt # _finetuned_on_AS2M_cpt1
+    beats_ckpt_path: /compute/babel-13-33/sbharad2/models/BEATs/BEATs_iter3_plus_AS2M_finetuned_on_AS2M_cpt1.pt # /u/sbharadwaj/BEATs_models/BEATs_iter3_plus_AS2M.pt # _finetuned_on_AS2M_cpt1
     specaug_config:
         apply_freq_mask: true
         freq_mask_width_range:
diff --git a/egs2/clotho_v2/asr1/conf/beats_bart_pt.yaml b/egs2/clotho_v2/asr1/conf/beats_bart_pt.yaml
index 86862c04de3..4c146c1cfd5 100644
--- a/egs2/clotho_v2/asr1/conf/beats_bart_pt.yaml
+++ b/egs2/clotho_v2/asr1/conf/beats_bart_pt.yaml
@@ -31,7 +31,7 @@ freeze_param: [
 
 encoder: beats
 encoder_conf:
-    beats_ckpt_path: /u/sbharadwaj/BEATs_models/BEATs_iter3_plus_AS2M_finetuned_on_AS2M_cpt1.pt # _finetuned_on_AS2M_cpt1
+    beats_ckpt_path: /compute/babel-13-33/sbharad2/models/BEATs/BEATs_iter3_plus_AS2M_finetuned_on_AS2M_cpt1.pt # /u/sbharadwaj/BEATs_models/BEATs_iter3_plus_AS2M_finetuned_on_AS2M_cpt1.pt # _finetuned_on_AS2M_cpt1
     specaug_config:
         apply_freq_mask: true
         freq_mask_width_range:
diff --git a/egs2/clotho_v2/asr1/conf/slurm.conf b/egs2/clotho_v2/asr1/conf/slurm.conf
index 9f397a1b55d..9ebbf6b3d90 100644
--- a/egs2/clotho_v2/asr1/conf/slurm.conf
+++ b/egs2/clotho_v2/asr1/conf/slurm.conf
@@ -12,5 +12,5 @@ option num_threads=* --cpus-per-task $0
 option num_threads=1 --cpus-per-task 1
 option num_nodes=* --nodes $0
 default gpu=0
-option gpu=0 -p cpu -A bbjs-delta-cpu
-option gpu=* -p gpuA40x4 -A bbjs-delta-gpu --gres=gpu:$0 -c $0
+option gpu=0 -p general
+option gpu=* -p general --gres=gpu:$0 -c $0
diff --git a/egs2/clotho_v2/asr1/local/data.sh b/egs2/clotho_v2/asr1/local/data.sh
index aa005c9059e..701054b560c 100755
--- a/egs2/clotho_v2/asr1/local/data.sh
+++ b/egs2/clotho_v2/asr1/local/data.sh
@@ -18,22 +18,27 @@ log "$0 $*"
 . ./path.sh
 . ./cmd.sh
 
-if [ $# -ne 0 ]; then
-    log "Error: No positional arguments are required."
-    exit 2
-fi
-
 if [ -z "${CLOTHO_V2}" ]; then
     log "Fill the value of 'CLOTHO_V2' of db.sh"
     exit 1
 fi
 
 log "stage 1: Data preparation"
+
 ## DOWNLOAD DATA
-CLOTHO_V2_ROOT_DIR="$(pwd)/downloads"
+# If there is no argument, the default download directory is set to currentdir/downloads
+if [ $# -ne 1 ]; then
+    CLOTHO_V2_ROOT_DIR="$(pwd)/downloads"
+    log "Using the default download directory: ${CLOTHO_V2_ROOT_DIR}"
+else
+    CLOTHO_V2_ROOT_DIR="$1/downloads"
+    log "Using the specified download directory: ${CLOTHO_V2_ROOT_DIR}"
+fi
+
+
 if [ ! -e "${CLOTHO_V2_ROOT_DIR}/download_done" ]; then
     log "stage 1: Data preparation - Installing aac-datasets"
-    if ! pip install aac-datasets; then
+    if ! pip3 install aac-datasets; then
         echo "Error: Installing aac-datasets failed."
         exit 1
     fi
diff --git a/egs2/clotho_v2/asr1/local/data_prep_clotho_chatgpt_mixup.py b/egs2/clotho_v2/asr1/local/data_prep_clotho_chatgpt_mixup.py
index 1be84bfda95..c37587d8f89 100644
--- a/egs2/clotho_v2/asr1/local/data_prep_clotho_chatgpt_mixup.py
+++ b/egs2/clotho_v2/asr1/local/data_prep_clotho_chatgpt_mixup.py
@@ -320,12 +320,13 @@ def __getitem__(self, index):
         return bundle
 
 
+mixup_captions_root_dir = sys.argv[1]
+clotho_audio_base_dir = sys.argv[2]
+
 # create directory paths for storing clotho mixup audio
-audio_mixup_write_dir = "downloads/CLOTHO_v2.1_audio_mixup/development/"
+audio_mixup_write_dir = f"{clotho_audio_base_dir}_audio_mixup/development/"  # "downloads/CLOTHO_v2.1_audio_mixup/development/"
 os.makedirs(os.path.dirname(audio_mixup_write_dir), exist_ok=True)
 
-mixup_captions_root_dir = sys.argv[1]
-clotho_audio_base_dir = sys.argv[2]
 # Create mixup dataset
 dset = ClothoMixupDataset(
     audio_dir=os.path.join(clotho_audio_base_dir, "clotho_audio_files/development"),
diff --git a/egs2/clotho_v2/asr1/run_ft.sh b/egs2/clotho_v2/asr1/run_ft.sh
index a012172175e..b4aaf9120fb 100755
--- a/egs2/clotho_v2/asr1/run_ft.sh
+++ b/egs2/clotho_v2/asr1/run_ft.sh
@@ -6,17 +6,31 @@ set -euo pipefail
 wandb_init_args=""
 other_args="$@"
 
-pre_trained_model_path=exp/asr_pt.initfix.bigbatch512.lr2e-4.weighted12layers.20241103.145125/valid.acc.ave_5best.pth
+expdir=/compute/babel-13-33/sbharad2/expdir
+dumpdir=/compute/babel-13-33/sbharad2/dumpdir
+local_data_opts=/compute/babel-13-33/sbharad2/expdir
+
+# pre_trained_model_path=exp/asr_pt.initfix.bigbatch512.lr2e-4.weighted12layers.20241103.145125/valid.acc.ave_5best.pth
+
+pt_tag=full.20241111.170406
+ckpt_name=13epoch
+
+pre_trained_model_path=${expdir}/asr_pt.${pt_tag}/${ckpt_name}.pth
+ft_tag=${pt_tag}.${ckpt_name}
+
+
 
 ./asr.sh \
-    --asr_tag ft_lr5e-5.initfix.bigbatch512.lr2e-4.weighted12layers.20241103.145125 \
+    --asr_tag ft.${ft_tag} \
+    --expdir ${expdir} \
+    --dumpdir ${dumpdir} \
     --feats_normalize uttmvn \
-    --stage 11 \
+    --stage 2 \
     --stop_stage 13 \
-    --asr_stats_dir exp/asr_stats_finetune \
+    --asr_stats_dir ${expdir}/asr_stats_finetune \
     --ngpu 2 \
     --gpu_inference true \
-    --nj 20 \
+    --nj 8 \
     --inference_nj 1 \
     --max_wav_duration 30 \
     --token_type hugging_face \
@@ -30,4 +44,5 @@ pre_trained_model_path=exp/asr_pt.initfix.bigbatch512.lr2e-4.weighted12layers.20
     --pretrained_model ${pre_trained_model_path} \
     --inference_asr_model valid.acc.ave_5best.pth \
     --asr_args "${wandb_init_args} ${other_args}" \
-    --local_score_opts exp/asr_ft.initfix.bigbatch512.lr2e-4.weighted12layers.20241103.145125/inference_beam_size10_ctc_weight0.0_hugging_face_decoderTrue_asr_model_valid.acc.ave_5best
+    --local_data_opts "${local_data_opts}" \
+    --local_score_opts ${expdir}/asr_ft.${ft_tag}/inference_beam_size10_ctc_weight0.0_hugging_face_decoderTrue_asr_model_${ckpt_name}
diff --git a/egs2/clotho_v2/asr1/run_inference.sh b/egs2/clotho_v2/asr1/run_inference.sh
new file mode 100755
index 00000000000..f9450152391
--- /dev/null
+++ b/egs2/clotho_v2/asr1/run_inference.sh
@@ -0,0 +1,45 @@
+#!/usr/bin/env bash
+# Runs pre-training
+set -euo pipefail
+
+# # Run inference on all checkpoints
+
+wandb_sweep_args="$@"
+
+# name_tag=ft.initfix.bigbatch512.lr2e-4.weighted12layers.20241103.145125
+# name_tag=pt.full.20241111.170406
+
+name_tag=pt.data_check.20241111.153210
+expdir=/compute/babel-13-33/sbharad2/expdir
+dumpdir=/compute/babel-13-33/sbharad2/dumpdir
+
+all_ckpts=$(find ${expdir}/asr_${name_tag} -name "*.pth"  -printf "%f\n" | awk '{print substr($0, 1, length($0) - 4)}')
+
+# for ckpt in valid.acc.best; do
+for ckpt in $all_ckpts ; do
+    if [[ $ckpt == "checkpoint" || $ckpt == "latest"  || $ckpt == "valid.acc.best" ]]; then
+        continue
+    fi
+    ./asr.sh \
+        --asr_tag $name_tag \
+        --expdir ${expdir} \
+        --dumpdir ${dumpdir} \
+        --feats_normalize uttmvn \
+        --stage 12 \
+        --stop_stage 13 \
+        --ngpu 1 \
+        --gpu_inference true \
+        --nj 20 \
+        --inference_nj 1 \
+        --max_wav_duration 30 \
+        --token_type hugging_face \
+        --use_lm false \
+        --hugging_face_model_name_or_path "facebook/bart-base" \
+        --inference_args "--beam_size 10 --ctc_weight 0.0 --hugging_face_decoder True" \
+        --train_set pretrain \
+        --valid_set validation \
+        --test_sets "evaluation" \
+        --inference_asr_model ${ckpt}.pth \
+        --asr_args "${wandb_sweep_args}" \
+        --local_score_opts "${expdir}/asr_${name_tag}/inference_beam_size10_ctc_weight0.0_hugging_face_decoderTrue_asr_model_${ckpt}"
+done
\ No newline at end of file
diff --git a/egs2/clotho_v2/asr1/run_pt.sh b/egs2/clotho_v2/asr1/run_pt.sh
index f64472baefa..d14dae97c46 100755
--- a/egs2/clotho_v2/asr1/run_pt.sh
+++ b/egs2/clotho_v2/asr1/run_pt.sh
@@ -8,15 +8,21 @@ timestamp=$(date "+%Y%m%d.%H%M%S")
 wandb_init_args=""
 other_args="$@"
 
+expdir=/compute/babel-13-33/sbharad2/expdir
+dumpdir=/compute/babel-13-33/sbharad2/dumpdir
+local_data_opts=/compute/babel-13-33/sbharad2/expdir
+
 ./asr.sh \
-    --asr_tag pt.initfix.bigbatch512.lr2e-4.weighted12layers.${timestamp} \
+    --asr_tag pt.full.${timestamp} \
     --asr_speech_fold_length 1600 \
+    --expdir ${expdir} \
+    --dumpdir ${dumpdir} \
     --feats_normalize uttmvn \
     --stage 11 \
     --stop_stage 13 \
     --ngpu 2 \
     --gpu_inference true \
-    --nj 20 \
+    --nj 10 \
     --inference_nj 1 \
     --max_wav_duration 30 \
     --token_type hugging_face \
@@ -26,45 +32,8 @@ other_args="$@"
     --train_set pretrain \
     --valid_set validation \
     --test_sets "validation evaluation" \
-    --asr_config conf/beats_bart_pt.yaml \
+    --asr_config conf/beats_bart_full.yaml \
     --inference_asr_model valid.acc.best.pth \
     --asr_args "${wandb_init_args} ${other_args}" \
+    --local_data_opts "${local_data_opts}" \
     --local_score_opts "exp/asr_pt.${timestamp}/inference_beam_size10_ctc_weight0.0_hugging_face_decoderTrue_asr_model_valid.acc.best"
-
-
-
-# # Run inference on all checkpoints
-
-# wandb_sweep_args="$@"
-
-# name_tag=ft.initfix.bigbatch512.lr2e-4.weighted12layers.20241103.145125
-
-# all_ckpts=$(find exp/asr_${name_tag} -name "*.pth"  -printf "%f\n" | awk '{print substr($0, 1, length($0) - 4)}')
-
-# # for ckpt in valid.acc.best; do
-# for ckpt in $all_ckpts ; do
-#     if [[ $ckpt == "checkpoint" || $ckpt == "latest"  || $ckpt == "valid.acc.best" ]]; then
-#         continue
-#     fi
-#     ./asr.sh \
-#         --asr_tag $name_tag \
-#         --feats_normalize uttmvn \
-#         --stage 12 \
-#         --stop_stage 13 \
-#         --ngpu 1 \
-#         --gpu_inference true \
-#         --nj 20 \
-#         --inference_nj 1 \
-#         --max_wav_duration 30 \
-#         --token_type hugging_face \
-#         --use_lm false \
-#         --hugging_face_model_name_or_path "facebook/bart-base" \
-#         --inference_args "--beam_size 10 --ctc_weight 0.0 --hugging_face_decoder True" \
-#         --train_set pretrain \
-#         --valid_set validation \
-#         --test_sets "evaluation" \
-#         --asr_config conf/dcase.yaml \
-#         --inference_asr_model ${ckpt}.pth \
-#         --asr_args "${wandb_sweep_args}" \
-#         --local_score_opts "exp/asr_${name_tag}/inference_beam_size10_ctc_weight0.0_hugging_face_decoderTrue_asr_model_${ckpt}"
-# done

From 32d539a33d07554d234f40d9b90e797d1d9174f5 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Tue, 12 Nov 2024 20:32:57 +0000
Subject: [PATCH 13/47] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 egs2/clotho_v2/asr1/run_inference.sh | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/egs2/clotho_v2/asr1/run_inference.sh b/egs2/clotho_v2/asr1/run_inference.sh
index f9450152391..38de696eeb1 100755
--- a/egs2/clotho_v2/asr1/run_inference.sh
+++ b/egs2/clotho_v2/asr1/run_inference.sh
@@ -42,4 +42,4 @@ for ckpt in $all_ckpts ; do
         --inference_asr_model ${ckpt}.pth \
         --asr_args "${wandb_sweep_args}" \
         --local_score_opts "${expdir}/asr_${name_tag}/inference_beam_size10_ctc_weight0.0_hugging_face_decoderTrue_asr_model_${ckpt}"
-done
\ No newline at end of file
+done

From af85cd59cea4029541b0bd4eabd31ef28bcf5da0 Mon Sep 17 00:00:00 2001
From: Shikhar Bharadwaj <sbharad2@andrew.cmu.edu>
Date: Wed, 13 Nov 2024 22:50:24 -0500
Subject: [PATCH 14/47] clean up

---
 egs2/clotho_v2/asr1/README.md        | 68 ++++++++++++++++++++++++++++
 egs2/clotho_v2/asr1/local/score.sh   |  2 +-
 egs2/clotho_v2/asr1/run_ft.sh        | 22 +++++----
 egs2/clotho_v2/asr1/run_inference.sh | 20 ++++----
 egs2/clotho_v2/asr1/run_pt.sh        | 19 ++++----
 egs2/clotho_v2/asr1/run_upload.sh    | 47 +++++++++++++++++++
 6 files changed, 148 insertions(+), 30 deletions(-)
 create mode 100644 egs2/clotho_v2/asr1/README.md
 create mode 100755 egs2/clotho_v2/asr1/run_upload.sh

diff --git a/egs2/clotho_v2/asr1/README.md b/egs2/clotho_v2/asr1/README.md
new file mode 100644
index 00000000000..57054793a5c
--- /dev/null
+++ b/egs2/clotho_v2/asr1/README.md
@@ -0,0 +1,68 @@
+# Clotho Audio Captioning RECIPE
+
+This recipe implements the DCASE 2023 Automated Audio Captioning (AAC) task with a BEATs encoder BART decoder model on the Clotho_v2 dataset, very close to what is described in [this paper](https://dcase.community/documents/challenge2023/technical_reports/DCASE2023_Wu_31_t6a.pdf) and reusing part of code from the [original implementation](https://github.com/slSeanWU/beats-conformer-bart-audio-captioner?tab=readme-ov-file). 
+We also reuse the code from the [BEATs repository](https://github.com/microsoft/unilm/tree/master/beats) for this implementation.
+
+# Training Details and Requirements
+The training is divided into an AAC pre-training and a fine-tuning stage.
+We do the AAC pre-training on [AudioCaps](https://aclanthology.org/N19-1011/) and [Clotho ChatGPT mixup](https://huggingface.co/datasets/slseanwu/clotho-chatgpt-mixup-50K) data.
+Fine-tuning is performed only on the standard development set of Clotho_v2 dataset.
+
+
+### Steps to run
+For running AAC pre-training
+1. Download AudioCaps and set the path to its root directory in db.sh (This recipe downloads clotho for you). 
+2. Download the BEATs checkpoint: [BEATs_iter3+](https://onedrive.live.com/?authkey=%21AGXnEG4l3mlIzfA&id=6B83B49411CA81A7%2125960&cid=6B83B49411CA81A7&parId=root&parQt=sharedby&o=OneUp) and change the `beats_ckpt_path` path in `conf/beats_bart_pt.yaml`
+3. Launch with `run_pt.sh`
+
+For running fine-tuning: 
+1. Set the value of `pt_tag` in `run_ft.sh` and
+2. Launch with `run_ft.sh`
+
+For running inference and scoring on all checkpoints `run_inference.sh` can be used by just changing the `name_tag`.
+
+## Trained checkpoints
+AAC Pre-trained model: https://huggingface.co/shikhar7ssu/dcase23.aac.pt
+
+Fine-tuned model: https://huggingface.co/shikhar7ssu/dcase23.aac
+
+
+### GPU Time
+AAC pre-training takes around ~12 hours on a single A40
+Fine-tuning, decoding and evalution takes ~2.5 hours on a single A40.
+All the scripts above are setup with 2 GPUs but that can be changed with the `ngpu` argument.
+
+
+<!-- Generated by scripts/utils/show_asr_result.sh -->
+# RESULTS
+## Environments
+- date: `Tue Nov  5 15:18:41 CST 2024`
+- python version: `3.9.18 | packaged by conda-forge | (main, Dec 23 2023, 16:33:10)  [GCC 12.3.0]`
+- espnet version: `espnet 202409`
+- pytorch version: `pytorch 2.4.0`
+- Git hash: `ac484495e6e3cbf58bd3d1175c607c3f05bf6898`
+  - Commit date: `Tue Nov 5 09:24:48 2024 -0600`
+
+<!-- Copied from the output produced by local/evaluation.py -->
+## exp/asr_ft_lr5e-5.initfix.bigbatch512.lr2e-4.weighted12layers.20241103.145125
+```
+=====================================================
+ Split: evaluation Evaluation over 1045 predictions.
+=====================================================
+ cider_d             : 0.40163786973038024
+ spice               : 0.12328679455316437
+ spider              : 0.2624623321417723
+ sbert_sim           : 0.5094392538748004
+ fer                 : 0.03827751196172249
+ fense               : 0.4915462823319092
+ meteor              : 0.17194480402701573
+ rouge_l             : 0.3550182495603068
+ fer.add_tail_prob   : 0.04409123957157135
+ fer.repeat_event_prob: 0.07279406487941742
+ fer.repeat_adv_prob : 0.0016611652681604028
+ fer.remove_conj_prob: 0.10949967056512833
+ fer.remove_verb_prob: 0.20523187518119812
+ fer.error_prob      : 0.319917231798172
+ spider_fl           : 0.2532289450419614
+=====================================================
+```
\ No newline at end of file
diff --git a/egs2/clotho_v2/asr1/local/score.sh b/egs2/clotho_v2/asr1/local/score.sh
index 147604a16ff..cea04f020c5 100755
--- a/egs2/clotho_v2/asr1/local/score.sh
+++ b/egs2/clotho_v2/asr1/local/score.sh
@@ -12,7 +12,7 @@ fi
 
 pip install aac-metrics
 aac-metrics-download
-
+# Currently this script just creates a text.result file in the same directory as the decode file.
 asr_expdir=$1
 splits=(evaluation)
 
diff --git a/egs2/clotho_v2/asr1/run_ft.sh b/egs2/clotho_v2/asr1/run_ft.sh
index b4aaf9120fb..2b2083779ae 100755
--- a/egs2/clotho_v2/asr1/run_ft.sh
+++ b/egs2/clotho_v2/asr1/run_ft.sh
@@ -6,26 +6,29 @@ set -euo pipefail
 wandb_init_args=""
 other_args="$@"
 
-expdir=/compute/babel-13-33/sbharad2/expdir
-dumpdir=/compute/babel-13-33/sbharad2/dumpdir
-local_data_opts=/compute/babel-13-33/sbharad2/expdir
+# expdir=/compute/babel-13-33/sbharad2/expdir
+# dumpdir=/compute/babel-13-33/sbharad2/dumpdir
+# local_data_opts=/compute/babel-13-33/sbharad2/expdir
+#     --expdir ${expdir} \
+#     --dumpdir ${dumpdir} \
+#     --local_data_opts "${local_data_opts}" \
 
 # pre_trained_model_path=exp/asr_pt.initfix.bigbatch512.lr2e-4.weighted12layers.20241103.145125/valid.acc.ave_5best.pth
+# pt_tag=full.20241111.170406
+# pt_tag=data_check.20241111.153210
 
-pt_tag=full.20241111.170406
-ckpt_name=13epoch
+pt_tag=beats.bartfrozen.20241111.164042
+
+ckpt_name=valid.acc.ave
 
 pre_trained_model_path=${expdir}/asr_pt.${pt_tag}/${ckpt_name}.pth
 ft_tag=${pt_tag}.${ckpt_name}
 
 
-
 ./asr.sh \
     --asr_tag ft.${ft_tag} \
-    --expdir ${expdir} \
-    --dumpdir ${dumpdir} \
     --feats_normalize uttmvn \
-    --stage 2 \
+    --stage 1 \
     --stop_stage 13 \
     --asr_stats_dir ${expdir}/asr_stats_finetune \
     --ngpu 2 \
@@ -44,5 +47,4 @@ ft_tag=${pt_tag}.${ckpt_name}
     --pretrained_model ${pre_trained_model_path} \
     --inference_asr_model valid.acc.ave_5best.pth \
     --asr_args "${wandb_init_args} ${other_args}" \
-    --local_data_opts "${local_data_opts}" \
     --local_score_opts ${expdir}/asr_ft.${ft_tag}/inference_beam_size10_ctc_weight0.0_hugging_face_decoderTrue_asr_model_${ckpt_name}
diff --git a/egs2/clotho_v2/asr1/run_inference.sh b/egs2/clotho_v2/asr1/run_inference.sh
index f9450152391..329f78bba5b 100755
--- a/egs2/clotho_v2/asr1/run_inference.sh
+++ b/egs2/clotho_v2/asr1/run_inference.sh
@@ -6,28 +6,28 @@ set -euo pipefail
 
 wandb_sweep_args="$@"
 
-# name_tag=ft.initfix.bigbatch512.lr2e-4.weighted12layers.20241103.145125
-# name_tag=pt.full.20241111.170406
 
-name_tag=pt.data_check.20241111.153210
-expdir=/compute/babel-13-33/sbharad2/expdir
-dumpdir=/compute/babel-13-33/sbharad2/dumpdir
+# expdir=/compute/babel-13-33/sbharad2/expdir
+# dumpdir=/compute/babel-13-33/sbharad2/dumpdir
+#         --expdir ${expdir} \
+#         --dumpdir ${dumpdir} \
 
+
+name_tag=ft.beats.bartfrozen.20241111.164042.valid.acc.ave
 all_ckpts=$(find ${expdir}/asr_${name_tag} -name "*.pth"  -printf "%f\n" | awk '{print substr($0, 1, length($0) - 4)}')
 
 # for ckpt in valid.acc.best; do
-for ckpt in $all_ckpts ; do
-    if [[ $ckpt == "checkpoint" || $ckpt == "latest"  || $ckpt == "valid.acc.best" ]]; then
+for ckpt in $all_ckpts; do
+    # We skip all symlinks
+    if [[ $ckpt == "checkpoint" || $ckpt == "latest"  || $ckpt == "valid.acc.best"  || $ckpt == "valid.acc.ave.pth" ]]; then
         continue
     fi
     ./asr.sh \
         --asr_tag $name_tag \
-        --expdir ${expdir} \
-        --dumpdir ${dumpdir} \
         --feats_normalize uttmvn \
         --stage 12 \
         --stop_stage 13 \
-        --ngpu 1 \
+        --ngpu 2 \
         --gpu_inference true \
         --nj 20 \
         --inference_nj 1 \
diff --git a/egs2/clotho_v2/asr1/run_pt.sh b/egs2/clotho_v2/asr1/run_pt.sh
index d14dae97c46..d678a087804 100755
--- a/egs2/clotho_v2/asr1/run_pt.sh
+++ b/egs2/clotho_v2/asr1/run_pt.sh
@@ -8,17 +8,19 @@ timestamp=$(date "+%Y%m%d.%H%M%S")
 wandb_init_args=""
 other_args="$@"
 
-expdir=/compute/babel-13-33/sbharad2/expdir
-dumpdir=/compute/babel-13-33/sbharad2/dumpdir
-local_data_opts=/compute/babel-13-33/sbharad2/expdir
+# expdir=/compute/babel-13-33/sbharad2/expdir
+# dumpdir=/compute/babel-13-33/sbharad2/dumpdir
+# local_data_opts=/compute/babel-13-33/sbharad2/expdir
+    # --local_data_opts "${local_data_opts}" \
+    # --expdir ${expdir} \
+    # --dumpdir ${dumpdir} \
+
 
 ./asr.sh \
-    --asr_tag pt.full.${timestamp} \
+    --asr_tag pt.${timestamp} \
     --asr_speech_fold_length 1600 \
-    --expdir ${expdir} \
-    --dumpdir ${dumpdir} \
     --feats_normalize uttmvn \
-    --stage 11 \
+    --stage 1 \
     --stop_stage 13 \
     --ngpu 2 \
     --gpu_inference true \
@@ -32,8 +34,7 @@ local_data_opts=/compute/babel-13-33/sbharad2/expdir
     --train_set pretrain \
     --valid_set validation \
     --test_sets "validation evaluation" \
-    --asr_config conf/beats_bart_full.yaml \
+    --asr_config conf/beats_bart_pt.yaml \
     --inference_asr_model valid.acc.best.pth \
     --asr_args "${wandb_init_args} ${other_args}" \
-    --local_data_opts "${local_data_opts}" \
     --local_score_opts "exp/asr_pt.${timestamp}/inference_beam_size10_ctc_weight0.0_hugging_face_decoderTrue_asr_model_valid.acc.best"
diff --git a/egs2/clotho_v2/asr1/run_upload.sh b/egs2/clotho_v2/asr1/run_upload.sh
new file mode 100755
index 00000000000..00d46cd180e
--- /dev/null
+++ b/egs2/clotho_v2/asr1/run_upload.sh
@@ -0,0 +1,47 @@
+#!/usr/bin/env bash
+set -euo pipefail
+# Runs model upload to huggingface
+
+# wandb_init_args="--use_wandb true --wandb_project DCASE_AAC --wandb_model_log_interval 0"
+wandb_init_args=""
+other_args="$@"
+
+expdir=/compute/babel-13-33/sbharad2/expdir
+dumpdir=/compute/babel-13-33/sbharad2/dumpdir
+local_data_opts=/compute/babel-13-33/sbharad2/expdir
+
+# asr_tag=pt.initfix.bigbatch512.lr2e-4.weighted12layers.20241103.145125
+# hf_repo=shikhar7ssu/dcase23.aac.pt
+
+asr_tag=ft_lr5e-5.initfix.bigbatch512.lr2e-4.weighted12layers.20241103.145125
+hf_repo=shikhar7ssu/dcase23.aac
+
+
+./asr.sh \
+    --asr_tag ${asr_tag} \
+    --hf_repo ${hf_repo} \
+    --skip_packing false \
+    --skip_upload_hf false \
+    --lang en \
+    --expdir ${expdir} \
+    --dumpdir ${dumpdir} \
+    --feats_normalize uttmvn \
+    --stage 14 \
+    --stop_stage 15 \
+    --asr_stats_dir ${expdir}/asr_stats_finetune \
+    --ngpu 2 \
+    --gpu_inference true \
+    --nj 8 \
+    --inference_nj 1 \
+    --max_wav_duration 30 \
+    --token_type hugging_face \
+    --use_lm false \
+    --hugging_face_model_name_or_path "facebook/bart-base" \
+    --inference_args "--beam_size 10 --ctc_weight 0.0 --hugging_face_decoder True" \
+    --train_set development_clotho \
+    --valid_set validation \
+    --test_sets "evaluation" \
+    --asr_config conf/beats_bart_ft.yaml \
+    --inference_asr_model valid.acc.ave.pth \
+    --asr_args "${wandb_init_args} ${other_args}" \
+    --local_data_opts "${local_data_opts}"

From 257ee2abb6b81143fa911b11443d2eb0958460a0 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Thu, 14 Nov 2024 03:52:31 +0000
Subject: [PATCH 15/47] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 egs2/clotho_v2/asr1/README.md | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/egs2/clotho_v2/asr1/README.md b/egs2/clotho_v2/asr1/README.md
index 57054793a5c..487e51503d5 100644
--- a/egs2/clotho_v2/asr1/README.md
+++ b/egs2/clotho_v2/asr1/README.md
@@ -1,6 +1,6 @@
 # Clotho Audio Captioning RECIPE
 
-This recipe implements the DCASE 2023 Automated Audio Captioning (AAC) task with a BEATs encoder BART decoder model on the Clotho_v2 dataset, very close to what is described in [this paper](https://dcase.community/documents/challenge2023/technical_reports/DCASE2023_Wu_31_t6a.pdf) and reusing part of code from the [original implementation](https://github.com/slSeanWU/beats-conformer-bart-audio-captioner?tab=readme-ov-file). 
+This recipe implements the DCASE 2023 Automated Audio Captioning (AAC) task with a BEATs encoder BART decoder model on the Clotho_v2 dataset, very close to what is described in [this paper](https://dcase.community/documents/challenge2023/technical_reports/DCASE2023_Wu_31_t6a.pdf) and reusing part of code from the [original implementation](https://github.com/slSeanWU/beats-conformer-bart-audio-captioner?tab=readme-ov-file).
 We also reuse the code from the [BEATs repository](https://github.com/microsoft/unilm/tree/master/beats) for this implementation.
 
 # Training Details and Requirements
@@ -11,11 +11,11 @@ Fine-tuning is performed only on the standard development set of Clotho_v2 datas
 
 ### Steps to run
 For running AAC pre-training
-1. Download AudioCaps and set the path to its root directory in db.sh (This recipe downloads clotho for you). 
+1. Download AudioCaps and set the path to its root directory in db.sh (This recipe downloads clotho for you).
 2. Download the BEATs checkpoint: [BEATs_iter3+](https://onedrive.live.com/?authkey=%21AGXnEG4l3mlIzfA&id=6B83B49411CA81A7%2125960&cid=6B83B49411CA81A7&parId=root&parQt=sharedby&o=OneUp) and change the `beats_ckpt_path` path in `conf/beats_bart_pt.yaml`
 3. Launch with `run_pt.sh`
 
-For running fine-tuning: 
+For running fine-tuning:
 1. Set the value of `pt_tag` in `run_ft.sh` and
 2. Launch with `run_ft.sh`
 
@@ -65,4 +65,4 @@ All the scripts above are setup with 2 GPUs but that can be changed with the `ng
  fer.error_prob      : 0.319917231798172
  spider_fl           : 0.2532289450419614
 =====================================================
-```
\ No newline at end of file
+```

From cd21d786d3c5965a7cf90929469f1a4a9a110f31 Mon Sep 17 00:00:00 2001
From: Shikhar Bharadwaj <sbharad2@andrew.cmu.edu>
Date: Thu, 14 Nov 2024 12:12:38 -0500
Subject: [PATCH 16/47] fix linting

---
 egs2/clotho_v2/asr1/local/data.sh             |  5 +-
 .../asr1/local/data_prep_audiocaps.py         | 10 +++-
 egs2/clotho_v2/asr1/local/data_prep_clotho.py |  3 +-
 .../local/data_prep_clotho_chatgpt_mixup.py   | 36 ++---------
 egs2/clotho_v2/asr1/local/evaluation.py       |  3 +-
 espnet2/asr/encoder/beats_encoder.py          | 60 ++++++++++++-------
 6 files changed, 56 insertions(+), 61 deletions(-)

diff --git a/egs2/clotho_v2/asr1/local/data.sh b/egs2/clotho_v2/asr1/local/data.sh
index 701054b560c..77229d7c873 100755
--- a/egs2/clotho_v2/asr1/local/data.sh
+++ b/egs2/clotho_v2/asr1/local/data.sh
@@ -85,11 +85,10 @@ fi
 
 SPLITS=(development_clotho validation evaluation) # map dev, val, eval to development, validation, evaluation
 N_REF=5
-for split_name in ${SPLITS[@]}; do
+for split_name in "${SPLITS[@]}"; do
     mkdir -p "data/${split_name}"
 done
 
-echo "$(which python)"
 # Prepare data in the Kaldi format, including three files:
 # text, wav.scp, utt2spk
 ##########
@@ -115,7 +114,7 @@ SPLITS+=(pretrain)
 # and 2 for val, eval - validation, evaluation
 
 # SORT ALL
-for split_name in ${SPLITS[@]}; do
+for split_name in "${SPLITS[@]}"; do
     for f in wav.scp utt2spk; do
         sort data/${split_name}/${f} -o data/${split_name}/${f}
     done
diff --git a/egs2/clotho_v2/asr1/local/data_prep_audiocaps.py b/egs2/clotho_v2/asr1/local/data_prep_audiocaps.py
index 822aa1d4d64..54853a782da 100644
--- a/egs2/clotho_v2/asr1/local/data_prep_audiocaps.py
+++ b/egs2/clotho_v2/asr1/local/data_prep_audiocaps.py
@@ -8,8 +8,10 @@
 
 if __name__ == "__main__":
     ROOT_DATA_DIR = sys.argv[1]
-    TRAIN_FILES = "train_split.json"  # this is a json with key "audiocap_ids" containing list of folder_ids
-    DATA_DIR = "data"  # this has folders each for a folder_id, containing audio.wav and metadata.json files
+    # this is a json with key "audiocap_ids" containing list of folder_ids
+    TRAIN_FILES = "train_split.json"
+    # this has folders each for a folder_id, containing audio.wav and metadata.json files
+    DATA_DIR = "data"
 
     LOCAL_DATA_DIR = "data/development_audiocaps"
 
@@ -57,5 +59,7 @@
                 if N_PROCESSED % 1000 == 0:
                     print(f"Processed {N_PROCESSED} audio files.")
     print(
-        f"Processed {N_PROCESSED} audio files. Skipped {N_VERY_LONG_AUDIO} long audio, {N_ZERO_LENGTH_AUDIO} empty audio and {N_ZERO_LENGTH_CAPTIONS} empty caption items."
+        f"Processed {N_PROCESSED} audio files. Skipped {N_VERY_LONG_AUDIO}"
+        f"long audio, {N_ZERO_LENGTH_AUDIO} empty audio and "
+        f"{N_ZERO_LENGTH_CAPTIONS} empty caption items."
     )
diff --git a/egs2/clotho_v2/asr1/local/data_prep_clotho.py b/egs2/clotho_v2/asr1/local/data_prep_clotho.py
index d3d29391802..6cddcf58a87 100644
--- a/egs2/clotho_v2/asr1/local/data_prep_clotho.py
+++ b/egs2/clotho_v2/asr1/local/data_prep_clotho.py
@@ -78,5 +78,6 @@
                 for f in text_f:
                     f.close()
         print(
-            f"For split {data_split}: Processed {N_PROCESSED} audio and their captions. {len(missing_audio)} audio files were missing."
+            f"For split {data_split}: Processed {N_PROCESSED} audio and their"
+            f" captions. {len(missing_audio)} audio files were missing."
         )
diff --git a/egs2/clotho_v2/asr1/local/data_prep_clotho_chatgpt_mixup.py b/egs2/clotho_v2/asr1/local/data_prep_clotho_chatgpt_mixup.py
index c37587d8f89..420f51cbd26 100644
--- a/egs2/clotho_v2/asr1/local/data_prep_clotho_chatgpt_mixup.py
+++ b/egs2/clotho_v2/asr1/local/data_prep_clotho_chatgpt_mixup.py
@@ -2,11 +2,11 @@
 import os
 import random
 import sys
+import time
 from collections import defaultdict
 from multiprocessing import Pool
 from string import punctuation
 from typing import List, Optional
-
 import librosa
 import numpy as np
 import pandas as pd
@@ -158,28 +158,6 @@ def mixup_audio(self, audio_a, audio_b, sample_name):
             audio_a = pyln.normalize.peak(audio_a, base_db)
             audio_b = pyln.normalize.peak(audio_b, base_db + mix_db_ratio)
 
-        # sf.write(
-        #     os.path.join("tmp", f"audio_{base_db:.2f}_{mix_db_ratio:.2f}_A_prenorm.wav"),
-        #     audio_a,
-        #     SAMPLE_RATE
-        # )
-        # sf.write(
-        #     os.path.join("tmp", f"audio_{base_db:.2f}_{mix_db_ratio:.2f}_B_prenorm.wav"),
-        #     audio_b,
-        #     SAMPLE_RATE
-        # )
-
-        # sf.write(
-        #     os.path.join("tmp", f"audio_{base_db:.2f}_{mix_db_ratio:.2f}_A.wav"),
-        #     audio_a,
-        #     SAMPLE_RATE
-        # )
-        # sf.write(
-        #     os.path.join("tmp", f"audio_{base_db:.2f}_{mix_db_ratio:.2f}_B.wav"),
-        #     audio_b,
-        #     SAMPLE_RATE
-        # )
-
         if len(audio_a) > len(audio_b):
             st_idx = random.choice(range(len(audio_a) - len(audio_b) + 1))
             audio_a[st_idx : st_idx + len(audio_b)] += mix_scale * audio_b
@@ -192,11 +170,10 @@ def mixup_audio(self, audio_a, audio_b, sample_name):
         try:
             sf.write(
                 mixed_audio_path,
-                # os.path.join("tmp", f"{base_db:.2f}_{mix_db_ratio:.2f}_mixed.wav"),
                 mixed_audio,
                 SAMPLE_RATE,
             )
-        except:
+        except Exception as e:
             print(f"Can not write audio {mixed_audio_path}")
             mixed_audio = None
             mixed_audio_path = None
@@ -229,8 +206,8 @@ def read_lemma_labels(self, sample_name, sample_idx):
         lemmas = self.lemma_data[sample_name][sample_idx]["caption_kw_labels"]
         labels = np.zeros((self.n_lemmas,))
 
-        for l in lemmas:
-            labels[l] = 1
+        for lem in lemmas:
+            labels[lem] = 1
 
         return labels
 
@@ -324,7 +301,8 @@ def __getitem__(self, index):
 clotho_audio_base_dir = sys.argv[2]
 
 # create directory paths for storing clotho mixup audio
-audio_mixup_write_dir = f"{clotho_audio_base_dir}_audio_mixup/development/"  # "downloads/CLOTHO_v2.1_audio_mixup/development/"
+# "downloads/CLOTHO_v2.1_audio_mixup/development/" < This is default.
+audio_mixup_write_dir = f"{clotho_audio_base_dir}_audio_mixup/development/"
 os.makedirs(os.path.dirname(audio_mixup_write_dir), exist_ok=True)
 
 # Create mixup dataset
@@ -348,7 +326,6 @@ def __getitem__(self, index):
 output_utt2spk_path = os.path.join(data_base, "utt2spk")
 
 os.makedirs(os.path.dirname(output_txt_path), exist_ok=True)
-import time
 
 start = time.time()
 
@@ -386,4 +363,3 @@ def _get_ith_item(i):
 
 print(f"Number of files skipped {N_ERROR}")
 print(f"Number of files processed {N_PROCESSED}")
-# TODO(shikhar): change location of downloaded data to downloads, check gitignore for more info.
diff --git a/egs2/clotho_v2/asr1/local/evaluation.py b/egs2/clotho_v2/asr1/local/evaluation.py
index 257f6f7f105..ef3f3cbd6ea 100644
--- a/egs2/clotho_v2/asr1/local/evaluation.py
+++ b/egs2/clotho_v2/asr1/local/evaluation.py
@@ -1,4 +1,5 @@
-"""Wrapper around aac_metrics (https://github.com/Labbeti/aac-metrics) to evaluate the predictions."""
+"""Wrapper around aac_metrics (https://github.com/Labbeti/aac-metrics)
+to evaluate the audio captioning predictions."""
 
 import argparse
 
diff --git a/espnet2/asr/encoder/beats_encoder.py b/espnet2/asr/encoder/beats_encoder.py
index b71945072bc..c3a2b71da06 100644
--- a/espnet2/asr/encoder/beats_encoder.py
+++ b/espnet2/asr/encoder/beats_encoder.py
@@ -103,15 +103,21 @@ class BeatsEncoder(AbsEncoder):
 
     (https://arxiv.org/abs/2212.09058)
     Args:
-        beats_ckpt_path: Path to a pretrained BEATs checkpoint. If `beats_config` is provided and it does
-            not match the config in the checkpoint, code might throw an error.
-        adapter_layers: Number of transformer-based adapter layers after encoder.
-        max_layer: Propagate input through all layers for encoding if None. Otherwise use upto `max_layer`.
-        use_weighted_representation: Use weighted representations from max_layer if True. Weights are randomly initialized.
-        beats_config: `BEATsConfig` object. If provided, we will try to override the config in the
-            checkpoint. This can be used to change dropouts etc for fine-tuning the model while starting
-            from a pretrained checkpoint.
-        specaug_config: Dictionary containing parameters for SpecAugment. If provided, SpecAugment will be applied.
+        beats_ckpt_path: Path to a pretrained BEATs checkpoint. If
+            `beats_config` is provided and it does not match the
+            config in the checkpoint, code might throw an error.
+        adapter_layers: Number of transformer-based adapter layers
+            after encoder.
+        max_layer: Propagate input through all layers for encoding
+            if None. Otherwise use upto `max_layer`.
+        use_weighted_representation: Use weighted representations
+            from max_layer if True. Weights are randomly initialized.
+        beats_config: `BEATsConfig` object. If provided, we will try
+            to override the config in the checkpoint. This can be used
+            to change dropouts etc for fine-tuning the model while
+            starting from a pretrained checkpoint.
+        specaug_config: Dictionary containing parameters for SpecAugment.
+            If provided, SpecAugment will be applied.
     """
 
     def __init__(
@@ -133,8 +139,10 @@ def __init__(
 
         # Four cases for loading Beats config:
         # 1. No checkpoint and no config: Default config
-        # 2. Checkpoint and no user-provided config: Load config from checkpoint
-        # 3. Checkpoint and user-provided config: Merge the two, but override with user-provided config
+        # 2. Checkpoint and no user-provided config: Load config from
+        #    checkpoint
+        # 3. Checkpoint and user-provided config: Merge the two, but
+        #    override with user-provided config
         # 4. No checkpoint and user-provided config: Use user-provided config
         config = BEATsConfig()  # Default config
         if beats_ckpt_path and beats_config:
@@ -180,7 +188,8 @@ def __init__(
         if self.use_weighted_representation:
             if self.max_layer is None:
                 logging.warning(
-                    f"max_layer must be provided when using weighted representations. Set to {config.encoder_layers}."
+                    f"max_layer must be provided when using weighted representations."
+                    f"Set to {config.encoder_layers}."
                 )
                 self.max_layer = config.encoder_layers
             self.layer_weights = nn.Parameter(
@@ -200,7 +209,8 @@ def reload_pretrained_parameters(self):
         The initialization occurs in three steps:
         1. ESPNet initializes all modules.
         2. This function initializes BEATs encoder overriding 1.
-        3. Optionally, if we have the pretrained checkpoint, we load the weights from the checkpoint overriding 2 and 1.
+        3. Optionally, if we have the pretrained checkpoint, we load the
+            weights from the checkpoint overriding 2 and 1.
         """
         logging.info("BEATs Initialization function called.")
         if self.post_extract_proj:
@@ -211,7 +221,8 @@ def reload_pretrained_parameters(self):
         if self.patch_embedding.bias is not None:
             torch.nn.init.constant_(self.patch_embedding.bias, 0)
 
-        # BEATs has different initialization from ESPNet for other modules, so override.
+        # BEATs has different initialization from ESPNet for other modules,
+        #  so override.
         self.encoder.apply(init_bert_params)
         if self.loaded_state_dict_ is not None:
 
@@ -220,9 +231,12 @@ def reload_pretrained_parameters(self):
             )
             # strict=False to ignore Weights in the predictor
             logging.info(
-                f"Loaded BEATs pretrained model. Following keys were missing in your custom model: {load_info.missing_keys}. "
-                f"Follwing keys could not be loaded from the pretrained checkpoint: {load_info.unexpected_keys}."
-                "It is expected to have 'predictor' listed above if you are fine-tuning with only the BEATs backbone."
+                f"Loaded BEATs pretrained model. Following keys were missing"
+                f" in your custom model: {load_info.missing_keys}. "
+                f"Follwing keys could not be loaded from the pretrained"
+                f"checkpoint: {load_info.unexpected_keys}."
+                "It is expected to have 'predictor' listed above if you are"
+                "fine-tuning with only the BEATs backbone."
             )
 
     def forward_padding_mask(
@@ -269,9 +283,9 @@ def forward(
     ) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]:
         """Wrapper for compatibility with ESPNets' AbsEncoder Interface."""
 
-        # NOTE(shikhar): If xs is not provided then the operation is costly, because this
-        # function tries to create a tensor of size maxlen x maxlen. Therfore, we unsqueeze
-        # and then squeeze tensors.
+        # NOTE(shikhar): If xs is not provided then the operation is costly,
+        # because this function tries to create a tensor of size maxlen x maxlen.
+        # Therfore, we unsqueeze and then squeeze tensors.
         mask = make_pad_mask(
             lengths=ilens, xs=xs_pad.unsqueeze(-1).unsqueeze(-1), length_dim=1
         ).to(xs_pad.device)
@@ -297,8 +311,6 @@ def extract_features(
             fbank = self.preprocess(source)
 
             if self.specaug is not None and self.training:
-                # NOTE(shikhar): We don't provide lengths because espnet implementation
-                # does not use it --> /espnet/espnet2/layers/mask_along_axis.py
                 fbank = self.specaug(fbank)[0]
 
         if padding_mask is not None:
@@ -1219,7 +1231,9 @@ def __init__(self, input_dim, output_dim, glu_type="sigmoid", bias_in_glu=True):
 
     def forward(self, x):
         """Forward pass"""
-        # to be consistent with GLU_Linear, we assume the input always has the #channel (#dim) in the last dimension of the tensor, so need to switch the dimension first for 1D-Conv case
+        # to be consistent with GLU_Linear, we assume the input always has the
+        # #channel (#dim) in the last dimension of the tensor, so need to
+        # switch the dimension first for 1D-Conv case
         x = self.linear(x)
 
         if self.glu_type == "bilinear":

From 0cf22727cad74e0caf19a65f7486fd81b416fdae Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Thu, 14 Nov 2024 17:14:17 +0000
Subject: [PATCH 17/47] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 egs2/clotho_v2/asr1/local/data_prep_clotho_chatgpt_mixup.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/egs2/clotho_v2/asr1/local/data_prep_clotho_chatgpt_mixup.py b/egs2/clotho_v2/asr1/local/data_prep_clotho_chatgpt_mixup.py
index 420f51cbd26..ccb353200a6 100644
--- a/egs2/clotho_v2/asr1/local/data_prep_clotho_chatgpt_mixup.py
+++ b/egs2/clotho_v2/asr1/local/data_prep_clotho_chatgpt_mixup.py
@@ -7,6 +7,7 @@
 from multiprocessing import Pool
 from string import punctuation
 from typing import List, Optional
+
 import librosa
 import numpy as np
 import pandas as pd

From 736f1d24efc16aed3d7f8ae7529705ff1addd912 Mon Sep 17 00:00:00 2001
From: Shikhar Bharadwaj <sbharad2@andrew.cmu.edu>
Date: Thu, 14 Nov 2024 13:35:23 -0500
Subject: [PATCH 18/47] fix remaining linting issue

---
 egs2/clotho_v2/asr1/local/data_prep_audiocaps.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/egs2/clotho_v2/asr1/local/data_prep_audiocaps.py b/egs2/clotho_v2/asr1/local/data_prep_audiocaps.py
index 54853a782da..0d2dc2c5638 100644
--- a/egs2/clotho_v2/asr1/local/data_prep_audiocaps.py
+++ b/egs2/clotho_v2/asr1/local/data_prep_audiocaps.py
@@ -10,7 +10,8 @@
     ROOT_DATA_DIR = sys.argv[1]
     # this is a json with key "audiocap_ids" containing list of folder_ids
     TRAIN_FILES = "train_split.json"
-    # this has folders each for a folder_id, containing audio.wav and metadata.json files
+    # this has folders each for a folder_id, containing audio.wav and
+    # metadata.json files
     DATA_DIR = "data"
 
     LOCAL_DATA_DIR = "data/development_audiocaps"

From 2e098be1ff475a86d21a3b2c1ccc13192dad6e3f Mon Sep 17 00:00:00 2001
From: Shikhar Bharadwaj <sbharad2@andrew.cmu.edu>
Date: Fri, 15 Nov 2024 10:12:24 -0500
Subject: [PATCH 19/47] handle comments
 https://github.com/espnet/espnet/pull/5915\#pullrequestreview-2437805856

---
 egs2/clotho_v2/asr1/README.md               |  8 +-
 egs2/clotho_v2/asr1/cmd.sh                  |  2 +-
 egs2/clotho_v2/asr1/conf/beats_bart_ft.yaml |  3 +-
 egs2/clotho_v2/asr1/local/Checker.ipynb     | 91 ---------------------
 egs2/clotho_v2/asr1/local/data.sh           | 58 +++++++------
 egs2/clotho_v2/asr1/run_ft.sh               | 14 +---
 egs2/clotho_v2/asr1/run_inference.sh        | 45 ----------
 egs2/clotho_v2/asr1/run_pt.sh               |  7 --
 egs2/clotho_v2/asr1/run_upload.sh           | 47 -----------
 9 files changed, 40 insertions(+), 235 deletions(-)
 delete mode 100644 egs2/clotho_v2/asr1/local/Checker.ipynb
 delete mode 100755 egs2/clotho_v2/asr1/run_inference.sh
 delete mode 100755 egs2/clotho_v2/asr1/run_upload.sh

diff --git a/egs2/clotho_v2/asr1/README.md b/egs2/clotho_v2/asr1/README.md
index 487e51503d5..9ebce036b51 100644
--- a/egs2/clotho_v2/asr1/README.md
+++ b/egs2/clotho_v2/asr1/README.md
@@ -15,16 +15,16 @@ For running AAC pre-training
 2. Download the BEATs checkpoint: [BEATs_iter3+](https://onedrive.live.com/?authkey=%21AGXnEG4l3mlIzfA&id=6B83B49411CA81A7%2125960&cid=6B83B49411CA81A7&parId=root&parQt=sharedby&o=OneUp) and change the `beats_ckpt_path` path in `conf/beats_bart_pt.yaml`
 3. Launch with `run_pt.sh`
 
-For running fine-tuning:
+For running fine-tuning, inference and scoring:
 1. Set the value of `pt_tag` in `run_ft.sh` and
 2. Launch with `run_ft.sh`
 
-For running inference and scoring on all checkpoints `run_inference.sh` can be used by just changing the `name_tag`.
 
 ## Trained checkpoints
-AAC Pre-trained model: https://huggingface.co/shikhar7ssu/dcase23.aac.pt
+AAC Pre-trained model: https://huggingface.co/espnet/DCASE23.AudioCaptioning.PreTrained 
 
-Fine-tuned model: https://huggingface.co/shikhar7ssu/dcase23.aac
+
+Fine-tuned model: https://huggingface.co/espnet/DCASE23.AudioCaptioning.FineTuned 
 
 
 ### GPU Time
diff --git a/egs2/clotho_v2/asr1/cmd.sh b/egs2/clotho_v2/asr1/cmd.sh
index 257c6c739df..2aae6919fef 100644
--- a/egs2/clotho_v2/asr1/cmd.sh
+++ b/egs2/clotho_v2/asr1/cmd.sh
@@ -28,7 +28,7 @@
 
 
 # Select the backend used by run.sh from "local", "stdout", "sge", "slurm", or "ssh"
-cmd_backend='slurm'
+cmd_backend='local'
 
 # Local machine, without any Job scheduling system
 if [ "${cmd_backend}" = local ]; then
diff --git a/egs2/clotho_v2/asr1/conf/beats_bart_ft.yaml b/egs2/clotho_v2/asr1/conf/beats_bart_ft.yaml
index 7e01b01f665..222a14e941a 100644
--- a/egs2/clotho_v2/asr1/conf/beats_bart_ft.yaml
+++ b/egs2/clotho_v2/asr1/conf/beats_bart_ft.yaml
@@ -29,7 +29,8 @@ freeze_param: [
 
 encoder: beats
 encoder_conf:
-    beats_ckpt_path: /compute/babel-13-33/sbharad2/models/BEATs/BEATs_iter3_plus_AS2M_finetuned_on_AS2M_cpt1.pt # /u/sbharadwaj/BEATs_models/BEATs_iter3_plus_AS2M.pt # _finetuned_on_AS2M_cpt1
+    # Please download the BEATs model from https://github.com/microsoft/unilm/tree/master/beats (iter3+, Beats finetuned model 1) and update the path below
+    beats_ckpt_path: /compute/babel-13-33/sbharad2/models/BEATs/BEATs_iter3_plus_AS2M_finetuned_on_AS2M_cpt1.pt
     specaug_config:
         apply_freq_mask: true
         freq_mask_width_range:
diff --git a/egs2/clotho_v2/asr1/local/Checker.ipynb b/egs2/clotho_v2/asr1/local/Checker.ipynb
deleted file mode 100644
index 231fa5bf132..00000000000
--- a/egs2/clotho_v2/asr1/local/Checker.ipynb
+++ /dev/null
@@ -1,91 +0,0 @@
-{
- "cells": [
-  {
-   "cell_type": "code",
-   "execution_count": 14,
-   "id": "1b6f0ae7-c574-4da9-8ed2-01513728be7e",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import pandas as pd\n",
-    "import os\n",
-    "from scipy.io import wavfile\n",
-    "from collections import defaultdict as dd"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 18,
-   "id": "fe8ed4fc-4561-4e71-aa1d-8f41e61c7633",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "base_path = \"/data/user_data/sbharad2/espnet/egs2/clotho_v2/asr1/local/CLOTHO_v2.1/\"\n",
-    "splits = ['development', 'evaluation','test','validation']\n",
-    "\n",
-    "def get_audio_duration(fname):\n",
-    "    samplerate, data = wavfile.read(fname)\n",
-    "    duration_s = len(data)/samplerate\n",
-    "    return duration_s\n",
-    "\n",
-    "# Audio length\n",
-    "audio_s = dd(lambda : [])\n",
-    "for split in splits:\n",
-    "    for fname in os.listdir(os.path.join(base_path,'clotho_audio_files',split)):\n",
-    "        ffname = os.path.join(base_path,'clotho_audio_files',split,fname)\n",
-    "        audio_s[split].append(get_audio_duration(ffname))"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 23,
-   "id": "59ea5e97-0f9d-4931-b28c-699513e38318",
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "development 30.0 15.0\n",
-      "evaluation 30.0 15.003174603174603\n",
-      "test 30.0 15.000090702947846\n",
-      "validation 29.95390022675737 15.05\n"
-     ]
-    }
-   ],
-   "source": [
-    "for split in splits:\n",
-    "    print(split,max(audio_s[split]),min(audio_s[split]))"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "b3884f12-1a4a-41ef-b4c9-e7c99af27086",
-   "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.8.19"
-  }
- },
- "nbformat": 4,
- "nbformat_minor": 5
-}
diff --git a/egs2/clotho_v2/asr1/local/data.sh b/egs2/clotho_v2/asr1/local/data.sh
index 77229d7c873..8b64b6a8453 100755
--- a/egs2/clotho_v2/asr1/local/data.sh
+++ b/egs2/clotho_v2/asr1/local/data.sh
@@ -25,39 +25,45 @@ fi
 
 log "stage 1: Data preparation"
 
-## DOWNLOAD DATA
-# If there is no argument, the default download directory is set to currentdir/downloads
-if [ $# -ne 1 ]; then
-    CLOTHO_V2_ROOT_DIR="$(pwd)/downloads"
-    log "Using the default download directory: ${CLOTHO_V2_ROOT_DIR}"
-else
-    CLOTHO_V2_ROOT_DIR="$1/downloads"
-    log "Using the specified download directory: ${CLOTHO_V2_ROOT_DIR}"
-fi
-
 
-if [ ! -e "${CLOTHO_V2_ROOT_DIR}/download_done" ]; then
-    log "stage 1: Data preparation - Installing aac-datasets"
-    if ! pip3 install aac-datasets; then
-        echo "Error: Installing aac-datasets failed."
-        exit 1
+## DOWNLOAD DATA if CLOTHO_V2 is set to downloads
+if [ "${CLOTHO_V2}" == "downloads" ]; then
+    # If there is no argument, the default download directory is set to currentdir/downloads
+    if [ $# -ne 1 ]; then
+        CLOTHO_V2_ROOT_DIR="$(pwd)/downloads"
+        log "Using the default download directory: ${CLOTHO_V2_ROOT_DIR}"
+    else
+        CLOTHO_V2_ROOT_DIR="$1/downloads"
+        log "Using the specified download directory: ${CLOTHO_V2_ROOT_DIR}"
     fi
-    echo "Downlaoding clotho into ${CLOTHO_V2_ROOT_DIR}."
-    mkdir -p "${CLOTHO_V2_ROOT_DIR}"
-    for split in val eval dev; do
-        echo "Downloading ${split} split."
-        if ! aac-datasets-download --root "${CLOTHO_V2_ROOT_DIR}" clotho --subsets "${split}"; then
-            echo "Error: Downloading Clotho dataset failed."
+
+
+    if [ ! -e "${CLOTHO_V2_ROOT_DIR}/download_done" ]; then
+        log "stage 1: Data preparation - Installing aac-datasets"
+        if ! pip3 install aac-datasets; then
+            echo "Error: Installing aac-datasets failed."
             exit 1
         fi
-    done
-    touch "${CLOTHO_V2_ROOT_DIR}/download_done"
+        echo "Downlaoding clotho into ${CLOTHO_V2_ROOT_DIR}."
+        mkdir -p "${CLOTHO_V2_ROOT_DIR}"
+        for split in val eval dev; do
+            echo "Downloading ${split} split."
+            if ! aac-datasets-download --root "${CLOTHO_V2_ROOT_DIR}" clotho --subsets "${split}"; then
+                echo "Error: Downloading Clotho dataset failed."
+                exit 1
+            fi
+        done
+        touch "${CLOTHO_V2_ROOT_DIR}/download_done"
+    else
+        echo "Clotho dataset is already downloaded. ${CLOTHO_V2_ROOT_DIR}/download_done exists."
+    fi
+
+    CLOTHO_V2_ROOT_DIR="${CLOTHO_V2_ROOT_DIR}/CLOTHO_v2.1"
 else
-    echo "Clotho dataset is already downloaded. ${CLOTHO_V2_ROOT_DIR}/download_done exists."
+    CLOTHO_V2_ROOT_DIR=${CLOTHO_V2}
+    log "Using the specified data directory: ${CLOTHO_V2_ROOT_DIR}"
 fi
 
-CLOTHO_V2_ROOT_DIR="${CLOTHO_V2_ROOT_DIR}/CLOTHO_v2.1"
-
 ## PREPARE DATA
 if [ ! -d ${CLOTHO_V2_ROOT_DIR} ]; then
     echo Cannot find ${CLOTHO_V2_ROOT_DIR} directory! Exiting...
diff --git a/egs2/clotho_v2/asr1/run_ft.sh b/egs2/clotho_v2/asr1/run_ft.sh
index 2b2083779ae..c03221c1aa3 100755
--- a/egs2/clotho_v2/asr1/run_ft.sh
+++ b/egs2/clotho_v2/asr1/run_ft.sh
@@ -2,22 +2,10 @@
 set -euo pipefail
 # Runs fine-tuning
 
-# wandb_init_args="--use_wandb true --wandb_project DCASE_AAC --wandb_model_log_interval 0"
 wandb_init_args=""
 other_args="$@"
 
-# expdir=/compute/babel-13-33/sbharad2/expdir
-# dumpdir=/compute/babel-13-33/sbharad2/dumpdir
-# local_data_opts=/compute/babel-13-33/sbharad2/expdir
-#     --expdir ${expdir} \
-#     --dumpdir ${dumpdir} \
-#     --local_data_opts "${local_data_opts}" \
-
-# pre_trained_model_path=exp/asr_pt.initfix.bigbatch512.lr2e-4.weighted12layers.20241103.145125/valid.acc.ave_5best.pth
-# pt_tag=full.20241111.170406
-# pt_tag=data_check.20241111.153210
-
-pt_tag=beats.bartfrozen.20241111.164042
+pt_tag=pt.initfix.bigbatch512.lr2e-4.weighted12layers.20241103.145125
 
 ckpt_name=valid.acc.ave
 
diff --git a/egs2/clotho_v2/asr1/run_inference.sh b/egs2/clotho_v2/asr1/run_inference.sh
deleted file mode 100755
index e1aa1a24fd0..00000000000
--- a/egs2/clotho_v2/asr1/run_inference.sh
+++ /dev/null
@@ -1,45 +0,0 @@
-#!/usr/bin/env bash
-# Runs pre-training
-set -euo pipefail
-
-# # Run inference on all checkpoints
-
-wandb_sweep_args="$@"
-
-
-# expdir=/compute/babel-13-33/sbharad2/expdir
-# dumpdir=/compute/babel-13-33/sbharad2/dumpdir
-#         --expdir ${expdir} \
-#         --dumpdir ${dumpdir} \
-
-
-name_tag=ft.beats.bartfrozen.20241111.164042.valid.acc.ave
-all_ckpts=$(find ${expdir}/asr_${name_tag} -name "*.pth"  -printf "%f\n" | awk '{print substr($0, 1, length($0) - 4)}')
-
-# for ckpt in valid.acc.best; do
-for ckpt in $all_ckpts; do
-    # We skip all symlinks
-    if [[ $ckpt == "checkpoint" || $ckpt == "latest"  || $ckpt == "valid.acc.best"  || $ckpt == "valid.acc.ave.pth" ]]; then
-        continue
-    fi
-    ./asr.sh \
-        --asr_tag $name_tag \
-        --feats_normalize uttmvn \
-        --stage 12 \
-        --stop_stage 13 \
-        --ngpu 2 \
-        --gpu_inference true \
-        --nj 20 \
-        --inference_nj 1 \
-        --max_wav_duration 30 \
-        --token_type hugging_face \
-        --use_lm false \
-        --hugging_face_model_name_or_path "facebook/bart-base" \
-        --inference_args "--beam_size 10 --ctc_weight 0.0 --hugging_face_decoder True" \
-        --train_set pretrain \
-        --valid_set validation \
-        --test_sets "evaluation" \
-        --inference_asr_model ${ckpt}.pth \
-        --asr_args "${wandb_sweep_args}" \
-        --local_score_opts "${expdir}/asr_${name_tag}/inference_beam_size10_ctc_weight0.0_hugging_face_decoderTrue_asr_model_${ckpt}"
-done
diff --git a/egs2/clotho_v2/asr1/run_pt.sh b/egs2/clotho_v2/asr1/run_pt.sh
index d678a087804..f19e18c1d17 100755
--- a/egs2/clotho_v2/asr1/run_pt.sh
+++ b/egs2/clotho_v2/asr1/run_pt.sh
@@ -8,13 +8,6 @@ timestamp=$(date "+%Y%m%d.%H%M%S")
 wandb_init_args=""
 other_args="$@"
 
-# expdir=/compute/babel-13-33/sbharad2/expdir
-# dumpdir=/compute/babel-13-33/sbharad2/dumpdir
-# local_data_opts=/compute/babel-13-33/sbharad2/expdir
-    # --local_data_opts "${local_data_opts}" \
-    # --expdir ${expdir} \
-    # --dumpdir ${dumpdir} \
-
 
 ./asr.sh \
     --asr_tag pt.${timestamp} \
diff --git a/egs2/clotho_v2/asr1/run_upload.sh b/egs2/clotho_v2/asr1/run_upload.sh
deleted file mode 100755
index 00d46cd180e..00000000000
--- a/egs2/clotho_v2/asr1/run_upload.sh
+++ /dev/null
@@ -1,47 +0,0 @@
-#!/usr/bin/env bash
-set -euo pipefail
-# Runs model upload to huggingface
-
-# wandb_init_args="--use_wandb true --wandb_project DCASE_AAC --wandb_model_log_interval 0"
-wandb_init_args=""
-other_args="$@"
-
-expdir=/compute/babel-13-33/sbharad2/expdir
-dumpdir=/compute/babel-13-33/sbharad2/dumpdir
-local_data_opts=/compute/babel-13-33/sbharad2/expdir
-
-# asr_tag=pt.initfix.bigbatch512.lr2e-4.weighted12layers.20241103.145125
-# hf_repo=shikhar7ssu/dcase23.aac.pt
-
-asr_tag=ft_lr5e-5.initfix.bigbatch512.lr2e-4.weighted12layers.20241103.145125
-hf_repo=shikhar7ssu/dcase23.aac
-
-
-./asr.sh \
-    --asr_tag ${asr_tag} \
-    --hf_repo ${hf_repo} \
-    --skip_packing false \
-    --skip_upload_hf false \
-    --lang en \
-    --expdir ${expdir} \
-    --dumpdir ${dumpdir} \
-    --feats_normalize uttmvn \
-    --stage 14 \
-    --stop_stage 15 \
-    --asr_stats_dir ${expdir}/asr_stats_finetune \
-    --ngpu 2 \
-    --gpu_inference true \
-    --nj 8 \
-    --inference_nj 1 \
-    --max_wav_duration 30 \
-    --token_type hugging_face \
-    --use_lm false \
-    --hugging_face_model_name_or_path "facebook/bart-base" \
-    --inference_args "--beam_size 10 --ctc_weight 0.0 --hugging_face_decoder True" \
-    --train_set development_clotho \
-    --valid_set validation \
-    --test_sets "evaluation" \
-    --asr_config conf/beats_bart_ft.yaml \
-    --inference_asr_model valid.acc.ave.pth \
-    --asr_args "${wandb_init_args} ${other_args}" \
-    --local_data_opts "${local_data_opts}"

From 2a6f3435a29a89d9ddc136ea7f3258abb89416c6 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Fri, 15 Nov 2024 15:13:36 +0000
Subject: [PATCH 20/47] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 egs2/clotho_v2/asr1/README.md | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/egs2/clotho_v2/asr1/README.md b/egs2/clotho_v2/asr1/README.md
index 9ebce036b51..d7fab4a350c 100644
--- a/egs2/clotho_v2/asr1/README.md
+++ b/egs2/clotho_v2/asr1/README.md
@@ -21,10 +21,10 @@ For running fine-tuning, inference and scoring:
 
 
 ## Trained checkpoints
-AAC Pre-trained model: https://huggingface.co/espnet/DCASE23.AudioCaptioning.PreTrained 
+AAC Pre-trained model: https://huggingface.co/espnet/DCASE23.AudioCaptioning.PreTrained
 
 
-Fine-tuned model: https://huggingface.co/espnet/DCASE23.AudioCaptioning.FineTuned 
+Fine-tuned model: https://huggingface.co/espnet/DCASE23.AudioCaptioning.FineTuned
 
 
 ### GPU Time

From 92d2e5793c53e234cdb423987ee65f8af4875e3c Mon Sep 17 00:00:00 2001
From: Shikhar Bharadwaj <sbharad2@andrew.cmu.edu>
Date: Fri, 15 Nov 2024 10:29:23 -0500
Subject: [PATCH 21/47] handle comment:
 https://github.com/espnet/espnet/pull/5915\#discussion_r1844016571

---
 egs2/TEMPLATE/asr1/db.sh      | 14 ++------------
 egs2/clotho_v2/asr1/README.md |  2 +-
 2 files changed, 3 insertions(+), 13 deletions(-)

diff --git a/egs2/TEMPLATE/asr1/db.sh b/egs2/TEMPLATE/asr1/db.sh
index 1073f873c12..f00fcc857bb 100755
--- a/egs2/TEMPLATE/asr1/db.sh
+++ b/egs2/TEMPLATE/asr1/db.sh
@@ -221,18 +221,8 @@ VOXCELEB=
 KSPONSPEECH=
 HIFITTS=downloads
 CLOTHO_V2=downloads
-
-# For only Delta environment
-if [[ "$(hostname -d)" == delta.ncsa.illinois.edu ]]; then
-    AUDIOCAPS=/old_scratch/bbjs/shared/corpora/audiocaps/
-    CLOTHO_CHATGPT_MIXUP=/old_scratch/bbjs/slseanwu/dcase23_aac/beats_baseline/chatgpt/
-fi
-
-# For only CMU Babel environment on compute nodes
-if [[ "$(hostname)" == babel* ]]; then
-    AUDIOCAPS=/compute/babel-13-33/sbharad2/audiocaps/
-    CLOTHO_CHATGPT_MIXUP=/compute/babel-13-33/sbharad2/chatgpt/
-fi
+AUDIOCAPS=
+CLOTHO_CHATGPT_MIXUP=
 
 # For only CMU TIR environment
 if [[ "$(hostname)" == tir* ]]; then
diff --git a/egs2/clotho_v2/asr1/README.md b/egs2/clotho_v2/asr1/README.md
index 9ebce036b51..9fa4f6f489c 100644
--- a/egs2/clotho_v2/asr1/README.md
+++ b/egs2/clotho_v2/asr1/README.md
@@ -11,7 +11,7 @@ Fine-tuning is performed only on the standard development set of Clotho_v2 datas
 
 ### Steps to run
 For running AAC pre-training
-1. Download AudioCaps and set the path to its root directory in db.sh (This recipe downloads clotho for you).
+1. Download AudioCaps and set the path to its root directory in db.sh (This recipe downloads clotho for you). Also download clotho mixup data (wav ids and captions) from [this repo](https://huggingface.co/datasets/slseanwu/clotho-chatgpt-mixup-50K) and set `CLOTHO_CHATGPT_MIXUP` in db.sh. The code will take care of mixing and creating the audio files.
 2. Download the BEATs checkpoint: [BEATs_iter3+](https://onedrive.live.com/?authkey=%21AGXnEG4l3mlIzfA&id=6B83B49411CA81A7%2125960&cid=6B83B49411CA81A7&parId=root&parQt=sharedby&o=OneUp) and change the `beats_ckpt_path` path in `conf/beats_bart_pt.yaml`
 3. Launch with `run_pt.sh`
 

From 45aab054a3c73c3f9264a908652561e0af2c2f96 Mon Sep 17 00:00:00 2001
From: Shikhar Bharadwaj <sbharad2@andrew.cmu.edu>
Date: Thu, 21 Nov 2024 15:07:30 -0500
Subject: [PATCH 22/47] add implementation from dcase23 code

---
 egs2/clotho_v2/asr1/cmd.sh                    |  2 +-
 egs2/clotho_v2/asr1/conf/beats_bart_ft.yaml   | 12 +--
 egs2/clotho_v2/asr1/conf/beats_bart_pt.yaml   | 17 ++--
 egs2/clotho_v2/asr1/conf/slurm.conf           |  2 +-
 egs2/clotho_v2/asr1/run_ft.sh                 | 15 +++-
 egs2/clotho_v2/asr1/run_inference.sh          | 38 +++++++++
 egs2/clotho_v2/asr1/run_pt.sh                 | 20 +++--
 egs2/clotho_v2/asr1/run_upload.sh             | 49 ++++++++++++
 .../hugging_face_transformers_decoder.py      | 11 ++-
 espnet2/asr/encoder/beats_encoder.py          | 78 +++++++++++++++----
 10 files changed, 198 insertions(+), 46 deletions(-)
 create mode 100755 egs2/clotho_v2/asr1/run_inference.sh
 create mode 100755 egs2/clotho_v2/asr1/run_upload.sh

diff --git a/egs2/clotho_v2/asr1/cmd.sh b/egs2/clotho_v2/asr1/cmd.sh
index 2aae6919fef..257c6c739df 100644
--- a/egs2/clotho_v2/asr1/cmd.sh
+++ b/egs2/clotho_v2/asr1/cmd.sh
@@ -28,7 +28,7 @@
 
 
 # Select the backend used by run.sh from "local", "stdout", "sge", "slurm", or "ssh"
-cmd_backend='local'
+cmd_backend='slurm'
 
 # Local machine, without any Job scheduling system
 if [ "${cmd_backend}" = local ]; then
diff --git a/egs2/clotho_v2/asr1/conf/beats_bart_ft.yaml b/egs2/clotho_v2/asr1/conf/beats_bart_ft.yaml
index 222a14e941a..ebd3c60190d 100644
--- a/egs2/clotho_v2/asr1/conf/beats_bart_ft.yaml
+++ b/egs2/clotho_v2/asr1/conf/beats_bart_ft.yaml
@@ -2,7 +2,7 @@ batch_type: folded
 unused_parameters: true
 batch_size: 64
 accum_grad: 2
-max_epoch: 50
+max_epoch: 40
 grad_clip: 2
 patience: none
 best_model_criterion:  # criterion to save best models
@@ -12,7 +12,7 @@ best_model_criterion:  # criterion to save best models
 keep_nbest_models: 5  # save nbest models and average these checkpoints
 use_amp: true    # whether to use automatic mixed precision
 num_att_plot: 0  # do not save attention plots to save time in the demo
-num_workers: 2   # number of workers in dataloader
+num_workers: 2   # number of workers in dataloader - was 128
 
 
 # BEATs implementation takes care of generating mel spectrogram, normalization and specaug
@@ -70,9 +70,11 @@ model_conf:
 
 optim: adamw
 optim_conf:
-    lr: 0.00005 # 5e-5
-    weight_decay: 0.001
+    lr: 0.00002 # 2e-5
+    weight_decay: 0.001 # 1e-3
 
 scheduler: warmuplr
 scheduler_conf:
-    warmup_steps: 4000
+    warmup_steps: 1000 #1k
+
+# https://github.com/slSeanWU/beats-conformer-bart-audio-captioner/blob/main/config/training_config.json
diff --git a/egs2/clotho_v2/asr1/conf/beats_bart_pt.yaml b/egs2/clotho_v2/asr1/conf/beats_bart_pt.yaml
index 4c146c1cfd5..9f1cf61a711 100644
--- a/egs2/clotho_v2/asr1/conf/beats_bart_pt.yaml
+++ b/egs2/clotho_v2/asr1/conf/beats_bart_pt.yaml
@@ -1,15 +1,15 @@
 batch_type: folded
 unused_parameters: true
-batch_size: 64
-accum_grad: 8
-max_epoch: 30
+batch_size: 32
+accum_grad: 2
+max_epoch: 20
 grad_clip: 2
 patience: none
 best_model_criterion:  # criterion to save best models
 -   - valid
     - acc
     - max
-keep_nbest_models: 5  # save nbest models and average these checkpoints
+keep_nbest_models: 10  # save nbest models and average these checkpoints
 use_amp: true    # whether to use automatic mixed precision
 num_att_plot: 0  # do not save attention plots to save time in the demo
 num_workers: 2   # number of workers in dataloader
@@ -43,8 +43,9 @@ encoder_conf:
         - 0
         - 0.12
         num_time_mask: 5
-    adapter_layers: 2
-    max_layer: 12
+    adapter_layers: 3
+    downsampling_rate: 3 # CNN downsampling over beats encoder
+    max_layer: 11 # 0 based index
     use_weighted_representation: true
 
 decoder: hugging_face_transformers
@@ -61,7 +62,7 @@ token_type: hugging_face
 model_conf:
     ctc_weight: 0.0  # No CTC, only attention branch
     lsm_weight: 0.1  # label smoothing weight
-    length_normalized_loss: true
+    length_normalized_loss: true  # length normalization
     # BART Dictionary customizations
     ignore_id: 1
     sym_blank: "<pad>"
@@ -76,4 +77,4 @@ optim_conf:
 
 scheduler: warmuplr
 scheduler_conf:
-    warmup_steps: 8000
+    warmup_steps: 2500 # 8000
diff --git a/egs2/clotho_v2/asr1/conf/slurm.conf b/egs2/clotho_v2/asr1/conf/slurm.conf
index 9ebbf6b3d90..3ae0a7ce304 100644
--- a/egs2/clotho_v2/asr1/conf/slurm.conf
+++ b/egs2/clotho_v2/asr1/conf/slurm.conf
@@ -12,5 +12,5 @@ option num_threads=* --cpus-per-task $0
 option num_threads=1 --cpus-per-task 1
 option num_nodes=* --nodes $0
 default gpu=0
-option gpu=0 -p general
+option gpu=0 -p array
 option gpu=* -p general --gres=gpu:$0 -c $0
diff --git a/egs2/clotho_v2/asr1/run_ft.sh b/egs2/clotho_v2/asr1/run_ft.sh
index c03221c1aa3..28cb4f516c8 100755
--- a/egs2/clotho_v2/asr1/run_ft.sh
+++ b/egs2/clotho_v2/asr1/run_ft.sh
@@ -5,18 +5,25 @@ set -euo pipefail
 wandb_init_args=""
 other_args="$@"
 
-pt_tag=pt.initfix.bigbatch512.lr2e-4.weighted12layers.20241103.145125
+expdir=/compute/babel-13-33/sbharad2/expdir
+dumpdir=/compute/babel-13-33/sbharad2/dumpdir
+
+pt_tag=scratch_bart.splen30s.lr2e-4.lossnormT.20241120.113058
 
 ckpt_name=valid.acc.ave
+ckpt_name=10epoch
 
 pre_trained_model_path=${expdir}/asr_pt.${pt_tag}/${ckpt_name}.pth
-ft_tag=${pt_tag}.${ckpt_name}
-
+ft_tag=${pt_tag}.match_contest.${ckpt_name}
+asr_speech_fold_length=4800 # 480000/16000 = 30 seconds
 
 ./asr.sh \
     --asr_tag ft.${ft_tag} \
+    --expdir ${expdir} \
+    --dumpdir ${dumpdir} \
     --feats_normalize uttmvn \
-    --stage 1 \
+    --stage 11 \
+    --asr_speech_fold_length ${asr_speech_fold_length} \
     --stop_stage 13 \
     --asr_stats_dir ${expdir}/asr_stats_finetune \
     --ngpu 2 \
diff --git a/egs2/clotho_v2/asr1/run_inference.sh b/egs2/clotho_v2/asr1/run_inference.sh
new file mode 100755
index 00000000000..85c35577447
--- /dev/null
+++ b/egs2/clotho_v2/asr1/run_inference.sh
@@ -0,0 +1,38 @@
+#!/usr/bin/env bash
+# Runs pre-training
+set -euo pipefail
+
+# # Run inference on all checkpoints
+
+wandb_sweep_args="$@"
+
+name_tag=ft.beats.bartfrozen.20241111.164042.valid.acc.ave
+all_ckpts=$(find ${expdir}/asr_${name_tag} -name "*.pth"  -printf "%f\n" | awk '{print substr($0, 1, length($0) - 4)}')
+
+# for ckpt in valid.acc.best; do
+for ckpt in $all_ckpts; do
+    # We skip all symlinks
+    if [[ $ckpt == "checkpoint" || $ckpt == "latest"  || $ckpt == "valid.acc.best"  || $ckpt == "valid.acc.ave.pth" ]]; then
+        continue
+    fi
+    ./asr.sh \
+        --asr_tag $name_tag \
+        --feats_normalize uttmvn \
+        --stage 12 \
+        --stop_stage 13 \
+        --ngpu 2 \
+        --gpu_inference true \
+        --nj 20 \
+        --inference_nj 1 \
+        --max_wav_duration 30 \
+        --token_type hugging_face \
+        --use_lm false \
+        --hugging_face_model_name_or_path "facebook/bart-base" \
+        --inference_args "--beam_size 10 --ctc_weight 0.0 --hugging_face_decoder True" \
+        --train_set pretrain \
+        --valid_set validation \
+        --test_sets "evaluation" \
+        --inference_asr_model ${ckpt}.pth \
+        --asr_args "${wandb_sweep_args}" \
+        --local_score_opts "${expdir}/asr_${name_tag}/inference_beam_size10_ctc_weight0.0_hugging_face_decoderTrue_asr_model_${ckpt}"
+done
diff --git a/egs2/clotho_v2/asr1/run_pt.sh b/egs2/clotho_v2/asr1/run_pt.sh
index f19e18c1d17..cd819326442 100755
--- a/egs2/clotho_v2/asr1/run_pt.sh
+++ b/egs2/clotho_v2/asr1/run_pt.sh
@@ -3,19 +3,25 @@
 set -euo pipefail
 
 timestamp=$(date "+%Y%m%d.%H%M%S")
+mynametag=pt.test_downsample_impl.test_weighting_impl.scratch_bart.3ladapter.match_contest.${timestamp}
 
-# wandb_init_args="--use_wandb true --wandb_project DCASE_AAC --wandb_model_log_interval 0"
-wandb_init_args=""
+wandb_init_args="--use_wandb true --wandb_entity shikhar --wandb_name ${mynametag} --wandb_project DCASE_AAC --wandb_model_log_interval 0"
 other_args="$@"
 
+expdir=/compute/babel-13-33/sbharad2/expdir
+dumpdir=/compute/babel-13-33/sbharad2/dumpdir
+# local_data_opts=/compute/babel-13-33/sbharad2/expdir
+asr_speech_fold_length=4800 # 480000/16000 = 30 seconds
 
 ./asr.sh \
-    --asr_tag pt.${timestamp} \
-    --asr_speech_fold_length 1600 \
+    --asr_tag ${mynametag} \
+    --expdir ${expdir} \
+    --dumpdir ${dumpdir} \
+    --asr_speech_fold_length ${asr_speech_fold_length} \
     --feats_normalize uttmvn \
-    --stage 1 \
+    --stage 11 \
     --stop_stage 13 \
-    --ngpu 2 \
+    --ngpu 1 \
     --gpu_inference true \
     --nj 10 \
     --inference_nj 1 \
@@ -30,4 +36,4 @@ other_args="$@"
     --asr_config conf/beats_bart_pt.yaml \
     --inference_asr_model valid.acc.best.pth \
     --asr_args "${wandb_init_args} ${other_args}" \
-    --local_score_opts "exp/asr_pt.${timestamp}/inference_beam_size10_ctc_weight0.0_hugging_face_decoderTrue_asr_model_valid.acc.best"
+    --local_score_opts "exp/asr_${mynametag}/inference_beam_size10_ctc_weight0.0_hugging_face_decoderTrue_asr_model_valid.acc.best"
diff --git a/egs2/clotho_v2/asr1/run_upload.sh b/egs2/clotho_v2/asr1/run_upload.sh
new file mode 100755
index 00000000000..1667bf91b4c
--- /dev/null
+++ b/egs2/clotho_v2/asr1/run_upload.sh
@@ -0,0 +1,49 @@
+#!/usr/bin/env bash
+set -euo pipefail
+# Runs model upload to huggingface
+
+# wandb_init_args="--use_wandb true --wandb_project DCASE_AAC --wandb_model_log_interval 0"
+wandb_init_args=""
+other_args="$@"
+
+expdir=/compute/babel-13-33/sbharad2/expdir
+dumpdir=/compute/babel-13-33/sbharad2/dumpdir
+local_data_opts=/compute/babel-13-33/sbharad2/expdir
+
+
+expdir=exp/
+# asr_tag=pt.initfix.bigbatch512.lr2e-4.weighted12layers.20241103.145125
+# hf_repo=shikhar7ssu/dcase23.aac.pt
+
+asr_tag=ft_lr5e-5.initfix.bigbatch512.lr2e-4.weighted12layers.20241103.145125
+hf_repo=espnet/DCASE23.AudioCaptioning.FineTuned
+
+
+./asr.sh \
+    --asr_tag ${asr_tag} \
+    --hf_repo ${hf_repo} \
+    --skip_packing false \
+    --skip_upload_hf false \
+    --lang en \
+    --expdir ${expdir} \
+    --dumpdir ${dumpdir} \
+    --feats_normalize uttmvn \
+    --stage 14 \
+    --stop_stage 15 \
+    --asr_stats_dir ${expdir}/asr_stats_finetune \
+    --ngpu 2 \
+    --gpu_inference true \
+    --nj 8 \
+    --inference_nj 1 \
+    --max_wav_duration 30 \
+    --token_type hugging_face \
+    --use_lm false \
+    --hugging_face_model_name_or_path "facebook/bart-base" \
+    --inference_args "--beam_size 10 --ctc_weight 0.0 --hugging_face_decoder True" \
+    --train_set development_clotho \
+    --valid_set validation \
+    --test_sets "evaluation" \
+    --asr_config conf/beats_bart_ft.yaml \
+    --inference_asr_model valid.acc.ave.pth \
+    --asr_args "${wandb_init_args} ${other_args}" \
+    --local_data_opts "${local_data_opts}"
diff --git a/espnet2/asr/decoder/hugging_face_transformers_decoder.py b/espnet2/asr/decoder/hugging_face_transformers_decoder.py
index 1af31b3679d..f12586d0c64 100644
--- a/espnet2/asr/decoder/hugging_face_transformers_decoder.py
+++ b/espnet2/asr/decoder/hugging_face_transformers_decoder.py
@@ -184,12 +184,15 @@ def forward(
         return x, ys_in_lens
 
     def reload_pretrained_parameters(self):
-        self.decoder.load_state_dict(self.decoder_pretrained_params)
+        # self.decoder.load_state_dict(self.decoder_pretrained_params)
 
-        if self.lm_head_pretrained_params is not None:
-            self.lm_head.load_state_dict(self.lm_head_pretrained_params)
+        # if self.lm_head_pretrained_params is not None:
+        # self.lm_head.load_state_dict(self.lm_head_pretrained_params)
 
-        logging.info("Pretrained Transformers model parameters reloaded!")
+        # logging.info("Pretrained Transformers model parameters reloaded!")
+        logging.info(
+            "Skipping loading of pretrained Transformers model parameters for DCASE!"
+        )
 
     def add_prefix_postfix(self, enc_out, hlens, ys_in_pad, ys_in_lens):
         args = {}
diff --git a/espnet2/asr/encoder/beats_encoder.py b/espnet2/asr/encoder/beats_encoder.py
index c3a2b71da06..6575168edad 100644
--- a/espnet2/asr/encoder/beats_encoder.py
+++ b/espnet2/asr/encoder/beats_encoder.py
@@ -112,6 +112,7 @@ class BeatsEncoder(AbsEncoder):
             if None. Otherwise use upto `max_layer`.
         use_weighted_representation: Use weighted representations
             from max_layer if True. Weights are randomly initialized.
+        downsampling_rate: Downsampling rate for the encoder. Applied if > 1.
         beats_config: `BEATsConfig` object. If provided, we will try
             to override the config in the checkpoint. This can be used
             to change dropouts etc for fine-tuning the model while
@@ -127,6 +128,7 @@ def __init__(
         adapter_layers: int = 0,
         max_layer: int = None,
         use_weighted_representation: bool = False,
+        downsampling_rate: int = 1,  # No downsampling by default
         beats_config: Optional[BEATsConfig] = None,
         specaug_config: Optional[Dict] = None,
     ) -> None:
@@ -181,20 +183,39 @@ def __init__(
         )
         self.dropout_input = nn.Dropout(config.dropout_input)
         assert not config.deep_norm or not config.layer_norm_first
-        # Add adapter layer to config.
+
         self.encoder = TransformerEncoder(config)
         self.layer_norm = LayerNorm(self.embed)
+
         self.use_weighted_representation = use_weighted_representation
         if self.use_weighted_representation:
             if self.max_layer is None:
                 logging.warning(
                     f"max_layer must be provided when using weighted representations."
-                    f"Set to {config.encoder_layers}."
+                    f"Set to {config.encoder_layers-1}."
                 )
-                self.max_layer = config.encoder_layers
+                self.max_layer = config.encoder_layers - 1  # 0 based index
             self.layer_weights = nn.Parameter(
-                torch.ones(self.max_layer) / self.max_layer, requires_grad=True
-            )  # Uniform init
+                torch.ones((self.max_layer + 1, 1)), requires_grad=True
+            )
+            # TODO(shikhar): Remove this after testing the current impl from Shih-Lun's code.
+            # self.register_parameter(
+            #     "encoder_repr_layer_weights",
+            #     nn.Parameter(torch.ones((self.max_layer + 1, 1))),
+            # )
+
+        # Downsampling modules
+        self.encoder_downsample_rate = downsampling_rate
+        self.downsample_conv = None
+        if self.encoder_downsample_rate > 1:
+            self.downsample_conv = nn.Conv1d(
+                in_channels=config.encoder_embed_dim,
+                out_channels=config.encoder_embed_dim,
+                kernel_size=int(
+                    round(self.encoder_downsample_rate * 1.5)
+                ),  # kernel multiplier from Shih-Lun's code
+                stride=self.encoder_downsample_rate,
+            )
 
         self.transformer_adapter = None
         if adapter_layers != 0:
@@ -244,12 +265,12 @@ def forward_padding_mask(
         features: torch.Tensor,
         padding_mask: torch.Tensor,
     ) -> torch.Tensor:
-        """Forward padding mask."""
+        """Forward padding mask. Featuires: BTC, padding_mask: BT."""
         extra = padding_mask.size(1) % features.size(1)
         if extra > 0:
             padding_mask = padding_mask[:, :-extra]
         padding_mask = padding_mask.view(padding_mask.size(0), features.size(1), -1)
-        padding_mask = padding_mask.all(-1)  # remove empty sequences
+        padding_mask = padding_mask.all(-1)  # remove totally empty sequences
         return padding_mask
 
     def preprocess(
@@ -323,6 +344,7 @@ def extract_features(
         features = self.layer_norm(features)
 
         if padding_mask is not None:
+            # features is BTC
             padding_mask = self.forward_padding_mask(features, padding_mask)
 
         if self.post_extract_proj is not None:
@@ -333,16 +355,40 @@ def extract_features(
             features, padding_mask=padding_mask, layer=max_layer
         )
 
+        # TODO(shikhar): Remove this after testing the current impl from Shih-Lun's code.
+        # if self.use_weighted_representation:
+        #     for i, (x_i, _) in enumerate(layer_results):
+        #         if i >= max_layer:
+        #             break
+        #         if i == 0:
+        #             features = self.layer_weights[i] * x_i
+        #         else:
+        #             features += self.layer_weights[i] * x_i
+        #     # T x B x C -> B x T x C
+        #     features = features.transpose(0, 1)
+
         if self.use_weighted_representation:
-            for i, (x_i, _) in enumerate(layer_results):
-                if i >= max_layer:
-                    break
-                if i == 0:
-                    features = self.layer_weights[i] * x_i
-                else:
-                    features += self.layer_weights[i] * x_i
-            # T x B x C -> B x T x C
-            features = features.transpose(0, 1)
+            repr_layer_weights = nn.functional.softmax(self.layer_weights, dim=-2)
+            assert (
+                max_layer > 0
+            ), "max_layer must be non-zero when using weighted representations."
+            features = (
+                torch.stack(
+                    [
+                        layer_result_i.transpose(0, 1)
+                        for layer_result_i, _ in layer_results[: max_layer + 1]
+                    ],
+                    dim=-2,
+                )
+                * repr_layer_weights
+            )
+            features = features.sum(dim=-2)  # BTC
+
+        if self.downsample_conv is not None:
+            features = self.downsample_conv(features.transpose(1, 2)).transpose(
+                1, 2
+            )  # BTC
+            padding_mask = self.forward_padding_mask(features, padding_mask)
 
         if self.transformer_adapter:
             features, _ = self.transformer_adapter(

From 896bc23667e4618f53a51991b65d90842ed4dbe1 Mon Sep 17 00:00:00 2001
From: Shikhar Bharadwaj <sbharad2@andrew.cmu.edu>
Date: Thu, 21 Nov 2024 21:25:16 -0500
Subject: [PATCH 23/47] changes for hyp-param tuning and audiocaps cleaning

---
 egs2/TEMPLATE/asr1/db.sh                           |  5 +++--
 egs2/clotho_v2/asr1/conf/beats_bart_pt.yaml        | 10 +++++-----
 egs2/clotho_v2/asr1/conf/slurm.conf                |  2 +-
 egs2/clotho_v2/asr1/local/data.sh                  |  2 +-
 egs2/clotho_v2/asr1/local/data_prep_audiocaps.py   | 10 +++++++++-
 .../asr1/local/data_prep_clotho_chatgpt_mixup.py   |  4 +++-
 egs2/clotho_v2/asr1/run_pt.sh                      |  9 +++++----
 .../decoder/hugging_face_transformers_decoder.py   | 14 +++++++-------
 8 files changed, 34 insertions(+), 22 deletions(-)

diff --git a/egs2/TEMPLATE/asr1/db.sh b/egs2/TEMPLATE/asr1/db.sh
index f00fcc857bb..936ed154f38 100755
--- a/egs2/TEMPLATE/asr1/db.sh
+++ b/egs2/TEMPLATE/asr1/db.sh
@@ -220,9 +220,10 @@ SPRING_INX=downloads
 VOXCELEB=
 KSPONSPEECH=
 HIFITTS=downloads
+
 CLOTHO_V2=downloads
-AUDIOCAPS=
-CLOTHO_CHATGPT_MIXUP=
+AUDIOCAPS=/compute/babel-13-33/sbharad2/audiocaps
+CLOTHO_CHATGPT_MIXUP=/compute/babel-13-33/sbharad2/chatgpt
 
 # For only CMU TIR environment
 if [[ "$(hostname)" == tir* ]]; then
diff --git a/egs2/clotho_v2/asr1/conf/beats_bart_pt.yaml b/egs2/clotho_v2/asr1/conf/beats_bart_pt.yaml
index 9f1cf61a711..392889146be 100644
--- a/egs2/clotho_v2/asr1/conf/beats_bart_pt.yaml
+++ b/egs2/clotho_v2/asr1/conf/beats_bart_pt.yaml
@@ -1,15 +1,15 @@
 batch_type: folded
 unused_parameters: true
-batch_size: 32
-accum_grad: 2
-max_epoch: 20
+batch_size: 64 # 32
+accum_grad: 8 # 2
+max_epoch: 80
 grad_clip: 2
 patience: none
 best_model_criterion:  # criterion to save best models
 -   - valid
     - acc
     - max
-keep_nbest_models: 10  # save nbest models and average these checkpoints
+keep_nbest_models: 5  # save nbest models and average these checkpoints
 use_amp: true    # whether to use automatic mixed precision
 num_att_plot: 0  # do not save attention plots to save time in the demo
 num_workers: 2   # number of workers in dataloader
@@ -77,4 +77,4 @@ optim_conf:
 
 scheduler: warmuplr
 scheduler_conf:
-    warmup_steps: 2500 # 8000
+    warmup_steps: 8000 # 2500
diff --git a/egs2/clotho_v2/asr1/conf/slurm.conf b/egs2/clotho_v2/asr1/conf/slurm.conf
index 3ae0a7ce304..e4ce8f0ad12 100644
--- a/egs2/clotho_v2/asr1/conf/slurm.conf
+++ b/egs2/clotho_v2/asr1/conf/slurm.conf
@@ -13,4 +13,4 @@ option num_threads=1 --cpus-per-task 1
 option num_nodes=* --nodes $0
 default gpu=0
 option gpu=0 -p array
-option gpu=* -p general --gres=gpu:$0 -c $0
+option gpu=* -p debug --gres=gpu:$0 -c $0
diff --git a/egs2/clotho_v2/asr1/local/data.sh b/egs2/clotho_v2/asr1/local/data.sh
index 8b64b6a8453..b0e879bc7d4 100755
--- a/egs2/clotho_v2/asr1/local/data.sh
+++ b/egs2/clotho_v2/asr1/local/data.sh
@@ -25,7 +25,7 @@ fi
 
 log "stage 1: Data preparation"
 
-
+set -ex
 ## DOWNLOAD DATA if CLOTHO_V2 is set to downloads
 if [ "${CLOTHO_V2}" == "downloads" ]; then
     # If there is no argument, the default download directory is set to currentdir/downloads
diff --git a/egs2/clotho_v2/asr1/local/data_prep_audiocaps.py b/egs2/clotho_v2/asr1/local/data_prep_audiocaps.py
index 0d2dc2c5638..70cd0660344 100644
--- a/egs2/clotho_v2/asr1/local/data_prep_audiocaps.py
+++ b/egs2/clotho_v2/asr1/local/data_prep_audiocaps.py
@@ -28,6 +28,10 @@
     missing_audio = []
     N_PROCESSED = 0
     N_ZERO_LENGTH_CAPTIONS = 0
+    # We also skip examples with less than 6 words
+    # according to Shih-Lun's paper Section 3.1:
+    # https://ieeexplore.ieee.org/stamp/stamp.jsp?arnumber=10447215
+    N_SMALL_LENGTH_CAPTIONS = 0
     N_ZERO_LENGTH_AUDIO = 0
     N_VERY_LONG_AUDIO = 0
     AUDIO_LENGTH_THRESHOLD = 30  # seconds
@@ -53,6 +57,9 @@
                 if len(text) == 0:
                     N_ZERO_LENGTH_CAPTIONS += 1
                     continue
+                if len(text.split()) < 6:
+                    N_SMALL_LENGTH_CAPTIONS += 1
+                    continue
                 print(f"{uttid} {audio_path}", file=wav_scp_f)
                 print(f"{uttid} {text}", file=text_f)
                 print(f"{uttid} dummy", file=utt2spk_f)
@@ -62,5 +69,6 @@
     print(
         f"Processed {N_PROCESSED} audio files. Skipped {N_VERY_LONG_AUDIO}"
         f"long audio, {N_ZERO_LENGTH_AUDIO} empty audio and "
-        f"{N_ZERO_LENGTH_CAPTIONS} empty caption items."
+        f"{N_ZERO_LENGTH_CAPTIONS} empty caption items. {N_SMALL_LENGTH_CAPTIONS} "
+        "were skipped because they had less than 6 words."
     )
diff --git a/egs2/clotho_v2/asr1/local/data_prep_clotho_chatgpt_mixup.py b/egs2/clotho_v2/asr1/local/data_prep_clotho_chatgpt_mixup.py
index ccb353200a6..ff34c5b6bfb 100644
--- a/egs2/clotho_v2/asr1/local/data_prep_clotho_chatgpt_mixup.py
+++ b/egs2/clotho_v2/asr1/local/data_prep_clotho_chatgpt_mixup.py
@@ -352,9 +352,11 @@ def _get_ith_item(i):
 ) as wav_scp_f, open(output_utt2spk_path, "w") as utt2spk_f:
     for i in range(len(dset)):
         uttid, audio_path, text = results[i]
-        if uttid is None:
+        if uttid is None or len(uttid) == 0:
             N_ERROR += 1
             continue
+        # Remove new lines and carriage returns
+        text = text.replace("\n", " ").replace("\r", " ").strip()
         print(f"{uttid} {audio_path}", file=wav_scp_f)
         print(f"{uttid} {text}", file=text_f)
         print(f"{uttid} dummy", file=utt2spk_f)
diff --git a/egs2/clotho_v2/asr1/run_pt.sh b/egs2/clotho_v2/asr1/run_pt.sh
index cd819326442..c44b227e1ff 100755
--- a/egs2/clotho_v2/asr1/run_pt.sh
+++ b/egs2/clotho_v2/asr1/run_pt.sh
@@ -3,14 +3,14 @@
 set -euo pipefail
 
 timestamp=$(date "+%Y%m%d.%H%M%S")
-mynametag=pt.test_downsample_impl.test_weighting_impl.scratch_bart.3ladapter.match_contest.${timestamp}
+mynametag=pt.cleanacaps.bs512-warmup8k.load_bart.3ladapter.${timestamp}
 
 wandb_init_args="--use_wandb true --wandb_entity shikhar --wandb_name ${mynametag} --wandb_project DCASE_AAC --wandb_model_log_interval 0"
 other_args="$@"
 
 expdir=/compute/babel-13-33/sbharad2/expdir
-dumpdir=/compute/babel-13-33/sbharad2/dumpdir
-# local_data_opts=/compute/babel-13-33/sbharad2/expdir
+dumpdir=/compute/babel-13-33/sbharad2/dumpdir_clean
+local_data_opts=/compute/babel-13-33/sbharad2/expdir
 asr_speech_fold_length=4800 # 480000/16000 = 30 seconds
 
 ./asr.sh \
@@ -18,10 +18,11 @@ asr_speech_fold_length=4800 # 480000/16000 = 30 seconds
     --expdir ${expdir} \
     --dumpdir ${dumpdir} \
     --asr_speech_fold_length ${asr_speech_fold_length} \
+    --local_data_opts ${local_data_opts} \
     --feats_normalize uttmvn \
     --stage 11 \
     --stop_stage 13 \
-    --ngpu 1 \
+    --ngpu 2 \
     --gpu_inference true \
     --nj 10 \
     --inference_nj 1 \
diff --git a/espnet2/asr/decoder/hugging_face_transformers_decoder.py b/espnet2/asr/decoder/hugging_face_transformers_decoder.py
index f12586d0c64..d6d83c58b44 100644
--- a/espnet2/asr/decoder/hugging_face_transformers_decoder.py
+++ b/espnet2/asr/decoder/hugging_face_transformers_decoder.py
@@ -184,15 +184,15 @@ def forward(
         return x, ys_in_lens
 
     def reload_pretrained_parameters(self):
-        # self.decoder.load_state_dict(self.decoder_pretrained_params)
+        self.decoder.load_state_dict(self.decoder_pretrained_params)
 
-        # if self.lm_head_pretrained_params is not None:
-        # self.lm_head.load_state_dict(self.lm_head_pretrained_params)
+        if self.lm_head_pretrained_params is not None:
+            self.lm_head.load_state_dict(self.lm_head_pretrained_params)
 
-        # logging.info("Pretrained Transformers model parameters reloaded!")
-        logging.info(
-            "Skipping loading of pretrained Transformers model parameters for DCASE!"
-        )
+        logging.info("Pretrained Transformers model parameters reloaded!")
+        # logging.info(
+        #     "Skipping loading of pretrained Transformers model parameters for DCASE!"
+        # )
 
     def add_prefix_postfix(self, enc_out, hlens, ys_in_pad, ys_in_lens):
         args = {}

From 6eea078b089b716d64f68abc578adf87970e37c6 Mon Sep 17 00:00:00 2001
From: Shikhar Bharadwaj <sbharad2@andrew.cmu.edu>
Date: Fri, 22 Nov 2024 23:35:29 -0500
Subject: [PATCH 24/47] ensure that scoring is run on text file and not on
 text.result

---
 egs2/clotho_v2/asr1/local/score.sh | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/egs2/clotho_v2/asr1/local/score.sh b/egs2/clotho_v2/asr1/local/score.sh
index cea04f020c5..e17384428b2 100755
--- a/egs2/clotho_v2/asr1/local/score.sh
+++ b/egs2/clotho_v2/asr1/local/score.sh
@@ -17,7 +17,7 @@ asr_expdir=$1
 splits=(evaluation)
 
 for split in ${splits[@]}; do
-	for decode_file in $((ls -d ${asr_expdir}/*/*/* && ls -d ${asr_expdir}/*/*) | grep "${split}/text"); do
+	for decode_file in $((ls -d ${asr_expdir}/*/*/* && ls -d ${asr_expdir}/*/*) | grep "${split}/text$"); do
 		echo "Decode file: ${decode_file} for split:${split}"
 		python local/evaluation.py --decode_file=${decode_file} --split=${split}
 	done

From b955bf90afe6eeb9f57e31dee1024eff3da38417 Mon Sep 17 00:00:00 2001
From: Shikhar Bharadwaj <sbharad2@andrew.cmu.edu>
Date: Fri, 29 Nov 2024 09:20:07 -0500
Subject: [PATCH 25/47] add model architecture config

---
 .../asr1/conf/bart_decoder_config.json        |  16 +++
 .../asr1/conf/wav2vec2_conformer_config.json  | 113 ++++++++++++++++++
 .../hugging_face_transformers_decoder.py      |  88 +++++++++++---
 espnet2/asr/encoder/beats_encoder.py          |  84 ++++++++-----
 espnet2/bin/asr_inference.py                  |  16 ++-
 5 files changed, 265 insertions(+), 52 deletions(-)
 create mode 100644 egs2/clotho_v2/asr1/conf/bart_decoder_config.json
 create mode 100644 egs2/clotho_v2/asr1/conf/wav2vec2_conformer_config.json

diff --git a/egs2/clotho_v2/asr1/conf/bart_decoder_config.json b/egs2/clotho_v2/asr1/conf/bart_decoder_config.json
new file mode 100644
index 00000000000..6b940da0f2c
--- /dev/null
+++ b/egs2/clotho_v2/asr1/conf/bart_decoder_config.json
@@ -0,0 +1,16 @@
+{
+    "vocab_size": 50265,
+    "max_position_embeddings": 256,
+    "decoder_layers": 6,
+    "decoder_ffn_dim": 3072,
+    "decoder_attention_heads": 16,
+    "d_model": 768,
+    "decoder_layerdrop": 0.0,
+    "activation_function": "gelu",
+    "dropout": 0.1,
+    "attention_dropout": 0.1,
+    "activation_dropout": 0.1,
+    "scale_embedding": false,
+    "use_cache": true,
+    "ignore_mismatched_sizes": true
+}
\ No newline at end of file
diff --git a/egs2/clotho_v2/asr1/conf/wav2vec2_conformer_config.json b/egs2/clotho_v2/asr1/conf/wav2vec2_conformer_config.json
new file mode 100644
index 00000000000..9e175c60a21
--- /dev/null
+++ b/egs2/clotho_v2/asr1/conf/wav2vec2_conformer_config.json
@@ -0,0 +1,113 @@
+{
+    "activation_dropout": 0.1,
+    "adapter_kernel_size": 3,
+    "adapter_stride": 2,
+    "add_adapter": false,
+    "apply_spec_augment": true,
+    "architectures": [
+      "Wav2Vec2ConformerForPreTraining"
+    ],
+    "attention_dropout": 0.1,
+    "bos_token_id": 1,
+    "classifier_proj_size": 256,
+    "codevector_dim": 768,
+    "conformer_conv_dropout": 0.1,
+    "contrastive_logits_temperature": 0.1,
+    "conv_bias": true,
+    "conv_depthwise_kernel_size": 31,
+    "conv_dim": [
+      512,
+      512,
+      512,
+      512,
+      512,
+      512,
+      512
+    ],
+    "conv_kernel": [
+      10,
+      3,
+      3,
+      3,
+      3,
+      2,
+      2
+    ],
+    "conv_stride": [
+      5,
+      2,
+      2,
+      2,
+      2,
+      2,
+      2
+    ],
+    "ctc_loss_reduction": "sum",
+    "ctc_zero_infinity": false,
+    "diversity_loss_weight": 0.1,
+    "do_stable_layer_norm": true,
+    "eos_token_id": 2,
+    "feat_extract_activation": "gelu",
+    "feat_extract_dropout": 0.0,
+    "feat_extract_norm": "layer",
+    "feat_proj_dropout": 0.1,
+    "feat_quantizer_dropout": 0.0,
+    "final_dropout": 0.1,
+    "gradient_checkpointing": false,
+    "hidden_act": "swish",
+    "hidden_dropout": 0.1,
+    "hidden_dropout_prob": 0.1,
+    "hidden_size": 768,
+    "initializer_range": 0.02,
+    "intermediate_size": 1536,
+    "layer_norm_eps": 1e-05,
+    "layerdrop": 0.0,
+    "mask_feature_length": 10,
+    "mask_feature_min_masks": 0,
+    "mask_feature_prob": 0.0,
+    "mask_time_length": 10,
+    "mask_time_min_masks": 2,
+    "mask_time_prob": 0.05,
+    "max_source_positions": 1536,
+    "model_type": "wav2vec2-conformer",
+    "num_adapter_layers": 3,
+    "num_attention_heads": 12,
+    "num_codevector_groups": 2,
+    "num_codevectors_per_group": 320,
+    "num_conv_pos_embedding_groups": 16,
+    "num_conv_pos_embeddings": 128,
+    "num_feat_extract_layers": 7,
+    "num_hidden_layers": 2,
+    "num_negatives": 100,
+    "output_hidden_size": 1024,
+    "pad_token_id": 0,
+    "position_embeddings_type": "relative",
+    "proj_codevector_dim": 768,
+    "rotary_embedding_base": 10000,
+    "tdnn_dilation": [
+      1,
+      2,
+      3,
+      1,
+      1
+    ],
+    "tdnn_dim": [
+      512,
+      512,
+      512,
+      512,
+      1500
+    ],
+    "tdnn_kernel": [
+      5,
+      3,
+      3,
+      1,
+      1
+    ],
+    "tokenizer_class": "Wav2Vec2CTCTokenizer",
+    "torch_dtype": "float32",
+    "transformers_version": "4.19.0.dev0",
+    "use_weighted_layer_sum": false,
+    "xvector_output_dim": 512
+  }
\ No newline at end of file
diff --git a/espnet2/asr/decoder/hugging_face_transformers_decoder.py b/espnet2/asr/decoder/hugging_face_transformers_decoder.py
index d6d83c58b44..6735dbc9f22 100644
--- a/espnet2/asr/decoder/hugging_face_transformers_decoder.py
+++ b/espnet2/asr/decoder/hugging_face_transformers_decoder.py
@@ -6,13 +6,14 @@
 
 import copy
 import logging
-from typing import Any, List, Tuple
-
+import os
+from typing import Any, List, Tuple, Optional
 import torch
 import torch.nn.functional as F
 from typeguard import typechecked
 
 from espnet2.asr.decoder.abs_decoder import AbsDecoder
+from espnet.asr.asr_utils import get_model_conf
 from espnet.nets.pytorch_backend.nets_utils import make_pad_mask
 
 try:
@@ -43,7 +44,34 @@ def __init__(
         causal_lm: bool = False,
         prefix: str = "",
         postfix: str = "",
+        config_path: Optional[str] = None,
+        load_pretrained_weights: bool = True,
+        ensure_untied_lm_head: bool = False,
     ):
+        """
+        Initializes the HuggingFaceTransformersDecoder.
+
+        Args:
+            vocab_size (int): The size of the vocabulary.
+            encoder_output_size (int): The size of the encoder output.
+            model_name_or_path (str): The name or path of the pre-trained Transformers model.
+            causal_lm (bool, optional): Whether to use a causal language model. Defaults to False.
+                This overrides the model_name_or_path if provided.
+            prefix (str, optional): Prefix to be added to the input tokens. Defaults to "".
+            postfix (str, optional): Postfix to be added to the input tokens. Defaults to "".
+            config_path (str, optional): Path to the configuration json file. Defaults to None. If
+                this is set, it can be used to override the default decoder configuration. More documentation here:
+                https://huggingface.co/docs/transformers/en/model_doc/auto#transformers.AutoModelForSeq2SeqLM.from_pretrained.kwargs
+            load_pretrained_weights (bool): Whether to load the pre-trained weights. Defaults to True.
+            ensure_untied_lm_head (bool): True ensures that the language model head is not
+                shared with the input token embeddings. When False, the original structure is
+                kept, ie, if the original Transformers implementation has tying of weights,
+                it is retained. Defaults to False.
+
+        Raises:
+            ImportError: If the `transformers` library is not available.
+            Exception: If the word embeddings attribute cannot be found in the model.
+        """
         super().__init__()
 
         if not is_transformers_available:
@@ -53,10 +81,19 @@ def __init__(
                 " && ./installers/install_transformers.sh`."
             )
 
+        self.load_pretrained_weights = load_pretrained_weights
+        self.ensure_untied_lm_head = ensure_untied_lm_head
+        self.overriding_architecture_config = (
+            vars(get_model_conf(model_path="", conf_path=config_path))
+            if config_path
+            else {}
+        )
         self.causal_lm = causal_lm
 
         if self.causal_lm:
-            model = AutoModelForCausalLM.from_pretrained(model_name_or_path)
+            model = AutoModelForCausalLM.from_pretrained(
+                model_name_or_path, **self.overriding_architecture_config
+            )
             self.decoder = get_hugging_face_model_network(model)
 
             if hasattr(self.decoder, "word_embeddings"):
@@ -87,7 +124,9 @@ def __init__(
                 tokenizer.encode(postfix, return_tensors="pt").long()
             ).detach()
         else:
-            model = AutoModelForSeq2SeqLM.from_pretrained(model_name_or_path)
+            model = AutoModelForSeq2SeqLM.from_pretrained(
+                model_name_or_path, **self.overriding_architecture_config
+            )
 
             if hasattr(model, "model"):
                 self.decoder = model.model.decoder
@@ -96,7 +135,10 @@ def __init__(
 
         model.resize_token_embeddings(vocab_size)
 
-        self.lm_head = get_hugging_face_model_lm_head(model)
+        if self.ensure_untied_lm_head:
+            self.lm_head = copy.deepcopy(get_hugging_face_model_lm_head(model))
+        else:
+            self.lm_head = get_hugging_face_model_lm_head(model)
 
         self.model_name_or_path = model_name_or_path
 
@@ -184,15 +226,17 @@ def forward(
         return x, ys_in_lens
 
     def reload_pretrained_parameters(self):
-        self.decoder.load_state_dict(self.decoder_pretrained_params)
-
-        if self.lm_head_pretrained_params is not None:
-            self.lm_head.load_state_dict(self.lm_head_pretrained_params)
+        if self.load_pretrained_weights:
+            self.decoder.load_state_dict(self.decoder_pretrained_params)
+            logging.info("Loaded pretrained Transformers decoder parameters!")
 
-        logging.info("Pretrained Transformers model parameters reloaded!")
-        # logging.info(
-        #     "Skipping loading of pretrained Transformers model parameters for DCASE!"
-        # )
+            if self.lm_head_pretrained_params is not None:
+                self.lm_head.load_state_dict(self.lm_head_pretrained_params)
+                logging.info("Loaded pretrained Transformers LM head parameters!")
+        else:
+            logging.info(
+                "Skipping the loading of pretrained Transformer model parameters!"
+            )
 
     def add_prefix_postfix(self, enc_out, hlens, ys_in_pad, ys_in_lens):
         args = {}
@@ -252,7 +296,7 @@ def score(self, ys, state, x, speech=None):
             **model_inputs,
             return_dict=True,
             output_attentions=False,
-            output_hidden_states=False
+            output_hidden_states=False,
         )
         next_token_logits = outputs.logits[:, -1, :]
         next_token_scores = torch.nn.functional.log_softmax(
@@ -278,7 +322,7 @@ def batch_score(
             **model_inputs,
             return_dict=True,
             output_attentions=False,
-            output_hidden_states=False
+            output_hidden_states=False,
         )
         next_token_logits = outputs.logits[:, -1, :]
         next_token_scores = torch.nn.functional.log_softmax(
@@ -309,3 +353,17 @@ def get_hugging_face_model_lm_head(model):
         raise Exception("Can not find the LM head attribute")
 
     return lm_head
+
+
+def read_config_file(config_path):
+    assert os.path.exists(
+        config_path
+    ), f"Config file supplied but not found: {config_path}"
+    config_dict = {}
+    with open(config_path, "r") as f:
+        config_dict = json.load(f)
+        config_dict["ignore_mismatched_sizes"] = (
+            True  # Loaded config could be different from pre-trained model.
+        )
+
+    return config_dict
diff --git a/espnet2/asr/encoder/beats_encoder.py b/espnet2/asr/encoder/beats_encoder.py
index 6575168edad..157ce678c11 100644
--- a/espnet2/asr/encoder/beats_encoder.py
+++ b/espnet2/asr/encoder/beats_encoder.py
@@ -28,6 +28,12 @@
 from espnet2.asr.specaug.specaug import SpecAug
 from espnet.nets.pytorch_backend.nets_utils import make_pad_mask
 
+from transformers.models.wav2vec2_conformer.modeling_wav2vec2_conformer import (
+    Wav2Vec2ConformerConfig,
+    Wav2Vec2ConformerEncoder,
+)
+from transformers.models.bart.modeling_bart import BartLearnedPositionalEmbedding
+
 if V(torch.__version__) >= V("1.6.0"):
     from torch.cuda.amp import autocast
 else:
@@ -106,31 +112,35 @@ class BeatsEncoder(AbsEncoder):
         beats_ckpt_path: Path to a pretrained BEATs checkpoint. If
             `beats_config` is provided and it does not match the
             config in the checkpoint, code might throw an error.
-        adapter_layers: Number of transformer-based adapter layers
-            after encoder.
         max_layer: Propagate input through all layers for encoding
             if None. Otherwise use upto `max_layer`.
+        downsampling_rate: Downsampling rate for the encoder. Applied if > 1.
+        adapter_config: Path to a config file for the wav2vec2 adapter.
         use_weighted_representation: Use weighted representations
             from max_layer if True. Weights are randomly initialized.
-        downsampling_rate: Downsampling rate for the encoder. Applied if > 1.
         beats_config: `BEATsConfig` object. If provided, we will try
             to override the config in the checkpoint. This can be used
             to change dropouts etc for fine-tuning the model while
             starting from a pretrained checkpoint.
         specaug_config: Dictionary containing parameters for SpecAugment.
             If provided, SpecAugment will be applied.
+        add_positional_information: Add learned positional embeddings.
+        max_positions: Maximum number of positions for positional embeddings.
+            Required if `add_positional_information` is True.
     """
 
     def __init__(
         self,
         input_size: int,
         beats_ckpt_path: str = None,
-        adapter_layers: int = 0,
         max_layer: int = None,
+        downsampling_rate: int = 1,
+        adapter_config: str = "",
         use_weighted_representation: bool = False,
-        downsampling_rate: int = 1,  # No downsampling by default
         beats_config: Optional[BEATsConfig] = None,
         specaug_config: Optional[Dict] = None,
+        add_positional_information: bool = False,
+        max_positions: Optional[int] = None,
     ) -> None:
         super().__init__()
 
@@ -191,18 +201,13 @@ def __init__(
         if self.use_weighted_representation:
             if self.max_layer is None:
                 logging.warning(
-                    f"max_layer must be provided when using weighted representations."
-                    f"Set to {config.encoder_layers-1}."
+                    f"max_layer must be provided when using weighted"
+                    f" representations. Set to {config.encoder_layers-1}."
                 )
                 self.max_layer = config.encoder_layers - 1  # 0 based index
             self.layer_weights = nn.Parameter(
                 torch.ones((self.max_layer + 1, 1)), requires_grad=True
             )
-            # TODO(shikhar): Remove this after testing the current impl from Shih-Lun's code.
-            # self.register_parameter(
-            #     "encoder_repr_layer_weights",
-            #     nn.Parameter(torch.ones((self.max_layer + 1, 1))),
-            # )
 
         # Downsampling modules
         self.encoder_downsample_rate = downsampling_rate
@@ -217,11 +222,26 @@ def __init__(
                 stride=self.encoder_downsample_rate,
             )
 
-        self.transformer_adapter = None
-        if adapter_layers != 0:
-            adapter_config = deepcopy(config)
-            adapter_config.encoder_layers = adapter_layers
-            self.transformer_adapter = TransformerEncoder(adapter_config)
+        # Adapter module
+        self.conformer_adapter = None
+        if adapter_config:
+            conformer_config = Wav2Vec2ConformerConfig.from_json_file(adapter_config)
+            self.conformer_adapter = Wav2Vec2ConformerEncoder(conformer_config)
+
+        # Positional embeddings applied before cross-attention with decoder.
+        self.cross_embed_positions = None
+        if add_positional_information:
+            assert (
+                max_positions is not None
+            ), "max_positions must be provided in the config."
+            learned_pos_dim = (
+                config.encoder_embed_dim
+                if not self.conformer_adapter
+                else self.conformer_adapter.config.hidden_size
+            )
+            self.cross_embed_positions = BartLearnedPositionalEmbedding(
+                max_positions, learned_pos_dim
+            )
 
     def reload_pretrained_parameters(self):
         """Initialization function for BEATs.
@@ -355,17 +375,10 @@ def extract_features(
             features, padding_mask=padding_mask, layer=max_layer
         )
 
-        # TODO(shikhar): Remove this after testing the current impl from Shih-Lun's code.
-        # if self.use_weighted_representation:
-        #     for i, (x_i, _) in enumerate(layer_results):
-        #         if i >= max_layer:
-        #             break
-        #         if i == 0:
-        #             features = self.layer_weights[i] * x_i
-        #         else:
-        #             features += self.layer_weights[i] * x_i
-        #     # T x B x C -> B x T x C
-        #     features = features.transpose(0, 1)
+        if max_layer is not None:
+            features = layer_results[max_layer][0].transpose(
+                0, 1
+            )  # use the output from the max_layer
 
         if self.use_weighted_representation:
             repr_layer_weights = nn.functional.softmax(self.layer_weights, dim=-2)
@@ -390,10 +403,17 @@ def extract_features(
             )  # BTC
             padding_mask = self.forward_padding_mask(features, padding_mask)
 
-        if self.transformer_adapter:
-            features, _ = self.transformer_adapter(
-                features, padding_mask=padding_mask, layer=None
-            )
+        if self.conformer_adapter:
+            # to handle incompatibility btw torch & huggingface
+            conformer_attn_mask = ~padding_mask
+            # run through conformer
+            features = self.conformer_adapter(
+                features,
+                attention_mask=conformer_attn_mask,
+            ).last_hidden_state
+
+        if self.cross_embed_positions is not None:
+            features = features + self.cross_embed_positions(features)
 
         return features, padding_mask
 
diff --git a/espnet2/bin/asr_inference.py b/espnet2/bin/asr_inference.py
index c30155d7fe5..c9d08ba67aa 100644
--- a/espnet2/bin/asr_inference.py
+++ b/espnet2/bin/asr_inference.py
@@ -49,6 +49,7 @@
 try:
     from transformers import AutoModelForCausalLM, AutoModelForSeq2SeqLM
     from transformers.file_utils import ModelOutput
+    from transformers.models.bart.modeling_bart import BartConfig, BartDecoder
 
     is_transformers_available = True
 except ImportError:
@@ -244,21 +245,27 @@ def __init__(
 
             if decoder.causal_lm:
                 hugging_face_model = AutoModelForCausalLM.from_pretrained(
-                    decoder.model_name_or_path
+                    decoder.model_name_or_path, **decoder.overriding_architecture_config
                 )
 
                 hugging_face_model.resize_token_embeddings(decoder.lm_head.out_features)
 
                 transformer = get_hugging_face_model_network(hugging_face_model)
                 transformer.load_state_dict(decoder.decoder.state_dict())
-
-                lm_head = get_hugging_face_model_lm_head(hugging_face_model)
+                if decoder.ensure_untied_lm_head:
+                    lm_head = copy.deepcopy(
+                        get_hugging_face_model_lm_head(hugging_face_model)
+                    )
+                else:
+                    lm_head = get_hugging_face_model_lm_head(hugging_face_model)
                 lm_head.load_state_dict(decoder.lm_head.state_dict())
             else:
                 hugging_face_model = AutoModelForSeq2SeqLM.from_pretrained(
-                    decoder.model_name_or_path
+                    decoder.model_name_or_path, **decoder.overriding_architecture_config
                 )
 
+                if decoder.ensure_untied_lm_head:
+                    hugging_face_model.lm_head = copy.deepcopy(decoder.lm_head)
                 hugging_face_model.lm_head.load_state_dict(decoder.lm_head.state_dict())
 
                 if hasattr(hugging_face_model, "model"):
@@ -290,7 +297,6 @@ def __init__(
                 hugging_face_decoder_conf["pad_token_id"] = (
                     hugging_face_model.config.eos_token_id
                 )
-
             beam_search = None
             beam_search_transducer = None
         else:

From b2ed4b9338b29b830abcdd4aa1cda28d250eed26 Mon Sep 17 00:00:00 2001
From: Shikhar Bharadwaj <sbharad2@andrew.cmu.edu>
Date: Fri, 29 Nov 2024 09:20:35 -0500
Subject: [PATCH 26/47] add normal init

---
 espnet2/tasks/asr.py              | 1 +
 espnet2/torch_utils/initialize.py | 2 ++
 2 files changed, 3 insertions(+)

diff --git a/espnet2/tasks/asr.py b/espnet2/tasks/asr.py
index a9b9bd1eb7d..a44c4785b23 100644
--- a/espnet2/tasks/asr.py
+++ b/espnet2/tasks/asr.py
@@ -260,6 +260,7 @@ def add_task_arguments(cls, parser: argparse.ArgumentParser):
                 "xavier_normal",
                 "kaiming_uniform",
                 "kaiming_normal",
+                "normal",
                 None,
             ],
         )
diff --git a/espnet2/torch_utils/initialize.py b/espnet2/torch_utils/initialize.py
index 938fce549f8..6371f883a40 100644
--- a/espnet2/torch_utils/initialize.py
+++ b/espnet2/torch_utils/initialize.py
@@ -79,6 +79,8 @@ def initialize(model: torch.nn.Module, init: str):
                     torch.nn.init.kaiming_uniform_(p.data, nonlinearity="relu")
                 elif init == "kaiming_normal":
                     torch.nn.init.kaiming_normal_(p.data, nonlinearity="relu")
+                elif init == "normal":
+                    torch.nn.init.normal_(p.data, mean=0.0, std=0.02)
                 else:
                     raise ValueError("Unknown initialization: " + init)
         # bias init

From aa3da9112b6df18331aae226ebba7dc2836d351d Mon Sep 17 00:00:00 2001
From: Shikhar Bharadwaj <sbharad2@andrew.cmu.edu>
Date: Fri, 29 Nov 2024 09:23:46 -0500
Subject: [PATCH 27/47] use all 5 clotho captions, add text preprocessing

---
 egs2/clotho_v2/asr1/local/data.sh             |  3 +-
 .../asr1/local/data_prep_audiocaps.py         |  6 ++-
 egs2/clotho_v2/asr1/local/data_prep_clotho.py | 49 +++++++++++++++----
 3 files changed, 45 insertions(+), 13 deletions(-)

diff --git a/egs2/clotho_v2/asr1/local/data.sh b/egs2/clotho_v2/asr1/local/data.sh
index b0e879bc7d4..3a2f57da2dc 100755
--- a/egs2/clotho_v2/asr1/local/data.sh
+++ b/egs2/clotho_v2/asr1/local/data.sh
@@ -25,7 +25,6 @@ fi
 
 log "stage 1: Data preparation"
 
-set -ex
 ## DOWNLOAD DATA if CLOTHO_V2 is set to downloads
 if [ "${CLOTHO_V2}" == "downloads" ]; then
     # If there is no argument, the default download directory is set to currentdir/downloads
@@ -89,7 +88,7 @@ if [ ! -d ${CLOTHO_CHATGPT_MIXUP} ]; then
     exit 1
 fi
 
-SPLITS=(development_clotho validation evaluation) # map dev, val, eval to development, validation, evaluation
+SPLITS=(development_clotho_all validation evaluation) # map dev, val, eval to development, validation, evaluation
 N_REF=5
 for split_name in "${SPLITS[@]}"; do
     mkdir -p "data/${split_name}"
diff --git a/egs2/clotho_v2/asr1/local/data_prep_audiocaps.py b/egs2/clotho_v2/asr1/local/data_prep_audiocaps.py
index 70cd0660344..e0ec767d5d5 100644
--- a/egs2/clotho_v2/asr1/local/data_prep_audiocaps.py
+++ b/egs2/clotho_v2/asr1/local/data_prep_audiocaps.py
@@ -3,9 +3,11 @@
 import json
 import os
 import sys
-
+from string import punctuation
 from scipy.io import wavfile
 
+strip_punct_table = str.maketrans("", "", punctuation)
+
 if __name__ == "__main__":
     ROOT_DATA_DIR = sys.argv[1]
     # this is a json with key "audiocap_ids" containing list of folder_ids
@@ -53,7 +55,7 @@
                 os.path.join(ROOT_DATA_DIR, DATA_DIR, train_id, "metadata.json")
             ) as f:
                 meta = json.load(f)
-                text = meta["caption"].strip()
+                text = meta["caption"].strip().lower().translate(strip_punct_table)
                 if len(text) == 0:
                     N_ZERO_LENGTH_CAPTIONS += 1
                     continue
diff --git a/egs2/clotho_v2/asr1/local/data_prep_clotho.py b/egs2/clotho_v2/asr1/local/data_prep_clotho.py
index 6cddcf58a87..91ad9173577 100644
--- a/egs2/clotho_v2/asr1/local/data_prep_clotho.py
+++ b/egs2/clotho_v2/asr1/local/data_prep_clotho.py
@@ -3,8 +3,11 @@
 import glob
 import os
 import sys
-
+import random
 import pandas as pd
+from string import punctuation
+
+strip_punct_table = str.maketrans("", "", punctuation)
 
 if __name__ == "__main__":
     ROOT_DATA_DIR = sys.argv[1]
@@ -23,8 +26,9 @@
         )
         captions_df = pd.read_csv(all_captions_path)
         N_PROCESSED = 0
+        N_CAPTIONS_PROCESSED = 0
         if data_split == "development":
-            dir_write_name = "development_clotho"
+            dir_write_name = "development_clotho_all"
         else:
             dir_write_name = data_split  # for validation and evaluation
         with open(
@@ -61,23 +65,50 @@
                         # Missing audio file
                         missing_audio += [audio_path]
                         continue
-                    print(f"{uttid} dummy", file=utt2spk_f)
-                    print(f"{uttid} {audio_path}", file=wav_scp_f)
                     if data_split == "evaluation":
+                        print(f"{uttid} dummy", file=utt2spk_f)
+                        print(f"{uttid} {audio_path}", file=wav_scp_f)
                         for i in range(1, N_REF + 1):
-                            text_i = row[f"caption_{i}"].strip()
+                            text_i = (
+                                row[f"caption_{i}"]
+                                .strip()
+                                .lower()
+                                .translate(strip_punct_table)
+                            )
+                            # Put each caption in a separate file
                             print(f"{uttid} {text_i}", file=text_f[i - 1])
+                            N_CAPTIONS_PROCESSED += 1
                     else:
-                        text = row["caption_1"].strip()
-                        print(f"{uttid} {text}", file=text_f[0])
+                        for caption_idx in range(1, N_REF + 1):
+                            text = (
+                                row[f"caption_{caption_idx}"]
+                                .strip()
+                                .lower()
+                                .translate(strip_punct_table)
+                            )
+                            # NOTE(shikhar): These files are sorted. Therefore adding
+                            # cpation index first to maintain the order of audio seen
+                            # during training.
+                            print(f"{caption_idx}_{uttid} dummy", file=utt2spk_f)
+                            print(f"{caption_idx}_{uttid} {audio_path}", file=wav_scp_f)
+                            print(
+                                f"{caption_idx}_{uttid} {text}", file=text_f[0]
+                            )  # put in same file but with different id
+                            N_CAPTIONS_PROCESSED += 1
+                            if data_split == "validation":
+                                # We don't need all captions for validation set.
+                                break
                     N_PROCESSED += 1
                     if N_PROCESSED % 1000 == 0:
-                        print(f"Processed {N_PROCESSED} audio files.")
+                        print(
+                            f"Processed {N_PROCESSED} audio files and"
+                            f" {N_CAPTIONS_PROCESSED} captions."
+                        )
                         continue
             finally:
                 for f in text_f:
                     f.close()
         print(
             f"For split {data_split}: Processed {N_PROCESSED} audio and their"
-            f" captions. {len(missing_audio)} audio files were missing."
+            f" {N_CAPTIONS_PROCESSED} captions. {len(missing_audio)} audio files were missing."
         )

From 9c8a90b9ada4f36f1ac38dc8fa971a1da1c1a7a7 Mon Sep 17 00:00:00 2001
From: Shikhar Bharadwaj <sbharad2@andrew.cmu.edu>
Date: Fri, 29 Nov 2024 09:25:33 -0500
Subject: [PATCH 28/47] add text preprocessing before evaluation

---
 egs2/clotho_v2/asr1/local/evaluation.py | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/egs2/clotho_v2/asr1/local/evaluation.py b/egs2/clotho_v2/asr1/local/evaluation.py
index ef3f3cbd6ea..11c14067760 100644
--- a/egs2/clotho_v2/asr1/local/evaluation.py
+++ b/egs2/clotho_v2/asr1/local/evaluation.py
@@ -4,6 +4,9 @@
 import argparse
 
 from aac_metrics import Evaluate
+from string import punctuation
+
+strip_punct_table = str.maketrans("", "", punctuation)
 
 parser = argparse.ArgumentParser()
 parser.add_argument(
@@ -27,7 +30,9 @@ def _parse_and_read_file(file_path: str):
         with open(file_path, "r", encoding="utf-8") as f:
             for line in f:
                 uttid, prediction = line.strip().split(maxsplit=1)
-                read_data[uttid] = prediction.strip()
+                read_data[uttid] = (
+                    prediction.strip().lower().translate(strip_punct_table)
+                )
     except Exception as exc:
         raise ValueError(f"Error reading file {file_path}") from exc
     if len(read_data) == 0:

From 6a929cf89a3d32eaf4538f839ff0ec86f0c0ba46 Mon Sep 17 00:00:00 2001
From: Shikhar Bharadwaj <sbharad2@andrew.cmu.edu>
Date: Fri, 29 Nov 2024 09:27:37 -0500
Subject: [PATCH 29/47] add pretraining and finetuning configs

---
 egs2/clotho_v2/asr1/conf/beats_bart_ft.yaml | 39 ++++++++++---------
 egs2/clotho_v2/asr1/conf/beats_bart_pt.yaml | 42 ++++++++++++---------
 2 files changed, 47 insertions(+), 34 deletions(-)

diff --git a/egs2/clotho_v2/asr1/conf/beats_bart_ft.yaml b/egs2/clotho_v2/asr1/conf/beats_bart_ft.yaml
index ebd3c60190d..94ba6bc997d 100644
--- a/egs2/clotho_v2/asr1/conf/beats_bart_ft.yaml
+++ b/egs2/clotho_v2/asr1/conf/beats_bart_ft.yaml
@@ -2,18 +2,17 @@ batch_type: folded
 unused_parameters: true
 batch_size: 64
 accum_grad: 2
-max_epoch: 40
-grad_clip: 2
+max_epoch: 8
+grad_clip: 1
 patience: none
-best_model_criterion:  # criterion to save best models
+best_model_criterion:
 -   - valid
     - acc
     - max
-keep_nbest_models: 5  # save nbest models and average these checkpoints
-use_amp: true    # whether to use automatic mixed precision
-num_att_plot: 0  # do not save attention plots to save time in the demo
-num_workers: 2   # number of workers in dataloader - was 128
-
+keep_nbest_models: 5
+use_amp: false  # whether to use automatic mixed precision
+num_att_plot: 0
+num_workers: 2 # dataloader workers
 
 # BEATs implementation takes care of generating mel spectrogram, normalization and specaug
 frontend: none
@@ -29,7 +28,8 @@ freeze_param: [
 
 encoder: beats
 encoder_conf:
-    # Please download the BEATs model from https://github.com/microsoft/unilm/tree/master/beats (iter3+, Beats finetuned model 1) and update the path below
+    # Please download the BEATs model from https://github.com/microsoft/unilm/tree/master/beats 
+    # (iter3+, Beats finetuned model 1) and update the path below
     beats_ckpt_path: /compute/babel-13-33/sbharad2/models/BEATs/BEATs_iter3_plus_AS2M_finetuned_on_AS2M_cpt1.pt
     specaug_config:
         apply_freq_mask: true
@@ -42,17 +42,24 @@ encoder_conf:
         - 0
         - 0.12
         num_time_mask: 5
-    adapter_layers: 2
-    max_layer: 12
-    use_weighted_representation: true
+    adapter_config: conf/wav2vec2_conformer_config.json
+    downsampling_rate: 3 # CNN downsampling over beats encoder
+    max_layer: 10 # 0 based index
+    use_weighted_representation: false
+    add_positional_information: true
+    max_positions: 1024 # These many positional embeddings will be learned
 
+# Pleae note that the decoder is not the standard BART-base, 
+# but a custom one whose config is defined in the file below
 decoder: hugging_face_transformers
 decoder_conf:
     model_name_or_path: facebook/bart-base
+    config_path: conf/bart_decoder_config.json
+    load_pretrained_weights: false
+    ensure_untied_lm_head: true
 
-# The initilization is only for the linear layer between the BEATs encoder and the BART decoder.
-# BEATs and BART weights override this for encoder and decoder respectively.
-init: kaiming_uniform
+# Initialization does not matter we use a pre-trained model
+init: normal
 
 token_type: hugging_face
 
@@ -76,5 +83,3 @@ optim_conf:
 scheduler: warmuplr
 scheduler_conf:
     warmup_steps: 1000 #1k
-
-# https://github.com/slSeanWU/beats-conformer-bart-audio-captioner/blob/main/config/training_config.json
diff --git a/egs2/clotho_v2/asr1/conf/beats_bart_pt.yaml b/egs2/clotho_v2/asr1/conf/beats_bart_pt.yaml
index 392889146be..3d9ffcdb7f8 100644
--- a/egs2/clotho_v2/asr1/conf/beats_bart_pt.yaml
+++ b/egs2/clotho_v2/asr1/conf/beats_bart_pt.yaml
@@ -1,20 +1,19 @@
 batch_type: folded
 unused_parameters: true
-batch_size: 64 # 32
-accum_grad: 8 # 2
-max_epoch: 80
-grad_clip: 2
+batch_size: 64
+accum_grad: 1
+max_epoch: 10
+grad_clip: 1
 patience: none
-best_model_criterion:  # criterion to save best models
+best_model_criterion:
 -   - valid
     - acc
     - max
-keep_nbest_models: 5  # save nbest models and average these checkpoints
-use_amp: true    # whether to use automatic mixed precision
-num_att_plot: 0  # do not save attention plots to save time in the demo
+keep_nbest_models: 5
+use_amp: false  # whether to use automatic mixed precision
+num_att_plot: 0
 num_workers: 2   # number of workers in dataloader
 
-
 # BEATs implementation takes care of generating mel spectrogram, normalization and specaug
 frontend: none
 input_size: 1 # important to set input_size to 1 if frontend is none
@@ -31,7 +30,9 @@ freeze_param: [
 
 encoder: beats
 encoder_conf:
-    beats_ckpt_path: /compute/babel-13-33/sbharad2/models/BEATs/BEATs_iter3_plus_AS2M_finetuned_on_AS2M_cpt1.pt # /u/sbharadwaj/BEATs_models/BEATs_iter3_plus_AS2M_finetuned_on_AS2M_cpt1.pt # _finetuned_on_AS2M_cpt1
+    # Please download the BEATs model from https://github.com/microsoft/unilm/tree/master/beats 
+    # (iter3+, Beats finetuned model 1) and update the path below
+    beats_ckpt_path: /compute/babel-13-33/sbharad2/models/BEATs/BEATs_iter3_plus_AS2M_finetuned_on_AS2M_cpt1.pt
     specaug_config:
         apply_freq_mask: true
         freq_mask_width_range:
@@ -43,18 +44,25 @@ encoder_conf:
         - 0
         - 0.12
         num_time_mask: 5
-    adapter_layers: 3
+    adapter_config: conf/wav2vec2_conformer_config.json
     downsampling_rate: 3 # CNN downsampling over beats encoder
-    max_layer: 11 # 0 based index
-    use_weighted_representation: true
+    max_layer: 10 # 0 based index
+    use_weighted_representation: false
+    add_positional_information: true
+    max_positions: 1024 # These many positional embeddings will be learned
 
+# Pleae note that the decoder is not the standard BART-base, 
+# but a custom one whose config is defined in the file below
 decoder: hugging_face_transformers
 decoder_conf:
     model_name_or_path: facebook/bart-base
+    config_path: conf/bart_decoder_config.json
+    load_pretrained_weights: false
+    ensure_untied_lm_head: true
 
-# The initilization is only for the linear layer between the BEATs encoder and the BART decoder.
-# BEATs and BART weights override this for encoder and decoder respectively.
-init: kaiming_uniform
+# This takes care of BART and adapter layers' initialization 
+# (essentially everything else other than BEATs encoder backbone)
+init: normal
 
 token_type: hugging_face
 
@@ -77,4 +85,4 @@ optim_conf:
 
 scheduler: warmuplr
 scheduler_conf:
-    warmup_steps: 8000 # 2500
+    warmup_steps: 2500 # 2.5k

From 9869ca4f240dfd99cfc7267ee32c61fd497c8830 Mon Sep 17 00:00:00 2001
From: Shikhar Bharadwaj <sbharad2@andrew.cmu.edu>
Date: Fri, 29 Nov 2024 09:57:50 -0500
Subject: [PATCH 30/47] add run file

---
 egs2/clotho_v2/asr1/run.sh | 55 ++++++++++++++++++++++++++++++++++++++
 1 file changed, 55 insertions(+)
 create mode 100755 egs2/clotho_v2/asr1/run.sh

diff --git a/egs2/clotho_v2/asr1/run.sh b/egs2/clotho_v2/asr1/run.sh
new file mode 100755
index 00000000000..28b1925411d
--- /dev/null
+++ b/egs2/clotho_v2/asr1/run.sh
@@ -0,0 +1,55 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+# Runs pre-training
+asr_speech_fold_length=4800 # 480000/16000 = 30 seconds
+
+./asr.sh \
+    --asr_tag pt \
+    --feats_normalize uttmvn \
+    --stage 1 \
+    --stop_stage 15 \
+    --ngpu 2 \
+    --gpu_inference true \
+    --nj 10 \
+    --inference_nj 1 \
+    --max_wav_duration 30 \
+    --token_type hugging_face \
+    --use_lm false \
+    --hugging_face_model_name_or_path "facebook/bart-base" \
+    --inference_args "--ctc_weight 0.0 --hugging_face_decoder True" \
+    --train_set pretrain \
+    --valid_set validation \
+    --test_sets "validation evaluation" \
+    --asr_config conf/beats_bart_pt.yaml \
+    --asr_speech_fold_length ${asr_speech_fold_length} \
+    --inference_asr_model latest.pth \
+    --local_score_opts "exp/asr_pt/ctc_weight0.0_hugging_face_decoderTrue_asr_model_latest"
+
+
+# Runs fine-tuning
+inference_ckpt=valid.acc.ave_5best
+
+./asr.sh \
+    --asr_tag ft \
+    --feats_normalize uttmvn \
+    --stage 1 \
+    --asr_speech_fold_length ${asr_speech_fold_length} \
+    --stop_stage 15 \
+    --asr_stats_dir exp/asr_stats_finetune \
+    --ngpu 2 \
+    --gpu_inference true \
+    --nj 8 \
+    --inference_nj 1 \
+    --max_wav_duration 30 \
+    --token_type hugging_face \
+    --use_lm false \
+    --hugging_face_model_name_or_path "facebook/bart-base" \
+    --inference_args "--ctc_weight 0.0 --hugging_face_decoder True" \
+    --train_set development_clotho_all \
+    --valid_set validation \
+    --test_sets "evaluation" \
+    --asr_config conf/beats_bart_ft.yaml \
+    --pretrained_model exp/asr_pt/latest.pth \
+    --inference_asr_model ${inference_ckpt}.pth \
+    --local_score_opts exp/asr_ft/ctc_weight0.0_hugging_face_decoderTrue_asr_model_${inference_ckpt}
\ No newline at end of file

From ac3c065c80c341f602b0b8655579f8a5619f79a6 Mon Sep 17 00:00:00 2001
From: Shikhar Bharadwaj <sbharad2@andrew.cmu.edu>
Date: Sat, 30 Nov 2024 09:24:34 -0500
Subject: [PATCH 31/47] add missing import

---
 espnet2/bin/asr_inference.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/espnet2/bin/asr_inference.py b/espnet2/bin/asr_inference.py
index c9d08ba67aa..0b0929fead4 100644
--- a/espnet2/bin/asr_inference.py
+++ b/espnet2/bin/asr_inference.py
@@ -1,5 +1,6 @@
 #!/usr/bin/env python3
 import argparse
+import copy
 import logging
 import sys
 from distutils.version import LooseVersion

From ea0ca7f9ed520d1f2e5f065501cab3e3c957f4ca Mon Sep 17 00:00:00 2001
From: Shikhar Bharadwaj <sbharad2@andrew.cmu.edu>
Date: Sat, 30 Nov 2024 09:32:12 -0500
Subject: [PATCH 32/47] restore to default slurm config, add Readme

---
 egs2/clotho_v2/asr1/README.md       | 58 ++++++++++++++---------------
 egs2/clotho_v2/asr1/conf/slurm.conf |  4 +-
 2 files changed, 29 insertions(+), 33 deletions(-)

diff --git a/egs2/clotho_v2/asr1/README.md b/egs2/clotho_v2/asr1/README.md
index 01ba638d981..c5b653e2dfa 100644
--- a/egs2/clotho_v2/asr1/README.md
+++ b/egs2/clotho_v2/asr1/README.md
@@ -1,6 +1,7 @@
 # Clotho Audio Captioning RECIPE
 
-This recipe implements the DCASE 2023 Automated Audio Captioning (AAC) task with a BEATs encoder BART decoder model on the Clotho_v2 dataset, very close to what is described in [this paper](https://dcase.community/documents/challenge2023/technical_reports/DCASE2023_Wu_31_t6a.pdf) and reusing part of code from the [original implementation](https://github.com/slSeanWU/beats-conformer-bart-audio-captioner?tab=readme-ov-file).
+This recipe implements the DCASE 2023 Automated Audio Captioning (AAC) task with a BEATs encoder BART decoder model on the Clotho_v2 dataset, very close to what is described in [this paper](https://arxiv.org/abs/2309.17352) and reusing part of code from the [original implementation](https://github.com/slSeanWU/beats-conformer-bart-audio-captioner?tab=readme-ov-file).
+More specifically, we provide the pre-training and fine-tuning code for second last row in Table 2 (without Instructor embedding) from the paper. 
 We also reuse the code from the [BEATs repository](https://github.com/microsoft/unilm/tree/master/beats) for this implementation.
 
 # Training Details and Requirements
@@ -10,59 +11,54 @@ Fine-tuning is performed only on the standard development set of Clotho_v2 datas
 
 
 ### Steps to run
-For running AAC pre-training
+
 1. Download AudioCaps and set the path to its root directory in db.sh (This recipe downloads clotho for you). Also download clotho mixup data (wav ids and captions) from [this repo](https://huggingface.co/datasets/slseanwu/clotho-chatgpt-mixup-50K) and set `CLOTHO_CHATGPT_MIXUP` in db.sh. The code will take care of mixing and creating the audio files.
 2. Download the BEATs checkpoint: [BEATs_iter3+](https://onedrive.live.com/?authkey=%21AGXnEG4l3mlIzfA&id=6B83B49411CA81A7%2125960&cid=6B83B49411CA81A7&parId=root&parQt=sharedby&o=OneUp) and change the `beats_ckpt_path` path in `conf/beats_bart_pt.yaml`
-3. Launch with `run_pt.sh`
-
-For running fine-tuning, inference and scoring:
-1. Set the value of `pt_tag` in `run_ft.sh` and
-2. Launch with `run_ft.sh`
+3. Launch with `run.sh`
 
 
 ## Trained checkpoints
 AAC Pre-trained model: https://huggingface.co/espnet/DCASE23.AudioCaptioning.PreTrained
 
-
 Fine-tuned model: https://huggingface.co/espnet/DCASE23.AudioCaptioning.FineTuned
 
 
 ### GPU Time
-AAC pre-training takes around ~12 hours on a single A40
-Fine-tuning, decoding and evalution takes ~2.5 hours on a single A40.
-All the scripts above are setup with 2 GPUs but that can be changed with the `ngpu` argument.
+AAC pre-training takes around ~4 hours on two A6000 gpus.
+Fine-tuning, decoding and evalution takes ~1 hour on two A6000.
+All the scripts above are setup with 2 GPUs but that can be changed with the `ngpu` argument. Please make sure to change the batch size accordingly (the provided setup requires 42.6GB GPU memory).
 
 
 <!-- Generated by scripts/utils/show_asr_result.sh -->
 # RESULTS
 ## Environments
-- date: `Tue Nov  5 15:18:41 CST 2024`
-- python version: `3.9.18 | packaged by conda-forge | (main, Dec 23 2023, 16:33:10)  [GCC 12.3.0]`
+- date: `Fri Nov 29 20:06:53 EST 2024`
+- python version: `3.9.20 (main, Oct  3 2024, 07:27:41)  [GCC 11.2.0]`
 - espnet version: `espnet 202409`
 - pytorch version: `pytorch 2.4.0`
-- Git hash: `ac484495e6e3cbf58bd3d1175c607c3f05bf6898`
-  - Commit date: `Tue Nov 5 09:24:48 2024 -0600`
+- Git hash: `65ea259e8effab5a43cdff87161a301dc0f20930`
+  - Commit date: `Fri Nov 29 10:54:44 2024 -0500`
 
 <!-- Copied from the output produced by local/evaluation.py -->
-## exp/asr_ft_lr5e-5.initfix.bigbatch512.lr2e-4.weighted12layers.20241103.145125
+## exp/asr_ft
 ```
 =====================================================
  Split: evaluation Evaluation over 1045 predictions.
 =====================================================
- cider_d             : 0.40163786973038024
- spice               : 0.12328679455316437
- spider              : 0.2624623321417723
- sbert_sim           : 0.5094392538748004
- fer                 : 0.03827751196172249
- fense               : 0.4915462823319092
- meteor              : 0.17194480402701573
- rouge_l             : 0.3550182495603068
- fer.add_tail_prob   : 0.04409123957157135
- fer.repeat_event_prob: 0.07279406487941742
- fer.repeat_adv_prob : 0.0016611652681604028
- fer.remove_conj_prob: 0.10949967056512833
- fer.remove_verb_prob: 0.20523187518119812
- fer.error_prob      : 0.319917231798172
- spider_fl           : 0.2532289450419614
+ cider_d             : 0.46045061153488653
+ spice               : 0.1345877073651595
+ spider              : 0.297519159450023
+ sbert_sim           : 0.5147198918907185
+ fer                 : 0.019138755980861243
+ fense               : 0.5048554784476391
+ meteor              : 0.18487611036251259
+ rouge_l             : 0.39408182293804006
+ fer.add_tail_prob   : 0.046671394258737564
+ fer.repeat_event_prob: 0.07274453341960907
+ fer.repeat_adv_prob : 0.0019690715707838535
+ fer.remove_conj_prob: 0.12462866306304932
+ fer.remove_verb_prob: 0.22472403943538666
+ fer.error_prob      : 0.34762266278266907
+ spider_fl           : 0.2923702923078383
 =====================================================
 ```
diff --git a/egs2/clotho_v2/asr1/conf/slurm.conf b/egs2/clotho_v2/asr1/conf/slurm.conf
index e4ce8f0ad12..bb24e661799 100644
--- a/egs2/clotho_v2/asr1/conf/slurm.conf
+++ b/egs2/clotho_v2/asr1/conf/slurm.conf
@@ -12,5 +12,5 @@ option num_threads=* --cpus-per-task $0
 option num_threads=1 --cpus-per-task 1
 option num_nodes=* --nodes $0
 default gpu=0
-option gpu=0 -p array
-option gpu=* -p debug --gres=gpu:$0 -c $0
+option gpu=0 -p cpu
+option gpu=* -p gpu --gres=gpu:$0 -c $0

From 6b95b4025750c500fa722b7c4fbd2ef30653dac2 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Sat, 30 Nov 2024 14:48:22 +0000
Subject: [PATCH 33/47] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 egs2/clotho_v2/asr1/README.md                         |  2 +-
 egs2/clotho_v2/asr1/conf/bart_decoder_config.json     |  2 +-
 egs2/clotho_v2/asr1/conf/beats_bart_ft.yaml           |  4 ++--
 egs2/clotho_v2/asr1/conf/beats_bart_pt.yaml           |  6 +++---
 .../asr1/conf/wav2vec2_conformer_config.json          |  2 +-
 egs2/clotho_v2/asr1/local/data_prep_audiocaps.py      |  1 +
 egs2/clotho_v2/asr1/local/data_prep_clotho.py         |  5 +++--
 egs2/clotho_v2/asr1/local/evaluation.py               |  2 +-
 egs2/clotho_v2/asr1/run.sh                            |  2 +-
 .../asr/decoder/hugging_face_transformers_decoder.py  |  3 ++-
 espnet2/asr/encoder/beats_encoder.py                  | 11 +++++------
 11 files changed, 21 insertions(+), 19 deletions(-)

diff --git a/egs2/clotho_v2/asr1/README.md b/egs2/clotho_v2/asr1/README.md
index c5b653e2dfa..79a6d4e7244 100644
--- a/egs2/clotho_v2/asr1/README.md
+++ b/egs2/clotho_v2/asr1/README.md
@@ -1,7 +1,7 @@
 # Clotho Audio Captioning RECIPE
 
 This recipe implements the DCASE 2023 Automated Audio Captioning (AAC) task with a BEATs encoder BART decoder model on the Clotho_v2 dataset, very close to what is described in [this paper](https://arxiv.org/abs/2309.17352) and reusing part of code from the [original implementation](https://github.com/slSeanWU/beats-conformer-bart-audio-captioner?tab=readme-ov-file).
-More specifically, we provide the pre-training and fine-tuning code for second last row in Table 2 (without Instructor embedding) from the paper. 
+More specifically, we provide the pre-training and fine-tuning code for second last row in Table 2 (without Instructor embedding) from the paper.
 We also reuse the code from the [BEATs repository](https://github.com/microsoft/unilm/tree/master/beats) for this implementation.
 
 # Training Details and Requirements
diff --git a/egs2/clotho_v2/asr1/conf/bart_decoder_config.json b/egs2/clotho_v2/asr1/conf/bart_decoder_config.json
index 6b940da0f2c..d5751dd5e8e 100644
--- a/egs2/clotho_v2/asr1/conf/bart_decoder_config.json
+++ b/egs2/clotho_v2/asr1/conf/bart_decoder_config.json
@@ -13,4 +13,4 @@
     "scale_embedding": false,
     "use_cache": true,
     "ignore_mismatched_sizes": true
-}
\ No newline at end of file
+}
diff --git a/egs2/clotho_v2/asr1/conf/beats_bart_ft.yaml b/egs2/clotho_v2/asr1/conf/beats_bart_ft.yaml
index 94ba6bc997d..bec8156b585 100644
--- a/egs2/clotho_v2/asr1/conf/beats_bart_ft.yaml
+++ b/egs2/clotho_v2/asr1/conf/beats_bart_ft.yaml
@@ -28,7 +28,7 @@ freeze_param: [
 
 encoder: beats
 encoder_conf:
-    # Please download the BEATs model from https://github.com/microsoft/unilm/tree/master/beats 
+    # Please download the BEATs model from https://github.com/microsoft/unilm/tree/master/beats
     # (iter3+, Beats finetuned model 1) and update the path below
     beats_ckpt_path: /compute/babel-13-33/sbharad2/models/BEATs/BEATs_iter3_plus_AS2M_finetuned_on_AS2M_cpt1.pt
     specaug_config:
@@ -49,7 +49,7 @@ encoder_conf:
     add_positional_information: true
     max_positions: 1024 # These many positional embeddings will be learned
 
-# Pleae note that the decoder is not the standard BART-base, 
+# Pleae note that the decoder is not the standard BART-base,
 # but a custom one whose config is defined in the file below
 decoder: hugging_face_transformers
 decoder_conf:
diff --git a/egs2/clotho_v2/asr1/conf/beats_bart_pt.yaml b/egs2/clotho_v2/asr1/conf/beats_bart_pt.yaml
index 3d9ffcdb7f8..c44a2a9b059 100644
--- a/egs2/clotho_v2/asr1/conf/beats_bart_pt.yaml
+++ b/egs2/clotho_v2/asr1/conf/beats_bart_pt.yaml
@@ -30,7 +30,7 @@ freeze_param: [
 
 encoder: beats
 encoder_conf:
-    # Please download the BEATs model from https://github.com/microsoft/unilm/tree/master/beats 
+    # Please download the BEATs model from https://github.com/microsoft/unilm/tree/master/beats
     # (iter3+, Beats finetuned model 1) and update the path below
     beats_ckpt_path: /compute/babel-13-33/sbharad2/models/BEATs/BEATs_iter3_plus_AS2M_finetuned_on_AS2M_cpt1.pt
     specaug_config:
@@ -51,7 +51,7 @@ encoder_conf:
     add_positional_information: true
     max_positions: 1024 # These many positional embeddings will be learned
 
-# Pleae note that the decoder is not the standard BART-base, 
+# Pleae note that the decoder is not the standard BART-base,
 # but a custom one whose config is defined in the file below
 decoder: hugging_face_transformers
 decoder_conf:
@@ -60,7 +60,7 @@ decoder_conf:
     load_pretrained_weights: false
     ensure_untied_lm_head: true
 
-# This takes care of BART and adapter layers' initialization 
+# This takes care of BART and adapter layers' initialization
 # (essentially everything else other than BEATs encoder backbone)
 init: normal
 
diff --git a/egs2/clotho_v2/asr1/conf/wav2vec2_conformer_config.json b/egs2/clotho_v2/asr1/conf/wav2vec2_conformer_config.json
index 9e175c60a21..23c3efba9a6 100644
--- a/egs2/clotho_v2/asr1/conf/wav2vec2_conformer_config.json
+++ b/egs2/clotho_v2/asr1/conf/wav2vec2_conformer_config.json
@@ -110,4 +110,4 @@
     "transformers_version": "4.19.0.dev0",
     "use_weighted_layer_sum": false,
     "xvector_output_dim": 512
-  }
\ No newline at end of file
+  }
diff --git a/egs2/clotho_v2/asr1/local/data_prep_audiocaps.py b/egs2/clotho_v2/asr1/local/data_prep_audiocaps.py
index e0ec767d5d5..4cbc86c9fff 100644
--- a/egs2/clotho_v2/asr1/local/data_prep_audiocaps.py
+++ b/egs2/clotho_v2/asr1/local/data_prep_audiocaps.py
@@ -4,6 +4,7 @@
 import os
 import sys
 from string import punctuation
+
 from scipy.io import wavfile
 
 strip_punct_table = str.maketrans("", "", punctuation)
diff --git a/egs2/clotho_v2/asr1/local/data_prep_clotho.py b/egs2/clotho_v2/asr1/local/data_prep_clotho.py
index 91ad9173577..c887b1e611b 100644
--- a/egs2/clotho_v2/asr1/local/data_prep_clotho.py
+++ b/egs2/clotho_v2/asr1/local/data_prep_clotho.py
@@ -2,11 +2,12 @@
 
 import glob
 import os
-import sys
 import random
-import pandas as pd
+import sys
 from string import punctuation
 
+import pandas as pd
+
 strip_punct_table = str.maketrans("", "", punctuation)
 
 if __name__ == "__main__":
diff --git a/egs2/clotho_v2/asr1/local/evaluation.py b/egs2/clotho_v2/asr1/local/evaluation.py
index 11c14067760..6571d979b68 100644
--- a/egs2/clotho_v2/asr1/local/evaluation.py
+++ b/egs2/clotho_v2/asr1/local/evaluation.py
@@ -2,9 +2,9 @@
 to evaluate the audio captioning predictions."""
 
 import argparse
+from string import punctuation
 
 from aac_metrics import Evaluate
-from string import punctuation
 
 strip_punct_table = str.maketrans("", "", punctuation)
 
diff --git a/egs2/clotho_v2/asr1/run.sh b/egs2/clotho_v2/asr1/run.sh
index 2403bbe76ed..64b4e3c8080 100755
--- a/egs2/clotho_v2/asr1/run.sh
+++ b/egs2/clotho_v2/asr1/run.sh
@@ -54,4 +54,4 @@ inference_ckpt=latest
     --asr_config conf/beats_bart_ft.yaml \
     --pretrained_model exp/asr_pt/${pt_ckpt}.pth \
     --inference_asr_model ${inference_ckpt}.pth \
-    --local_score_opts exp/asr_ft/inference_ctc_weight0.0_hugging_face_decoderTrue_asr_model_${inference_ckpt}
\ No newline at end of file
+    --local_score_opts exp/asr_ft/inference_ctc_weight0.0_hugging_face_decoderTrue_asr_model_${inference_ckpt}
diff --git a/espnet2/asr/decoder/hugging_face_transformers_decoder.py b/espnet2/asr/decoder/hugging_face_transformers_decoder.py
index 6735dbc9f22..af15b9518eb 100644
--- a/espnet2/asr/decoder/hugging_face_transformers_decoder.py
+++ b/espnet2/asr/decoder/hugging_face_transformers_decoder.py
@@ -7,7 +7,8 @@
 import copy
 import logging
 import os
-from typing import Any, List, Tuple, Optional
+from typing import Any, List, Optional, Tuple
+
 import torch
 import torch.nn.functional as F
 from typeguard import typechecked
diff --git a/espnet2/asr/encoder/beats_encoder.py b/espnet2/asr/encoder/beats_encoder.py
index 157ce678c11..58a8b13a6e5 100644
--- a/espnet2/asr/encoder/beats_encoder.py
+++ b/espnet2/asr/encoder/beats_encoder.py
@@ -23,16 +23,15 @@
 import torchaudio.compliance.kaldi as ta_kaldi
 from packaging.version import parse as V
 from torch.nn import LayerNorm, Parameter
-
-from espnet2.asr.encoder.abs_encoder import AbsEncoder
-from espnet2.asr.specaug.specaug import SpecAug
-from espnet.nets.pytorch_backend.nets_utils import make_pad_mask
-
+from transformers.models.bart.modeling_bart import BartLearnedPositionalEmbedding
 from transformers.models.wav2vec2_conformer.modeling_wav2vec2_conformer import (
     Wav2Vec2ConformerConfig,
     Wav2Vec2ConformerEncoder,
 )
-from transformers.models.bart.modeling_bart import BartLearnedPositionalEmbedding
+
+from espnet2.asr.encoder.abs_encoder import AbsEncoder
+from espnet2.asr.specaug.specaug import SpecAug
+from espnet.nets.pytorch_backend.nets_utils import make_pad_mask
 
 if V(torch.__version__) >= V("1.6.0"):
     from torch.cuda.amp import autocast

From 32092fa75f370dabd160294ad487b2d5b0cc74c8 Mon Sep 17 00:00:00 2001
From: Shikhar Bharadwaj <sbharad2@andrew.cmu.edu>
Date: Sat, 30 Nov 2024 12:06:20 -0500
Subject: [PATCH 34/47] restore db.sh

---
 egs2/TEMPLATE/asr1/db.sh | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/egs2/TEMPLATE/asr1/db.sh b/egs2/TEMPLATE/asr1/db.sh
index 7f5b8e2da90..1393b47bc35 100755
--- a/egs2/TEMPLATE/asr1/db.sh
+++ b/egs2/TEMPLATE/asr1/db.sh
@@ -221,10 +221,9 @@ SPRING_INX=downloads
 VOXCELEB=
 KSPONSPEECH=
 HIFITTS=downloads
-
 CLOTHO_V2=downloads
-AUDIOCAPS=/compute/babel-13-33/sbharad2/audiocaps
-CLOTHO_CHATGPT_MIXUP=/compute/babel-13-33/sbharad2/chatgpt
+AUDIOCAPS=
+CLOTHO_CHATGPT_MIXUP=
 
 # For only CMU TIR environment
 if [[ "$(hostname)" == tir* ]]; then

From dded1d5a9cd9e33bb2462e9a74726ce0d872b815 Mon Sep 17 00:00:00 2001
From: Shikhar Bharadwaj <sbharad2@andrew.cmu.edu>
Date: Mon, 2 Dec 2024 10:51:34 -0500
Subject: [PATCH 35/47] fix ci errors: line length, and transformers import

---
 egs2/clotho_v2/asr1/local/data_prep_clotho.py |  3 +-
 .../hugging_face_transformers_decoder.py      | 35 +++++++++++--------
 espnet2/asr/encoder/beats_encoder.py          | 27 +++++++++++---
 3 files changed, 45 insertions(+), 20 deletions(-)

diff --git a/egs2/clotho_v2/asr1/local/data_prep_clotho.py b/egs2/clotho_v2/asr1/local/data_prep_clotho.py
index c887b1e611b..8e3a314103d 100644
--- a/egs2/clotho_v2/asr1/local/data_prep_clotho.py
+++ b/egs2/clotho_v2/asr1/local/data_prep_clotho.py
@@ -111,5 +111,6 @@
                     f.close()
         print(
             f"For split {data_split}: Processed {N_PROCESSED} audio and their"
-            f" {N_CAPTIONS_PROCESSED} captions. {len(missing_audio)} audio files were missing."
+            f" {N_CAPTIONS_PROCESSED} captions. {len(missing_audio)} audio files"
+            "were missing."
         )
diff --git a/espnet2/asr/decoder/hugging_face_transformers_decoder.py b/espnet2/asr/decoder/hugging_face_transformers_decoder.py
index af15b9518eb..b35cee8f896 100644
--- a/espnet2/asr/decoder/hugging_face_transformers_decoder.py
+++ b/espnet2/asr/decoder/hugging_face_transformers_decoder.py
@@ -55,23 +55,30 @@ def __init__(
         Args:
             vocab_size (int): The size of the vocabulary.
             encoder_output_size (int): The size of the encoder output.
-            model_name_or_path (str): The name or path of the pre-trained Transformers model.
-            causal_lm (bool, optional): Whether to use a causal language model. Defaults to False.
-                This overrides the model_name_or_path if provided.
-            prefix (str, optional): Prefix to be added to the input tokens. Defaults to "".
-            postfix (str, optional): Postfix to be added to the input tokens. Defaults to "".
-            config_path (str, optional): Path to the configuration json file. Defaults to None. If
-                this is set, it can be used to override the default decoder configuration. More documentation here:
-                https://huggingface.co/docs/transformers/en/model_doc/auto#transformers.AutoModelForSeq2SeqLM.from_pretrained.kwargs
-            load_pretrained_weights (bool): Whether to load the pre-trained weights. Defaults to True.
-            ensure_untied_lm_head (bool): True ensures that the language model head is not
-                shared with the input token embeddings. When False, the original structure is
-                kept, ie, if the original Transformers implementation has tying of weights,
-                it is retained. Defaults to False.
+            model_name_or_path (str): The name or path of the pre-trained
+                 Transformers model.
+            causal_lm (bool, optional): Whether to use a causal language
+                model. Defaults to False. This overrides the
+                model_name_or_path if provided.
+            prefix (str, optional): Prefix to be added to the input
+                tokens. Defaults to "".
+            postfix (str, optional): Postfix to be added to the input
+                tokens. Defaults to "".
+            config_path (str, optional): Path to the configuration json
+                file. Defaults to None. If this is set, it can be used to
+                override the default decoder configuration.
+            load_pretrained_weights (bool): Whether to load the pre-trained
+                weights. Defaults to True.
+            ensure_untied_lm_head (bool): True ensures that the language model
+                head is not shared with the input token embeddings. When False,
+                the original structure is kept, ie, if the original Transformers
+                implementation has tying of weights, it is retained. Defaults
+                to False.
 
         Raises:
             ImportError: If the `transformers` library is not available.
-            Exception: If the word embeddings attribute cannot be found in the model.
+            Exception: If the word embeddings attribute cannot be found in
+                the model.
         """
         super().__init__()
 
diff --git a/espnet2/asr/encoder/beats_encoder.py b/espnet2/asr/encoder/beats_encoder.py
index 58a8b13a6e5..14adcbdd113 100644
--- a/espnet2/asr/encoder/beats_encoder.py
+++ b/espnet2/asr/encoder/beats_encoder.py
@@ -23,11 +23,18 @@
 import torchaudio.compliance.kaldi as ta_kaldi
 from packaging.version import parse as V
 from torch.nn import LayerNorm, Parameter
-from transformers.models.bart.modeling_bart import BartLearnedPositionalEmbedding
-from transformers.models.wav2vec2_conformer.modeling_wav2vec2_conformer import (
-    Wav2Vec2ConformerConfig,
-    Wav2Vec2ConformerEncoder,
-)
+
+try:
+    from transformers.models.bart.modeling_bart import BartLearnedPositionalEmbedding
+    from transformers.models.wav2vec2_conformer.modeling_wav2vec2_conformer import (
+        Wav2Vec2ConformerConfig,
+        Wav2Vec2ConformerEncoder,
+    )
+
+    is_transformers_available = True
+except ImportError:
+    is_transformers_available = False
+
 
 from espnet2.asr.encoder.abs_encoder import AbsEncoder
 from espnet2.asr.specaug.specaug import SpecAug
@@ -155,6 +162,16 @@ def __init__(
         # 3. Checkpoint and user-provided config: Merge the two, but
         #    override with user-provided config
         # 4. No checkpoint and user-provided config: Use user-provided config
+        if adapter_config or add_positional_information:
+            # We need transformers library for adapter and positional embeddings
+            if not is_transformers_available:
+                raise ImportError(
+                    "`transformers` is not available. Please install it "
+                    " via `pip install transformers` or"
+                    " `cd /path/to/espnet/tools && "
+                    ". ./activate_python.sh"
+                    " && ./installers/install_transformers.sh`."
+                )
         config = BEATsConfig()  # Default config
         if beats_ckpt_path and beats_config:
             logging.warning(

From 79cf4972132c00104aad1f40fa77c1b24fb9259a Mon Sep 17 00:00:00 2001
From: Shikhar Bharadwaj <sbharad2@andrew.cmu.edu>
Date: Mon, 2 Dec 2024 18:29:34 -0500
Subject: [PATCH 36/47] remove unnecessary runner files and imports

---
 egs2/clotho_v2/asr1/run_ft.sh        | 45 ----------------------------
 egs2/clotho_v2/asr1/run_inference.sh | 38 -----------------------
 egs2/clotho_v2/asr1/run_pt.sh        | 40 -------------------------
 espnet2/bin/asr_inference.py         |  2 +-
 4 files changed, 1 insertion(+), 124 deletions(-)
 delete mode 100755 egs2/clotho_v2/asr1/run_ft.sh
 delete mode 100755 egs2/clotho_v2/asr1/run_inference.sh
 delete mode 100755 egs2/clotho_v2/asr1/run_pt.sh

diff --git a/egs2/clotho_v2/asr1/run_ft.sh b/egs2/clotho_v2/asr1/run_ft.sh
deleted file mode 100755
index 28cb4f516c8..00000000000
--- a/egs2/clotho_v2/asr1/run_ft.sh
+++ /dev/null
@@ -1,45 +0,0 @@
-#!/usr/bin/env bash
-set -euo pipefail
-# Runs fine-tuning
-
-wandb_init_args=""
-other_args="$@"
-
-expdir=/compute/babel-13-33/sbharad2/expdir
-dumpdir=/compute/babel-13-33/sbharad2/dumpdir
-
-pt_tag=scratch_bart.splen30s.lr2e-4.lossnormT.20241120.113058
-
-ckpt_name=valid.acc.ave
-ckpt_name=10epoch
-
-pre_trained_model_path=${expdir}/asr_pt.${pt_tag}/${ckpt_name}.pth
-ft_tag=${pt_tag}.match_contest.${ckpt_name}
-asr_speech_fold_length=4800 # 480000/16000 = 30 seconds
-
-./asr.sh \
-    --asr_tag ft.${ft_tag} \
-    --expdir ${expdir} \
-    --dumpdir ${dumpdir} \
-    --feats_normalize uttmvn \
-    --stage 11 \
-    --asr_speech_fold_length ${asr_speech_fold_length} \
-    --stop_stage 13 \
-    --asr_stats_dir ${expdir}/asr_stats_finetune \
-    --ngpu 2 \
-    --gpu_inference true \
-    --nj 8 \
-    --inference_nj 1 \
-    --max_wav_duration 30 \
-    --token_type hugging_face \
-    --use_lm false \
-    --hugging_face_model_name_or_path "facebook/bart-base" \
-    --inference_args "--beam_size 10 --ctc_weight 0.0 --hugging_face_decoder True" \
-    --train_set development_clotho \
-    --valid_set validation \
-    --test_sets "evaluation" \
-    --asr_config conf/beats_bart_ft.yaml \
-    --pretrained_model ${pre_trained_model_path} \
-    --inference_asr_model valid.acc.ave_5best.pth \
-    --asr_args "${wandb_init_args} ${other_args}" \
-    --local_score_opts ${expdir}/asr_ft.${ft_tag}/inference_beam_size10_ctc_weight0.0_hugging_face_decoderTrue_asr_model_${ckpt_name}
diff --git a/egs2/clotho_v2/asr1/run_inference.sh b/egs2/clotho_v2/asr1/run_inference.sh
deleted file mode 100755
index 85c35577447..00000000000
--- a/egs2/clotho_v2/asr1/run_inference.sh
+++ /dev/null
@@ -1,38 +0,0 @@
-#!/usr/bin/env bash
-# Runs pre-training
-set -euo pipefail
-
-# # Run inference on all checkpoints
-
-wandb_sweep_args="$@"
-
-name_tag=ft.beats.bartfrozen.20241111.164042.valid.acc.ave
-all_ckpts=$(find ${expdir}/asr_${name_tag} -name "*.pth"  -printf "%f\n" | awk '{print substr($0, 1, length($0) - 4)}')
-
-# for ckpt in valid.acc.best; do
-for ckpt in $all_ckpts; do
-    # We skip all symlinks
-    if [[ $ckpt == "checkpoint" || $ckpt == "latest"  || $ckpt == "valid.acc.best"  || $ckpt == "valid.acc.ave.pth" ]]; then
-        continue
-    fi
-    ./asr.sh \
-        --asr_tag $name_tag \
-        --feats_normalize uttmvn \
-        --stage 12 \
-        --stop_stage 13 \
-        --ngpu 2 \
-        --gpu_inference true \
-        --nj 20 \
-        --inference_nj 1 \
-        --max_wav_duration 30 \
-        --token_type hugging_face \
-        --use_lm false \
-        --hugging_face_model_name_or_path "facebook/bart-base" \
-        --inference_args "--beam_size 10 --ctc_weight 0.0 --hugging_face_decoder True" \
-        --train_set pretrain \
-        --valid_set validation \
-        --test_sets "evaluation" \
-        --inference_asr_model ${ckpt}.pth \
-        --asr_args "${wandb_sweep_args}" \
-        --local_score_opts "${expdir}/asr_${name_tag}/inference_beam_size10_ctc_weight0.0_hugging_face_decoderTrue_asr_model_${ckpt}"
-done
diff --git a/egs2/clotho_v2/asr1/run_pt.sh b/egs2/clotho_v2/asr1/run_pt.sh
deleted file mode 100755
index c44b227e1ff..00000000000
--- a/egs2/clotho_v2/asr1/run_pt.sh
+++ /dev/null
@@ -1,40 +0,0 @@
-#!/usr/bin/env bash
-# Runs pre-training
-set -euo pipefail
-
-timestamp=$(date "+%Y%m%d.%H%M%S")
-mynametag=pt.cleanacaps.bs512-warmup8k.load_bart.3ladapter.${timestamp}
-
-wandb_init_args="--use_wandb true --wandb_entity shikhar --wandb_name ${mynametag} --wandb_project DCASE_AAC --wandb_model_log_interval 0"
-other_args="$@"
-
-expdir=/compute/babel-13-33/sbharad2/expdir
-dumpdir=/compute/babel-13-33/sbharad2/dumpdir_clean
-local_data_opts=/compute/babel-13-33/sbharad2/expdir
-asr_speech_fold_length=4800 # 480000/16000 = 30 seconds
-
-./asr.sh \
-    --asr_tag ${mynametag} \
-    --expdir ${expdir} \
-    --dumpdir ${dumpdir} \
-    --asr_speech_fold_length ${asr_speech_fold_length} \
-    --local_data_opts ${local_data_opts} \
-    --feats_normalize uttmvn \
-    --stage 11 \
-    --stop_stage 13 \
-    --ngpu 2 \
-    --gpu_inference true \
-    --nj 10 \
-    --inference_nj 1 \
-    --max_wav_duration 30 \
-    --token_type hugging_face \
-    --use_lm false \
-    --hugging_face_model_name_or_path "facebook/bart-base" \
-    --inference_args "--beam_size 10 --ctc_weight 0.0 --hugging_face_decoder True" \
-    --train_set pretrain \
-    --valid_set validation \
-    --test_sets "validation evaluation" \
-    --asr_config conf/beats_bart_pt.yaml \
-    --inference_asr_model valid.acc.best.pth \
-    --asr_args "${wandb_init_args} ${other_args}" \
-    --local_score_opts "exp/asr_${mynametag}/inference_beam_size10_ctc_weight0.0_hugging_face_decoderTrue_asr_model_valid.acc.best"
diff --git a/espnet2/bin/asr_inference.py b/espnet2/bin/asr_inference.py
index 0b0929fead4..ad88258fec3 100644
--- a/espnet2/bin/asr_inference.py
+++ b/espnet2/bin/asr_inference.py
@@ -50,7 +50,6 @@
 try:
     from transformers import AutoModelForCausalLM, AutoModelForSeq2SeqLM
     from transformers.file_utils import ModelOutput
-    from transformers.models.bart.modeling_bart import BartConfig, BartDecoder
 
     is_transformers_available = True
 except ImportError:
@@ -298,6 +297,7 @@ def __init__(
                 hugging_face_decoder_conf["pad_token_id"] = (
                     hugging_face_model.config.eos_token_id
                 )
+
             beam_search = None
             beam_search_transducer = None
         else:

From da52e553789833c8d3b4663deddc62e52e0358b2 Mon Sep 17 00:00:00 2001
From: Shikhar Bharadwaj <sbharad2@andrew.cmu.edu>
Date: Tue, 3 Dec 2024 07:39:17 -0500
Subject: [PATCH 37/47] restore default cmd.sh

---
 egs2/clotho_v2/asr1/cmd.sh | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/egs2/clotho_v2/asr1/cmd.sh b/egs2/clotho_v2/asr1/cmd.sh
index 257c6c739df..2aae6919fef 100644
--- a/egs2/clotho_v2/asr1/cmd.sh
+++ b/egs2/clotho_v2/asr1/cmd.sh
@@ -28,7 +28,7 @@
 
 
 # Select the backend used by run.sh from "local", "stdout", "sge", "slurm", or "ssh"
-cmd_backend='slurm'
+cmd_backend='local'
 
 # Local machine, without any Job scheduling system
 if [ "${cmd_backend}" = local ]; then

From 27f5c9d66077b41468dec23f36dbf1f8422ad0ee Mon Sep 17 00:00:00 2001
From: Shikhar Bharadwaj <sbharad2@andrew.cmu.edu>
Date: Tue, 3 Dec 2024 07:50:54 -0500
Subject: [PATCH 38/47] restore default slurm.conf

---
 egs2/clotho_v2/asr1/conf/slurm.conf | 8 +++-----
 1 file changed, 3 insertions(+), 5 deletions(-)

diff --git a/egs2/clotho_v2/asr1/conf/slurm.conf b/egs2/clotho_v2/asr1/conf/slurm.conf
index bb24e661799..3b229673638 100644
--- a/egs2/clotho_v2/asr1/conf/slurm.conf
+++ b/egs2/clotho_v2/asr1/conf/slurm.conf
@@ -1,16 +1,14 @@
 # Default configuration
 command sbatch --export=PATH
 option name=* --job-name $0
-default time=2-00:00:00
 option time=* --time $0
-default mem=10G
 option mem=* --mem-per-cpu $0
-option mem=10G --mem-per-cpu 10G
 option mem=0
-default num_threads=4
 option num_threads=* --cpus-per-task $0
 option num_threads=1 --cpus-per-task 1
 option num_nodes=* --nodes $0
 default gpu=0
 option gpu=0 -p cpu
-option gpu=* -p gpu --gres=gpu:$0 -c $0
+option gpu=* -p gpu --gres=gpu:$0 -c $0  # Recommend allocating more CPU than, or equal to the number of GPU
+# note: the --max-jobs-run option is supported as a special case
+# by slurm.pl and you don't have to handle it in the config file.

From 9fdef3af266c0f93a5a6f68639a09a6e8c53d4d8 Mon Sep 17 00:00:00 2001
From: Shikhar Bharadwaj <sbharad2@andrew.cmu.edu>
Date: Fri, 6 Dec 2024 10:16:29 -0500
Subject: [PATCH 39/47] remove upload script

---
 egs2/clotho_v2/asr1/run_upload.sh | 49 -------------------------------
 1 file changed, 49 deletions(-)
 delete mode 100755 egs2/clotho_v2/asr1/run_upload.sh

diff --git a/egs2/clotho_v2/asr1/run_upload.sh b/egs2/clotho_v2/asr1/run_upload.sh
deleted file mode 100755
index 1667bf91b4c..00000000000
--- a/egs2/clotho_v2/asr1/run_upload.sh
+++ /dev/null
@@ -1,49 +0,0 @@
-#!/usr/bin/env bash
-set -euo pipefail
-# Runs model upload to huggingface
-
-# wandb_init_args="--use_wandb true --wandb_project DCASE_AAC --wandb_model_log_interval 0"
-wandb_init_args=""
-other_args="$@"
-
-expdir=/compute/babel-13-33/sbharad2/expdir
-dumpdir=/compute/babel-13-33/sbharad2/dumpdir
-local_data_opts=/compute/babel-13-33/sbharad2/expdir
-
-
-expdir=exp/
-# asr_tag=pt.initfix.bigbatch512.lr2e-4.weighted12layers.20241103.145125
-# hf_repo=shikhar7ssu/dcase23.aac.pt
-
-asr_tag=ft_lr5e-5.initfix.bigbatch512.lr2e-4.weighted12layers.20241103.145125
-hf_repo=espnet/DCASE23.AudioCaptioning.FineTuned
-
-
-./asr.sh \
-    --asr_tag ${asr_tag} \
-    --hf_repo ${hf_repo} \
-    --skip_packing false \
-    --skip_upload_hf false \
-    --lang en \
-    --expdir ${expdir} \
-    --dumpdir ${dumpdir} \
-    --feats_normalize uttmvn \
-    --stage 14 \
-    --stop_stage 15 \
-    --asr_stats_dir ${expdir}/asr_stats_finetune \
-    --ngpu 2 \
-    --gpu_inference true \
-    --nj 8 \
-    --inference_nj 1 \
-    --max_wav_duration 30 \
-    --token_type hugging_face \
-    --use_lm false \
-    --hugging_face_model_name_or_path "facebook/bart-base" \
-    --inference_args "--beam_size 10 --ctc_weight 0.0 --hugging_face_decoder True" \
-    --train_set development_clotho \
-    --valid_set validation \
-    --test_sets "evaluation" \
-    --asr_config conf/beats_bart_ft.yaml \
-    --inference_asr_model valid.acc.ave.pth \
-    --asr_args "${wandb_init_args} ${other_args}" \
-    --local_data_opts "${local_data_opts}"

From 1515739e1066274ad490d2b48377d9f1a6fb593f Mon Sep 17 00:00:00 2001
From: Shikhar Bharadwaj <sbharad2@andrew.cmu.edu>
Date: Fri, 6 Dec 2024 21:30:45 -0500
Subject: [PATCH 40/47] rename variables, remove unnecessary utility function
 from hf decoder and correct assertion in beats

---
 egs2/clotho_v2/asr1/conf/beats_bart_ft.yaml   |  4 +-
 egs2/clotho_v2/asr1/conf/beats_bart_pt.yaml   |  4 +-
 .../hugging_face_transformers_decoder.py      | 46 +++++++---------
 espnet2/asr/encoder/beats_encoder.py          | 53 +++++++++++--------
 4 files changed, 53 insertions(+), 54 deletions(-)

diff --git a/egs2/clotho_v2/asr1/conf/beats_bart_ft.yaml b/egs2/clotho_v2/asr1/conf/beats_bart_ft.yaml
index bec8156b585..c5637d4a8d7 100644
--- a/egs2/clotho_v2/asr1/conf/beats_bart_ft.yaml
+++ b/egs2/clotho_v2/asr1/conf/beats_bart_ft.yaml
@@ -54,9 +54,9 @@ encoder_conf:
 decoder: hugging_face_transformers
 decoder_conf:
     model_name_or_path: facebook/bart-base
-    config_path: conf/bart_decoder_config.json
+    overriding_architecture_config: conf/bart_decoder_config.json
     load_pretrained_weights: false
-    ensure_untied_lm_head: true
+    separate_lm_head: true
 
 # Initialization does not matter we use a pre-trained model
 init: normal
diff --git a/egs2/clotho_v2/asr1/conf/beats_bart_pt.yaml b/egs2/clotho_v2/asr1/conf/beats_bart_pt.yaml
index c44a2a9b059..05596dfb58b 100644
--- a/egs2/clotho_v2/asr1/conf/beats_bart_pt.yaml
+++ b/egs2/clotho_v2/asr1/conf/beats_bart_pt.yaml
@@ -56,9 +56,9 @@ encoder_conf:
 decoder: hugging_face_transformers
 decoder_conf:
     model_name_or_path: facebook/bart-base
-    config_path: conf/bart_decoder_config.json
+    overriding_architecture_config: conf/bart_decoder_config.json
     load_pretrained_weights: false
-    ensure_untied_lm_head: true
+    separate_lm_head: true
 
 # This takes care of BART and adapter layers' initialization
 # (essentially everything else other than BEATs encoder backbone)
diff --git a/espnet2/asr/decoder/hugging_face_transformers_decoder.py b/espnet2/asr/decoder/hugging_face_transformers_decoder.py
index b35cee8f896..9e71a5da8ba 100644
--- a/espnet2/asr/decoder/hugging_face_transformers_decoder.py
+++ b/espnet2/asr/decoder/hugging_face_transformers_decoder.py
@@ -7,7 +7,7 @@
 import copy
 import logging
 import os
-from typing import Any, List, Optional, Tuple
+from typing import Any, List, Optional, Tuple, Union
 
 import torch
 import torch.nn.functional as F
@@ -45,9 +45,9 @@ def __init__(
         causal_lm: bool = False,
         prefix: str = "",
         postfix: str = "",
-        config_path: Optional[str] = None,
+        overriding_architecture_config: Optional[Union[str, dict]] = {},
         load_pretrained_weights: bool = True,
-        ensure_untied_lm_head: bool = False,
+        separate_lm_head: bool = False,
     ):
         """
         Initializes the HuggingFaceTransformersDecoder.
@@ -64,12 +64,13 @@ def __init__(
                 tokens. Defaults to "".
             postfix (str, optional): Postfix to be added to the input
                 tokens. Defaults to "".
-            config_path (str, optional): Path to the configuration json
-                file. Defaults to None. If this is set, it can be used to
-                override the default decoder configuration.
+            overriding_architecture_config (str or dict, optional): Path to the configuration json
+                file or the json dictionary itself. Defaults to None. If this
+                is set, it can be used to override the default decoder
+                configuration.
             load_pretrained_weights (bool): Whether to load the pre-trained
                 weights. Defaults to True.
-            ensure_untied_lm_head (bool): True ensures that the language model
+            separate_lm_head (bool): True ensures that the language model
                 head is not shared with the input token embeddings. When False,
                 the original structure is kept, ie, if the original Transformers
                 implementation has tying of weights, it is retained. Defaults
@@ -90,12 +91,15 @@ def __init__(
             )
 
         self.load_pretrained_weights = load_pretrained_weights
-        self.ensure_untied_lm_head = ensure_untied_lm_head
-        self.overriding_architecture_config = (
-            vars(get_model_conf(model_path="", conf_path=config_path))
-            if config_path
-            else {}
-        )
+        self.separate_lm_head = separate_lm_head
+
+        self.overriding_architecture_config = overriding_architecture_config
+        if type(overriding_architecture_config) == str:
+            # It is path to a json config file
+            self.overriding_architecture_config = vars(
+                get_model_conf(model_path="", conf_path=overriding_architecture_config)
+            )
+
         self.causal_lm = causal_lm
 
         if self.causal_lm:
@@ -143,7 +147,7 @@ def __init__(
 
         model.resize_token_embeddings(vocab_size)
 
-        if self.ensure_untied_lm_head:
+        if self.separate_lm_head:
             self.lm_head = copy.deepcopy(get_hugging_face_model_lm_head(model))
         else:
             self.lm_head = get_hugging_face_model_lm_head(model)
@@ -361,17 +365,3 @@ def get_hugging_face_model_lm_head(model):
         raise Exception("Can not find the LM head attribute")
 
     return lm_head
-
-
-def read_config_file(config_path):
-    assert os.path.exists(
-        config_path
-    ), f"Config file supplied but not found: {config_path}"
-    config_dict = {}
-    with open(config_path, "r") as f:
-        config_dict = json.load(f)
-        config_dict["ignore_mismatched_sizes"] = (
-            True  # Loaded config could be different from pre-trained model.
-        )
-
-    return config_dict
diff --git a/espnet2/asr/encoder/beats_encoder.py b/espnet2/asr/encoder/beats_encoder.py
index 14adcbdd113..b1f3bdf6f61 100644
--- a/espnet2/asr/encoder/beats_encoder.py
+++ b/espnet2/asr/encoder/beats_encoder.py
@@ -49,7 +49,7 @@ def autocast(enabled=True):
         yield
 
 
-class BEATsConfig:
+class BeatsConfig:
     def __init__(self, cfg=None):
         self.input_patch_size: int = 16  # patch size of patch embedding
         self.embed_dim: int = 512  # patch embedding dimension
@@ -115,7 +115,7 @@ class BeatsEncoder(AbsEncoder):
 
     (https://arxiv.org/abs/2212.09058)
     Args:
-        beats_ckpt_path: Path to a pretrained BEATs checkpoint. If
+        beats_ckpt_path: Path to a pretrained Beats checkpoint. If
             `beats_config` is provided and it does not match the
             config in the checkpoint, code might throw an error.
         max_layer: Propagate input through all layers for encoding
@@ -124,7 +124,7 @@ class BeatsEncoder(AbsEncoder):
         adapter_config: Path to a config file for the wav2vec2 adapter.
         use_weighted_representation: Use weighted representations
             from max_layer if True. Weights are randomly initialized.
-        beats_config: `BEATsConfig` object. If provided, we will try
+        beats_config: `BeatsConfig` object. If provided, we will try
             to override the config in the checkpoint. This can be used
             to change dropouts etc for fine-tuning the model while
             starting from a pretrained checkpoint.
@@ -143,7 +143,7 @@ def __init__(
         downsampling_rate: int = 1,
         adapter_config: str = "",
         use_weighted_representation: bool = False,
-        beats_config: Optional[BEATsConfig] = None,
+        beats_config: Optional[BeatsConfig] = None,
         specaug_config: Optional[Dict] = None,
         add_positional_information: bool = False,
         max_positions: Optional[int] = None,
@@ -172,7 +172,7 @@ def __init__(
                     ". ./activate_python.sh"
                     " && ./installers/install_transformers.sh`."
                 )
-        config = BEATsConfig()  # Default config
+        config = BeatsConfig()  # Default config
         if beats_ckpt_path and beats_config:
             logging.warning(
                 "Both pretrained checkpoint and config are provided."
@@ -181,11 +181,11 @@ def __init__(
         self.loaded_state_dict_ = None
         if beats_ckpt_path is not None:
             self.loaded_state_dict_ = torch.load(beats_ckpt_path)
-            logging.info(f"Loaded BEATs pretrained config from {beats_ckpt_path}.")
-            config = BEATsConfig(self.loaded_state_dict_["cfg"])
+            logging.info(f"Loaded Beats pretrained config from {beats_ckpt_path}.")
+            config = BeatsConfig(self.loaded_state_dict_["cfg"])
         if beats_config is not None:
-            config.update(beats_config)
-            logging.info("Overriding BEATs config with user-provided config.")
+            config.update(vars(beats_config))
+            logging.info("Overriding Beats config with user-provided config.")
 
         self.specaug = None
         if specaug_config is not None:
@@ -260,16 +260,16 @@ def __init__(
             )
 
     def reload_pretrained_parameters(self):
-        """Initialization function for BEATs.
+        """Initialization function for Beats.
 
         This must be called last in the initialization procedure.
         The initialization occurs in three steps:
-        1. ESPNet initializes all modules.
-        2. This function initializes BEATs encoder overriding 1.
+        1. ESPnet initializes all modules.
+        2. This function initializes Beats encoder overriding 1.
         3. Optionally, if we have the pretrained checkpoint, we load the
             weights from the checkpoint overriding 2 and 1.
         """
-        logging.info("BEATs Initialization function called.")
+        logging.info("Beats Initialization function called.")
         if self.post_extract_proj:
             torch.nn.init.xavier_normal_(self.post_extract_proj.weight)
             if self.post_extract_proj.bias is not None:
@@ -278,7 +278,7 @@ def reload_pretrained_parameters(self):
         if self.patch_embedding.bias is not None:
             torch.nn.init.constant_(self.patch_embedding.bias, 0)
 
-        # BEATs has different initialization from ESPNet for other modules,
+        # Beats has different initialization from ESPnet for other modules,
         #  so override.
         self.encoder.apply(init_bert_params)
         if self.loaded_state_dict_ is not None:
@@ -288,12 +288,12 @@ def reload_pretrained_parameters(self):
             )
             # strict=False to ignore Weights in the predictor
             logging.info(
-                f"Loaded BEATs pretrained model. Following keys were missing"
+                f"Loaded Beats pretrained model. Following keys were missing"
                 f" in your custom model: {load_info.missing_keys}. "
                 f"Follwing keys could not be loaded from the pretrained"
                 f"checkpoint: {load_info.unexpected_keys}."
                 "It is expected to have 'predictor' listed above if you are"
-                "fine-tuning with only the BEATs backbone."
+                "fine-tuning with only the Beats backbone."
             )
 
     def forward_padding_mask(
@@ -338,7 +338,16 @@ def forward(
         ilens: torch.Tensor,
         prev_states: torch.Tensor = None,
     ) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]:
-        """Wrapper for compatibility with ESPNets' AbsEncoder Interface."""
+        """Wrapper for compatibility with ESPnets' AbsEncoder Interface.
+        Args:
+            xs_pad: (B, T, D)
+            ilens: (B,)
+            prev_states: None
+        Returns:
+            audio_representation: (B, T, D)
+            output_lens: (B,)
+            masks: None
+        """
 
         # NOTE(shikhar): If xs is not provided then the operation is costly,
         # because this function tries to create a tensor of size maxlen x maxlen.
@@ -346,16 +355,16 @@ def forward(
         mask = make_pad_mask(
             lengths=ilens, xs=xs_pad.unsqueeze(-1).unsqueeze(-1), length_dim=1
         ).to(xs_pad.device)
-        # Adjust shapes to be compatible with BEATs code
+        # Adjust shapes to be compatible with Beats code
         xs_pad, mask = xs_pad.squeeze(-1).squeeze(-1), mask.squeeze(-1).squeeze(-1)
         # masks = None
-        last_hidden, mask = self.extract_features(
+        audio_representation, mask = self.extract_features(
             xs_pad,
             mask,
             max_layer=self.max_layer,
         )
         output_lens = (~mask).sum(-1)
-        return last_hidden, output_lens, None
+        return audio_representation, output_lens, None
 
     def extract_features(
         self,
@@ -399,8 +408,8 @@ def extract_features(
         if self.use_weighted_representation:
             repr_layer_weights = nn.functional.softmax(self.layer_weights, dim=-2)
             assert (
-                max_layer > 0
-            ), "max_layer must be non-zero when using weighted representations."
+                max_layer is not None
+            ), "max_layer must not be None when using weighted representations."
             features = (
                 torch.stack(
                     [

From ceed0e92a93120d2b0c339c5f5cb2494a61da837 Mon Sep 17 00:00:00 2001
From: Shikhar Bharadwaj <sbharad2@andrew.cmu.edu>
Date: Fri, 6 Dec 2024 21:34:39 -0500
Subject: [PATCH 41/47] hf decoder tests: overriding config and skipping param
 loading

---
 .../test_hugging_face_transformers_decoder.py | 39 +++++++++++++++++++
 1 file changed, 39 insertions(+)

diff --git a/test/espnet2/asr/decoder/test_hugging_face_transformers_decoder.py b/test/espnet2/asr/decoder/test_hugging_face_transformers_decoder.py
index ccd15bd6d4b..342ae532ce3 100644
--- a/test/espnet2/asr/decoder/test_hugging_face_transformers_decoder.py
+++ b/test/espnet2/asr/decoder/test_hugging_face_transformers_decoder.py
@@ -45,6 +45,45 @@ def test_reload_pretrained_parameters():
     assert torch.equal(saved_param, new_param)
 
 
+@pytest.mark.execution_timeout(30)
+def test_skip_reload_pretrained_parameters():
+    decoder = HuggingFaceTransformersDecoder(
+        5000, 32, "akreal/tiny-random-mbart", load_pretrained_weights=False
+    )
+    # 1. Parameters loaded from hf in init should not be zero.
+    # 2. After zeroing them out they should remain zero post reloading.
+    saved_param = decoder.parameters().__next__().detach().clone()
+    decoder.parameters().__next__().data *= 0
+    zeroed_param = decoder.parameters().__next__().detach().clone()
+    assert not torch.equal(saved_param, zeroed_param)
+
+    decoder.reload_pretrained_parameters()
+    param_after_maybe_reloading = decoder.parameters().__next__().detach().clone()
+    assert torch.equal(
+        param_after_maybe_reloading, zeroed_param
+    )  # Reloading should be skipped
+
+
+@pytest.mark.execution_timeout(30)
+def test_override_hf_decoder_config():
+    overriding_architecture_config = {"d_model": 8, "ignore_mismatched_sizes": True}
+    decoder = HuggingFaceTransformersDecoder(
+        5000,
+        32,
+        "akreal/tiny-random-mbart",
+        overriding_architecture_config=overriding_architecture_config,
+    )
+    assert decoder.decoder.embed_tokens.weight.shape == (5000, 8)
+
+    # Check that without overriding embedding matrix has shape=(5000, 32)
+    decoder = HuggingFaceTransformersDecoder(
+        5000,
+        32,
+        "akreal/tiny-random-mbart",
+    )
+    assert decoder.decoder.embed_tokens.weight.shape == (5000, 16)
+
+
 @pytest.mark.parametrize(
     "model_name_or_path",
     [

From 2acac6567456324f0e19670da470ce0dbb2b7525 Mon Sep 17 00:00:00 2001
From: Shikhar Bharadwaj <sbharad2@andrew.cmu.edu>
Date: Fri, 6 Dec 2024 21:35:22 -0500
Subject: [PATCH 42/47] add test for bets: fwd-bkwd, config override

---
 .../espnet2/asr/encoder/test_beats_encoder.py | 101 ++++++++++++++++++
 1 file changed, 101 insertions(+)
 create mode 100644 test/espnet2/asr/encoder/test_beats_encoder.py

diff --git a/test/espnet2/asr/encoder/test_beats_encoder.py b/test/espnet2/asr/encoder/test_beats_encoder.py
new file mode 100644
index 00000000000..7eabe15d66d
--- /dev/null
+++ b/test/espnet2/asr/encoder/test_beats_encoder.py
@@ -0,0 +1,101 @@
+import pytest
+import torch
+from packaging.version import parse as V
+
+from espnet2.asr.encoder.beats_encoder import BeatsConfig, BeatsEncoder
+
+is_torch_1_12_1_plus = V(torch.__version__) >= V("1.12.1")
+
+
+def test_override_beats_config():
+    if not is_torch_1_12_1_plus:
+        return
+
+    beats_config = BeatsConfig(cfg={"encoder_layers": 2})
+    beats_model = BeatsEncoder(
+        input_size=1,
+        beats_config=beats_config,
+    )
+    assert (
+        len(beats_model.encoder.layers) == 2
+    ), f"Number of layers should be 2. It is {len(beats_model.encoder.layers)}"
+
+
+# Each parameter value creates a variant of the model
+@pytest.mark.timeout(30)
+@pytest.mark.parametrize("downsampling_rate", [1, 2])
+@pytest.mark.parametrize("use_weighted_representation", [False, True])
+@pytest.mark.parametrize(
+    "add_positional_information, max_positions", [(False, None), (True, 12800)]
+)
+def test_forward_pass(
+    downsampling_rate,
+    use_weighted_representation,
+    add_positional_information,
+    max_positions,
+):
+    if not is_torch_1_12_1_plus:
+        return
+
+    beats_config = BeatsConfig(cfg={"encoder_layers": 2})  # Smaller model
+    beats_model = BeatsEncoder(
+        input_size=1,
+        beats_config=beats_config,
+        downsampling_rate=downsampling_rate,
+        use_weighted_representation=use_weighted_representation,
+        add_positional_information=add_positional_information,
+        max_positions=max_positions,
+    )
+    x = torch.randn((2, 32_000, 1))  # B,T,1
+    x_lens = torch.LongTensor([16_000, 24_000])
+    output_rep, output_len, _ = beats_model(x, x_lens)
+
+    # Check output representation
+    assert (
+        output_rep.size(0) == 2
+    ), f"Representation batch size should be 2. It is {output_rep.size(0)}"
+    correct_length = 47 if downsampling_rate == 2 else 96
+    assert (
+        output_rep.size(1) == correct_length
+    ), f"Representation length should be {correct_length}. It is {output_rep.size(1)}"
+    assert output_rep.size(2) == 768, f"Output dim should be 768"
+
+    # Check output length
+    assert (
+        output_len.size(0) == 2
+    ), f"Output length batch size should be 2. It is {output_len.size(0)}"
+    correct_length = (25, 38) if downsampling_rate == 2 else (50, 75)
+    assert (
+        tuple(output_len.tolist()) == correct_length
+    ), f"Output length vector should be {correct_length}. It is {output_len}"
+
+
+@pytest.mark.timeout(30)
+@pytest.mark.parametrize("downsampling_rate", [1, 2])
+@pytest.mark.parametrize("use_weighted_representation", [False, True])
+@pytest.mark.parametrize(
+    "add_positional_information, max_positions", [(False, None), (True, 12800)]
+)
+def test_backward_pass(
+    downsampling_rate,
+    use_weighted_representation,
+    add_positional_information,
+    max_positions,
+):
+    if not is_torch_1_12_1_plus:
+        return
+
+    beats_config = BeatsConfig(cfg={"encoder_layers": 2})  # Smaller model
+    beats_model = BeatsEncoder(
+        input_size=1,
+        beats_config=beats_config,
+        downsampling_rate=downsampling_rate,
+        use_weighted_representation=use_weighted_representation,
+        add_positional_information=add_positional_information,
+        max_positions=max_positions,
+    )
+    x = torch.randn((2, 32_000, 1), requires_grad=True)  # B,T,1
+    x_lens = torch.LongTensor([16_000, 24_000])
+    output_rep, output_len, _ = beats_model(x, x_lens)
+
+    output_rep.sum().backward()

From c76215accad62175afbec1a59991cd4675890be8 Mon Sep 17 00:00:00 2001
From: Shikhar Bharadwaj <sbharad2@andrew.cmu.edu>
Date: Fri, 6 Dec 2024 21:39:39 -0500
Subject: [PATCH 43/47] restore run.sh to default values

---
 egs2/clotho_v2/asr1/run.sh | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/egs2/clotho_v2/asr1/run.sh b/egs2/clotho_v2/asr1/run.sh
index 64b4e3c8080..61c4904e9f7 100755
--- a/egs2/clotho_v2/asr1/run.sh
+++ b/egs2/clotho_v2/asr1/run.sh
@@ -12,7 +12,7 @@ inference_ckpt=valid.acc.ave_5best
     --stop_stage 15 \
     --ngpu 2 \
     --gpu_inference true \
-    --nj 10 \
+    --nj 8 \
     --inference_nj 1 \
     --max_wav_duration 30 \
     --token_type hugging_face \
@@ -35,7 +35,7 @@ inference_ckpt=latest
 ./asr.sh \
     --asr_tag ft \
     --feats_normalize uttmvn \
-    --stage 15 \
+    --stage 1 \
     --stop_stage 15 \
     --asr_speech_fold_length ${asr_speech_fold_length} \
     --asr_stats_dir exp/asr_stats_finetune \

From 41bffd36c03bdb00bad35e0feb398dcdf535965a Mon Sep 17 00:00:00 2001
From: Shikhar Bharadwaj <sbharad2@andrew.cmu.edu>
Date: Fri, 6 Dec 2024 22:03:20 -0500
Subject: [PATCH 44/47] satisfy ci

---
 espnet2/asr/decoder/hugging_face_transformers_decoder.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/espnet2/asr/decoder/hugging_face_transformers_decoder.py b/espnet2/asr/decoder/hugging_face_transformers_decoder.py
index 9e71a5da8ba..ba20e623466 100644
--- a/espnet2/asr/decoder/hugging_face_transformers_decoder.py
+++ b/espnet2/asr/decoder/hugging_face_transformers_decoder.py
@@ -94,7 +94,7 @@ def __init__(
         self.separate_lm_head = separate_lm_head
 
         self.overriding_architecture_config = overriding_architecture_config
-        if type(overriding_architecture_config) == str:
+        if isinstance(overriding_architecture_config, str):
             # It is path to a json config file
             self.overriding_architecture_config = vars(
                 get_model_conf(model_path="", conf_path=overriding_architecture_config)

From d5382f1960845d21d6f82d9fd685f8cb713b41cd Mon Sep 17 00:00:00 2001
From: Yuzhu Wang <yuzhu.wang@tuni.fi>
Date: Sat, 7 Dec 2024 11:57:10 +0200
Subject: [PATCH 45/47] fix: gradient truncation bug in pit_solver.py

---
 espnet2/enh/loss/wrappers/pit_solver.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/espnet2/enh/loss/wrappers/pit_solver.py b/espnet2/enh/loss/wrappers/pit_solver.py
index fe9e11f2a26..d31d3f42307 100644
--- a/espnet2/enh/loss/wrappers/pit_solver.py
+++ b/espnet2/enh/loss/wrappers/pit_solver.py
@@ -98,9 +98,9 @@ def pair_loss(permutation):
                     perm0 = perm_.unsqueeze(1)
                 stats[k] = new_v.gather(1, perm0.to(device=new_v.device)).unbind(1)
         else:
-            loss = torch.tensor(
+            loss = torch.stack(
                 [
-                    torch.tensor(
+                    torch.stack(
                         [
                             pre_hook(
                                 self.criterion,

From 194688354a271932a9b5cfd571cb57f031707f2f Mon Sep 17 00:00:00 2001
From: Shikhar Bharadwaj <sbharad2@andrew.cmu.edu>
Date: Sat, 7 Dec 2024 09:30:40 -0500
Subject: [PATCH 46/47] line len

---
 .../decoder/hugging_face_transformers_decoder.py   | 14 ++++++++------
 1 file changed, 8 insertions(+), 6 deletions(-)

diff --git a/espnet2/asr/decoder/hugging_face_transformers_decoder.py b/espnet2/asr/decoder/hugging_face_transformers_decoder.py
index ba20e623466..cf451bed3b0 100644
--- a/espnet2/asr/decoder/hugging_face_transformers_decoder.py
+++ b/espnet2/asr/decoder/hugging_face_transformers_decoder.py
@@ -64,10 +64,10 @@ def __init__(
                 tokens. Defaults to "".
             postfix (str, optional): Postfix to be added to the input
                 tokens. Defaults to "".
-            overriding_architecture_config (str or dict, optional): Path to the configuration json
-                file or the json dictionary itself. Defaults to None. If this
-                is set, it can be used to override the default decoder
-                configuration.
+            overriding_architecture_config (str or dict, optional): Path to the
+                configuration json file or the json dictionary itself. Defaults
+                to None. If this is set, it can be used to override the
+                default decoder configuration.
             load_pretrained_weights (bool): Whether to load the pre-trained
                 weights. Defaults to True.
             separate_lm_head (bool): True ensures that the language model
@@ -85,8 +85,10 @@ def __init__(
 
         if not is_transformers_available:
             raise ImportError(
-                "`transformers` is not available. Please install it via `pip install"
-                " transformers` or `cd /path/to/espnet/tools && . ./activate_python.sh"
+                "`transformers` is not available. Please install it "
+                "via `pip install"
+                " transformers` or `cd /path/to/espnet/tools "
+                "&& . ./activate_python.sh"
                 " && ./installers/install_transformers.sh`."
             )
 

From ec75a491bef0fd0969e9c1a567a008757847e0d5 Mon Sep 17 00:00:00 2001
From: Shikhar Bharadwaj <sbharad2@andrew.cmu.edu>
Date: Sat, 7 Dec 2024 10:15:03 -0500
Subject: [PATCH 47/47] fix hf asr inference: var renaming

---
 espnet2/bin/asr_inference.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/espnet2/bin/asr_inference.py b/espnet2/bin/asr_inference.py
index ad88258fec3..eea8d8b0d18 100644
--- a/espnet2/bin/asr_inference.py
+++ b/espnet2/bin/asr_inference.py
@@ -252,7 +252,7 @@ def __init__(
 
                 transformer = get_hugging_face_model_network(hugging_face_model)
                 transformer.load_state_dict(decoder.decoder.state_dict())
-                if decoder.ensure_untied_lm_head:
+                if decoder.separate_lm_head:
                     lm_head = copy.deepcopy(
                         get_hugging_face_model_lm_head(hugging_face_model)
                     )
@@ -264,7 +264,7 @@ def __init__(
                     decoder.model_name_or_path, **decoder.overriding_architecture_config
                 )
 
-                if decoder.ensure_untied_lm_head:
+                if decoder.separate_lm_head:
                     hugging_face_model.lm_head = copy.deepcopy(decoder.lm_head)
                 hugging_face_model.lm_head.load_state_dict(decoder.lm_head.state_dict())