push to HF

.gitattributes

@@ -2,13 +2,11 @@
 *.arrow filter=lfs diff=lfs merge=lfs -text
 *.bin filter=lfs diff=lfs merge=lfs -text
 *.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
 *.ftz filter=lfs diff=lfs merge=lfs -text
 *.gz filter=lfs diff=lfs merge=lfs -text
 *.h5 filter=lfs diff=lfs merge=lfs -text
 *.joblib filter=lfs diff=lfs merge=lfs -text
 *.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
 *.model filter=lfs diff=lfs merge=lfs -text
 *.msgpack filter=lfs diff=lfs merge=lfs -text
 *.npy filter=lfs diff=lfs merge=lfs -text
@@ -16,16 +14,14 @@
 *.onnx filter=lfs diff=lfs merge=lfs -text
 *.ot filter=lfs diff=lfs merge=lfs -text
 *.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
 *.pickle filter=lfs diff=lfs merge=lfs -text
 *.pkl filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
 *.pt filter=lfs diff=lfs merge=lfs -text
 *.pth filter=lfs diff=lfs merge=lfs -text
 *.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
 saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
 *.tflite filter=lfs diff=lfs merge=lfs -text
 *.tgz filter=lfs diff=lfs merge=lfs -text
 *.wasm filter=lfs diff=lfs merge=lfs -text
@@ -33,3 +29,6 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+src/epoch=3-step=7459.ckpt filter=lfs diff=lfs merge=lfs -text
+src/wav2vec_small.pt filter=lfs diff=lfs merge=lfs -text
+data/p326_split filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,541 @@
+# generated by: https://github.com/michaelliao/gitignore-online-generator
+
+#################### Python.gitignore ####################
+
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/#use-with-ide
+.pdm.toml
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/
+
+#################### Archives.gitignore ####################
+
+# It's better to unpack these files and commit the raw source because
+# git has its own built in compression methods.
+*.7z
+*.jar
+*.rar
+*.zip
+*.gz
+*.gzip
+*.tgz
+*.bzip
+*.bzip2
+*.bz2
+*.xz
+*.lzma
+*.cab
+*.xar
+
+# Packing-only formats
+*.iso
+*.tar
+
+# Package management formats
+*.dmg
+*.xpi
+*.gem
+*.egg
+*.deb
+*.rpm
+*.msi
+*.msm
+*.msp
+*.txz
+
+#################### Backup.gitignore ####################
+
+*.bak
+*.gho
+*.ori
+*.orig
+*.tmp
+
+#################### Emacs.gitignore ####################
+
+# -*- mode: gitignore; -*-
+*~
+\#*\#
+/.emacs.desktop
+/.emacs.desktop.lock
+*.elc
+auto-save-list
+tramp
+.\#*
+
+# Org-mode
+.org-id-locations
+*_archive
+
+# flymake-mode
+*_flymake.*
+
+# eshell files
+/eshell/history
+/eshell/lastdir
+
+# elpa packages
+/elpa/
+
+# reftex files
+*.rel
+
+# AUCTeX auto folder
+/auto/
+
+# cask packages
+.cask/
+dist/
+
+# Flycheck
+flycheck_*.el
+
+# server auth directory
+/server/
+
+# projectiles files
+.projectile
+
+# directory configuration
+.dir-locals.el
+
+# network security
+/network-security.data
+
+
+#################### JetBrains.gitignore ####################
+
+# Covers JetBrains IDEs: IntelliJ, RubyMine, PhpStorm, AppCode, PyCharm, CLion, Android Studio, WebStorm and Rider
+# Reference: https://intellij-support.jetbrains.com/hc/en-us/articles/206544839
+
+# User-specific stuff
+.idea/**/workspace.xml
+.idea/**/tasks.xml
+.idea/**/usage.statistics.xml
+.idea/**/dictionaries
+.idea/**/shelf
+
+# AWS User-specific
+.idea/**/aws.xml
+
+# Generated files
+.idea/**/contentModel.xml
+
+# Sensitive or high-churn files
+.idea/**/dataSources/
+.idea/**/dataSources.ids
+.idea/**/dataSources.local.xml
+.idea/**/sqlDataSources.xml
+.idea/**/dynamic.xml
+.idea/**/uiDesigner.xml
+.idea/**/dbnavigator.xml
+
+# Gradle
+.idea/**/gradle.xml
+.idea/**/libraries
+
+# Gradle and Maven with auto-import
+# When using Gradle or Maven with auto-import, you should exclude module files,
+# since they will be recreated, and may cause churn.  Uncomment if using
+# auto-import.
+# .idea/artifacts
+# .idea/compiler.xml
+# .idea/jarRepositories.xml
+# .idea/modules.xml
+# .idea/*.iml
+# .idea/modules
+# *.iml
+# *.ipr
+
+# CMake
+cmake-build-*/
+
+# Mongo Explorer plugin
+.idea/**/mongoSettings.xml
+
+# File-based project format
+*.iws
+
+# IntelliJ
+out/
+
+# mpeltonen/sbt-idea plugin
+.idea_modules/
+
+# JIRA plugin
+atlassian-ide-plugin.xml
+
+# Cursive Clojure plugin
+.idea/replstate.xml
+
+# SonarLint plugin
+.idea/sonarlint/
+
+# Crashlytics plugin (for Android Studio and IntelliJ)
+com_crashlytics_export_strings.xml
+crashlytics.properties
+crashlytics-build.properties
+fabric.properties
+
+# Editor-based Rest Client
+.idea/httpRequests
+
+# Android studio 3.1+ serialized cache file
+.idea/caches/build_file_checksums.ser
+
+#################### Linux.gitignore ####################
+
+*~
+
+# temporary files which can be created if a process still has a handle open of a deleted file
+.fuse_hidden*
+
+# KDE directory preferences
+.directory
+
+# Linux trash folder which might appear on any partition or disk
+.Trash-*
+
+# .nfs files are created when an open file is removed but is still being accessed
+.nfs*
+
+#################### NotepadPP.gitignore ####################
+
+# Notepad++ backups #
+*.bak
+
+#################### PuTTY.gitignore ####################
+
+# Private key
+*.ppk
+
+#################### SublimeText.gitignore ####################
+
+# Cache files for Sublime Text
+*.tmlanguage.cache
+*.tmPreferences.cache
+*.stTheme.cache
+
+# Workspace files are user-specific
+*.sublime-workspace
+
+# Project files should be checked into the repository, unless a significant
+# proportion of contributors will probably not be using Sublime Text
+# *.sublime-project
+
+# SFTP configuration file
+sftp-config.json
+sftp-config-alt*.json
+
+# Package control specific files
+Package Control.last-run
+Package Control.ca-list
+Package Control.ca-bundle
+Package Control.system-ca-bundle
+Package Control.cache/
+Package Control.ca-certs/
+Package Control.merged-ca-bundle
+Package Control.user-ca-bundle
+oscrypto-ca-bundle.crt
+bh_unicode_properties.cache
+
+# Sublime-github package stores a github token in this file
+# https://packagecontrol.io/packages/sublime-github
+GitHub.sublime-settings
+
+#################### Vim.gitignore ####################
+
+# Swap
+[._]*.s[a-v][a-z]
+!*.svg  # comment out if you don't need vector files
+[._]*.sw[a-p]
+[._]s[a-rt-v][a-z]
+[._]ss[a-gi-z]
+[._]sw[a-p]
+
+# Session
+Session.vim
+Sessionx.vim
+
+# Temporary
+.netrwhist
+*~
+# Auto-generated tag files
+tags
+# Persistent undo
+[._]*.un~
+
+#################### VirtualEnv.gitignore ####################
+
+# Virtualenv
+# http://iamzed.com/2009/05/07/a-primer-on-virtualenv/
+.Python
+[Bb]in
+[Ii]nclude
+[Ll]ib
+[Ll]ib64
+[Ll]ocal
+[Ss]cripts
+pyvenv.cfg
+.venv
+pip-selfcheck.json
+
+#################### VisualStudioCode.gitignore ####################
+
+.vscode/*
+!.vscode/settings.json
+!.vscode/tasks.json
+!.vscode/launch.json
+!.vscode/extensions.json
+!.vscode/*.code-snippets
+
+# Local History for Visual Studio Code
+.history/
+
+# Built Visual Studio Code Extensions
+*.vsix
+
+#################### Windows.gitignore ####################
+
+# Windows thumbnail cache files
+Thumbs.db
+Thumbs.db:encryptable
+ehthumbs.db
+ehthumbs_vista.db
+
+# Dump file
+*.stackdump
+
+# Folder config file
+[Dd]esktop.ini
+
+# Recycle Bin used on file shares
+$RECYCLE.BIN/
+
+# Windows Installer files
+*.cab
+*.msi
+*.msix
+*.msm
+*.msp
+
+# Windows shortcuts
+*.lnk
+
+#################### macOS.gitignore ####################
+
+# General
+.DS_Store
+.AppleDouble
+.LSOverride
+
+# Icon must end with two \r
+Icon
+
+
+# Thumbnails
+._*
+
+# Files that might appear in the root of a volume
+.DocumentRevisions-V100
+.fseventsd
+.Spotlight-V100
+.TemporaryItems
+.Trashes
+.VolumeIcon.icns
+.com.apple.timemachine.donotpresent
+
+# Directories potentially created on remote AFP share
+.AppleDB
+.AppleDesktop
+Network Trash Folder
+Temporary Items
+.apdisk
+
+#################### Custom.gitignore ####################
+
+# add your custom gitignore here:
+!.gitignore
+!.gitsubmodules
+
+# ignore data
+data/
+exp/
+!src/lightning_module.py
+*.wav
+# ignore plots
+*.png
+# ignore csv
+*.csv
+
+
+!config/template.yaml
+config
+!local
+
+## Currently
+src/wav2vec_small.pt
+*.ckpt
+*.bak
+
+fine_tuned/
+
+# 
+user/
+
+.vscode

README.md

@@ -1,12 +1,127 @@
+# Laronix Data Collection
+
+This repository contains information about the Laronix data collection process, which involves collecting parallel data from AVA users. The dataset consists of two main sessions: scripted data and conversational data.
+
+## Dataset
+
+The dataset is organized as follows:
+
+### 1. Scripted Data
+
+The scripted data session includes 200 sentences collected from 5 articles. The references for both the audio and text versions of these sentences have already been uploaded or will be uploaded to the Laronix Recording system. (Ask [Kevin](kevin@laronix.com) for these files) The distribution of sentences from each article is as follows:
+
+- Arthur the Rat: 56 sentences
+- Cinder: 19 sentences
+- Rainbow: 26 sentences
+- Sentences: 59 sentences
+- VCTK: 40 sentences
+
+### 2. Conversational Data
+
+The conversational data session focuses on natural conversations and involves the following components:
+
+#### a. Q&A
+
+In this component, a set of 50 sentences will be provided, consisting of questions and answers. During the recording, the partner will ask the questions (Q), and the patient will provide the answers (A). Both the questions and answers will be recorded.
+
+#### b. Freestyle
+
+The patients will have the freedom to talk about a given topic. They will be asked to respond with 5 to 10 sentences. The structure for this component can be referenced from the [IELTS speaking test](https://www.ieltsbuddy.com/IELTS-speaking-questions-with-answers.html).
+
+## Data Inclusion Criteria
+
++ No hearing loss or history of active cancer.
++ 6 weeks of practice with AVA.
+
+## Document for Laronix Recording System
+
+The Laronix recording system is designed for data collection from potential users of the AVA Device, which replaces their voice cord.
+
+### Input:
+
+- Audio signal
+- Reference ID
+- Reference text
+- Reference Phoneme per minute
+
+### Output:
+
+- wav_pause_plot: Wave signal plot with pauses detected by VAD algorithm (SNR = 40dB)
+- Predicted Mean Opinion Score: Score estimating data quality on the MOS scale using an ML prediction model (1-5)
+- Hypotheses: Text predicted by Automatic Speech Recognition model (wav2vev2.0 + CTC)
+- WER: Word Error Rate (lower is better)
+- Predicted Phonemes
+- PPM: Phonemes per minute
+- Message: Feedback from the system
+
+## User Instruction
+
+Please follow the instructions provided at the top of the APP page.
+
+```
+- Laronix_AUTOMOS
+    - data
+        - Template
+            - ref_wav/
+                - 1.wav
+                - 2.wav
+                - ...
+            - ref_txt.txt
+            - ref.csv # audio prosody features reference <generate by script>
+    - exp
+        - Template
+            - Audio_to_evaluate # RAW WAV DATA
+            - log.csv # Recording log 
+            - output # wav.file  <generate by script>
+    - model
+        - epoch=3-step=7459.ckpt # MOS estimate model
+        - wav2vec_small.pt # WER model
+    - local
+        - get_ref_PPM.py # script for generating data/<ref_dir>/ref.csv
+        - post_processing.py # script for generating exp/<ref_dir>/output/*.wav
+```
+
 ---
-title: Laronix Recording
-emoji: 🐠
-colorFrom: gray
+title: Laronix Automos
+emoji: 🏃
+colorFrom: blue
 colorTo: blue
 sdk: gradio
-sdk_version: 3.47.1
+sdk_version: 3.2
 app_file: app.py
 pinned: false
+license: afl-3.0
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+
+# Laronix_AutoMOS
+
+## Usage:
+### Step 1: Prepare data and text
+`<todo>`
+### Step 2: Preprocessing
+```
+## Generating *.csv, Voice/Unvoice Plot (optional) and config (optional)
+python local/get_ref_PPM.py --ref_txt <ref_text> \
+                            --ref_wavs <ref_wavs> \
+                            --output_dir <output_dir> \
+                            --to_config <True/False> \
+                            --UV_flag <True/False> \
+                            --UV_thre <UV_thre>}
+```
+### Step 3: Launch recording session:
+
+```
+## Start app.py
+python app.py <config.yaml>
+```
++ **Find logging below and lick URL to start**
+```
+Launch examples
+Running on local URL:  http://127.0.0.1:7860/
+...
+    (Logs...)
+...
+Running on public URL: https://87abe771e93229da.gradio.app
+```

app.py

@@ -0,0 +1,468 @@
+"""
+TODO:
+    + [x] Load Configuration
+    + [ ] Checking
+    + [ ] Better saving directory
+"""
+import numpy as np
+from pathlib import Path
+import jiwer
+import pdb
+import torch.nn as nn
+import torch
+import torchaudio
+import gradio as gr
+from logging import PlaceHolder
+from transformers import Wav2Vec2Processor, Wav2Vec2ForCTC
+import yaml
+from transformers import pipeline
+import librosa
+import librosa.display
+import matplotlib.pyplot as plt
+
+
+# local import
+import sys
+
+sys.path.append("src")
+import lightning_module
+
+# Load automos
+config_yaml = sys.argv[1]
+with open(config_yaml, "r") as f:
+    # pdb.set_trace()
+    try:
+        config = yaml.safe_load(f)
+    except FileExistsError:
+        print("Config file Loading Error")
+        exit()
+
+# Auto load examples
+
+with open(config["ref_txt"], "r") as f:
+    refs = f.readlines()
+refs_ids = [x.split()[0] for x in refs]
+refs_txt = [" ".join(x.split()[1:]) for x in refs]
+ref_feature = np.loadtxt(config["ref_feature"], delimiter=",", dtype="str")
+ref_wavs = [str(x) for x in sorted(Path(config["ref_wavs"]).glob("**/*.wav"))]
+
+dummy_wavs = [None for x in np.arange(len(ref_wavs))]
+
+refs_ppm = np.array(ref_feature[:, -1][1:], dtype="str")
+
+reference_id = gr.Textbox(value="ID", placeholder="Utter ID", label="Reference_ID")
+
+reference_textbox = gr.Textbox(
+    value="Input reference here",
+    placeholder="Input reference here",
+    label="Reference",
+)
+reference_PPM = gr.Textbox(placeholder="Pneumatic Voice's PPM", label="Ref PPM")
+
+# Set up interface
+# remove dummpy wavs, ue the same ref_wavs for eval wavs
+print("Preparing Examples")
+examples = [
+    [w, w_, i, x, y] for w, w_, i, x, y in zip(ref_wavs, ref_wavs, refs_ids, refs_txt, refs_ppm)
+]
+
+p = pipeline(
+    "automatic-speech-recognition",
+    model="KevinGeng/whipser_medium_en_PAL300_step25",
+    device=0,
+)
+
+# WER part
+transformation = jiwer.Compose(
+    [
+        jiwer.RemovePunctuation(),
+        jiwer.ToLowerCase(),
+        jiwer.RemoveWhiteSpace(replace_by_space=True),
+        jiwer.RemoveMultipleSpaces(),
+        jiwer.ReduceToListOfListOfWords(word_delimiter=" "),
+    ]
+)
+
+# WPM part
+processor = Wav2Vec2Processor.from_pretrained("facebook/wav2vec2-xlsr-53-espeak-cv-ft")
+phoneme_model = Wav2Vec2ForCTC.from_pretrained("facebook/wav2vec2-xlsr-53-espeak-cv-ft")
+
+
+class ChangeSampleRate(nn.Module):
+    def __init__(self, input_rate: int, output_rate: int):
+        super().__init__()
+        self.output_rate = output_rate
+        self.input_rate = input_rate
+
+    def forward(self, wav: torch.tensor) -> torch.tensor:
+        # Only accepts 1-channel waveform input
+        wav = wav.view(wav.size(0), -1)
+        new_length = wav.size(-1) * self.output_rate // self.input_rate
+        indices = torch.arange(new_length) * (self.input_rate / self.output_rate)
+        round_down = wav[:, indices.long()]
+        round_up = wav[:, (indices.long() + 1).clamp(max=wav.size(-1) - 1)]
+        output = round_down * (1.0 - indices.fmod(1.0)).unsqueeze(0) + (
+            round_up * indices.fmod(1.0).unsqueeze(0)
+        )
+        return output
+
+
+# MOS model
+model = lightning_module.BaselineLightningModule.load_from_checkpoint(
+    "src/epoch=3-step=7459.ckpt"
+).eval()
+
+# Get Speech Interval
+
+def get_speech_interval(signal, db):
+    audio_interv = librosa.effects.split(signal, top_db=db)
+    pause_end = [x[0] for x in audio_interv[1:]]
+    pause_start = [x[1] for x in audio_interv[0:-1]]
+    pause_interv = [[x, y] for x, y in zip(pause_start, pause_end)]
+    return audio_interv, pause_interv
+
+# plot UV
+
+
+def plot_UV(signal, audio_interv, sr):
+    fig, ax = plt.subplots(nrows=2, sharex=True)
+    librosa.display.waveshow(signal, sr=sr, ax=ax[0])
+    uv_flag = np.zeros(len(signal))
+    for i in audio_interv:
+        uv_flag[i[0] : i[1]] = 1
+
+    ax[1].plot(np.arange(len(signal)) / sr, uv_flag, "r")
+    ax[1].set_ylim([-0.1, 1.1])
+    return fig
+
+def calc_mos(_, audio_path, id, ref, pre_ppm, fig=None):
+    if audio_path == None:
+        audio_path = _
+        print("using ref audio as eval audio since it's empty")
+
+    wav, sr = torchaudio.load(audio_path)
+    if wav.shape[0] != 1:
+        wav = wav[0, :]
+    print(wav.shape)
+
+    osr = 16000
+    batch = wav.unsqueeze(0).repeat(10, 1, 1)
+    csr = ChangeSampleRate(sr, osr)
+    out_wavs = csr(wav)
+
+    # ASR
+    trans = jiwer.ToLowerCase()(p(audio_path)["text"])
+
+    # WER
+    wer = jiwer.wer(
+        ref,
+        trans,
+        truth_transform=transformation,
+        hypothesis_transform=transformation,
+    )
+    # MOS
+    batch = {
+        "wav": out_wavs,
+        "domains": torch.tensor([0]),
+        "judge_id": torch.tensor([288]),
+    }
+    with torch.no_grad():
+        output = model(batch)
+    predic_mos = output.mean(dim=1).squeeze().detach().numpy() * 2 + 3
+
+    # Phonemes per minute (PPM)
+    with torch.no_grad():
+        logits = phoneme_model(out_wavs).logits
+    phone_predicted_ids = torch.argmax(logits, dim=-1)
+    phone_transcription = processor.batch_decode(phone_predicted_ids)
+    lst_phonemes = phone_transcription[0].split(" ")
+
+    # VAD for pause detection
+    wav_vad = torchaudio.functional.vad(wav, sample_rate=sr)
+    # pdb.set_trace()
+    a_h, p_h = get_speech_interval(wav_vad.numpy(), db=40)
+    # print(a_h)
+    # print(len(a_h))
+    fig_h = plot_UV(wav_vad.numpy().squeeze(), a_h, sr=sr)
+    ppm = len(lst_phonemes) / (wav_vad.shape[-1] / sr) * 60
+
+    error_msg = "!!! ERROR MESSAGE !!!\n"
+    if audio_path == _ or audio_path == None:
+        error_msg += "ERROR: Fail recording, Please start from the beginning again."
+        return (
+            fig_h,
+            predic_mos,
+            trans,
+            wer,
+            phone_transcription,
+            ppm,
+            error_msg,
+        )
+    if ppm >= float(pre_ppm) + float(config["thre"]["maxppm"]):
+        error_msg += "ERROR: Please speak slower.\n"
+    elif ppm <= float(pre_ppm) - float(config["thre"]["minppm"]):
+        error_msg += "ERROR: Please speak faster.\n"
+    elif predic_mos <= float(config["thre"]["AUTOMOS"]):
+        error_msg += "ERROR: Naturalness is too low, Please try again.\n"
+    elif wer >= float(config["thre"]["WER"]):
+        error_msg += "ERROR: Intelligibility is too low, Please try again\n"
+    else:
+        error_msg = (
+            "GOOD JOB! Please 【Save the Recording】.\nYou can start recording the next sample."
+        )
+
+    return (
+        fig_h,
+        predic_mos,
+        trans,
+        wer,
+        phone_transcription,
+        ppm,
+        error_msg,
+    )
+
+with open("src/description.html", "r", encoding="utf-8") as f:
+    description = f.read()
+# description
+
+refs_ppm = np.array(ref_feature[:, -1][1:], dtype="str")
+
+reference_id = gr.Textbox(value="ID", placeholder="Utter ID", label="Reference_ID", visible=False)
+reference_textbox = gr.Textbox(
+    value="Input reference here",
+    placeholder="Input reference here",
+    label="Reference",
+)
+reference_PPM = gr.Textbox(placeholder="Pneumatic Voice's PPM", label="Ref PPM", visible=False)
+
+# Flagging setup
+
+# Interface
+# Participant Information
+def record_part_info(name, gender, first_lng):
+    message = "Participant information is successfully collected."
+    id_str = "%s_%s_%s" % (name, gender[0], first_lng[0])
+
+    if name == None:
+        message = "ERROR: Name Information incomplete!"
+        id_str = "ERROR"
+
+    if gender == None:
+        message = "ERROR: Please select gender"
+        id_str = "ERROR"
+
+    if len(gender) > 1:
+        message = "ERROR: Please select one gender only"
+        id_str = "ERROR"
+
+    if first_lng == None:
+        message = "ERROR: Please select your english proficiency"
+        id_str = "ERROR"
+
+    if len(first_lng) > 1:
+        message = "ERROR: Please select one english proficiency only"
+        id_str = "ERROR"
+
+    return message, id_str
+
+
+# information page not using now
+name = gr.Textbox(placeholder="Name", label="Name")
+gender = gr.CheckboxGroup(["Male", "Female"], label="gender")
+first_lng = gr.CheckboxGroup(
+    [
+        "B1 Intermediate",
+        "B2: Upper Intermediate",
+        "C1: Advanced",
+        "C2: Proficient",
+    ],
+    label="English Proficiency (CEFR)",
+)
+
+msg = gr.Textbox(placeholder="Evaluation for valid participant", label="message")
+id_str = gr.Textbox(placeholder="participant id", label="participant_id")
+
+info = gr.Interface(
+    fn=record_part_info,
+    inputs=[name, gender, first_lng],
+    outputs=[msg, id_str],
+    title="Participant Information Page",
+    allow_flagging="never",
+    css="body {background-color: blue}",
+)
+# Experiment
+if config["exp_id"] == None:
+    config["exp_id"] = Path(config_yaml).stem
+
+## This is the theme for the interface
+css = """
+.ref_text textarea {font-size: 40px !important}
+.message textarea {font-size: 40px !important}
+"""
+
+my_theme = gr.themes.Default().set(
+    button_primary_background_fill="#75DA99",
+    button_primary_background_fill_dark="#DEF2D7",
+    button_primary_text_color="black",
+    button_secondary_text_color="black",
+)
+
+# Callback for saving the recording
+callback = gr.CSVLogger()
+
+with gr.Blocks(css=css, theme=my_theme) as demo:
+    with gr.Column():
+        with gr.Row():
+            ref_audio = gr.Audio(
+                source="microphone",
+                type="filepath",
+                label="Reference_Audio",
+                container=True,
+                interactive=False,
+                visible=False,
+            )
+            with gr.Row():
+                eval_audio = gr.Audio(
+                    source="microphone",
+                    type="filepath",
+                    container=True,
+                    label="Audio_to_Evaluate",
+                )
+                b_redo = gr.ClearButton(
+                    value="Redo", variant="stop", components=[eval_audio], size="sm"
+                )
+                reference_textbox = gr.Textbox(
+                    value="Input reference here",
+                    placeholder="Input reference here",
+                    label="Reference",
+                    interactive=True,
+                    elem_classes="ref_text",
+                )
+                with gr.Accordion("Input for Development", open=False):
+                    reference_id = gr.Textbox(
+                        value="ID",
+                        placeholder="Utter ID",
+                        label="Reference_ID",
+                        visible=True,
+                    )
+                    reference_PPM = gr.Textbox(
+                        placeholder="Pneumatic Voice's PPM",
+                        label="Ref PPM",
+                        visible=True,
+                    )
+        with gr.Row():
+            b = gr.Button(value="1.Submit", variant="primary", elem_classes="submit")
+
+            # TODO
+            # b_more = gr.Button(value="Show More", elem_classes="verbose")
+        with gr.Row():
+            inputs = [
+                ref_audio,
+                eval_audio,
+                reference_id,
+                reference_textbox,
+                reference_PPM,
+            ]
+            e = gr.Examples(examples, inputs, examples_per_page=5)
+
+    with gr.Column():
+        with gr.Row():
+            ## output block
+            msg = gr.Textbox(
+                placeholder="Recording Feedback",
+                label="Message",
+                interactive=False,
+                elem_classes="message",
+            )
+            with gr.Accordion("Output for Development", open=False):
+                wav_plot = gr.Plot(PlaceHolder="Wav/Pause Plot", label="wav_pause_plot", visible=True)
+
+                predict_mos = gr.Textbox(
+                    placeholder="Predicted MOS",
+                    label="Predicted MOS",
+                    visible=True,
+                )
+
+                hyp = gr.Textbox(placeholder="Hypothesis", label="Hypothesis", visible=True)
+
+                wer = gr.Textbox(placeholder="Word Error Rate", label="WER", visible=True)
+
+                predict_pho = gr.Textbox(
+                    placeholder="Predicted Phonemes",
+                    label="Predicted Phonemes",
+                    visible=True,
+                )
+
+                ppm = gr.Textbox(
+                    placeholder="Phonemes per minutes",
+                    label="PPM",
+                    visible=True,
+                )
+            outputs = [
+                wav_plot,
+                predict_mos,
+                hyp,
+                wer,
+                predict_pho,
+                ppm,
+                msg,
+            ]
+
+            # b = gr.Button("Submit")
+            b.click(fn=calc_mos, inputs=inputs, outputs=outputs, api_name="Submit")
+
+        # Logger
+        callback.setup(
+            components=[
+                eval_audio,
+                reference_id,
+                reference_textbox,
+                reference_PPM,
+                predict_mos,
+                hyp,
+                wer,
+                ppm,
+                msg],
+            flagging_dir="./exp/%s" % config["exp_id"],
+        )
+
+        with gr.Row():
+            b2 = gr.Button("2. Save the Recording", variant="primary", elem_id="save")
+            js_confirmed_saving = "(x) => confirm('Recording Saved!')"
+            # eval_audio,
+            b2.click(
+                lambda *args: callback.flag(args),
+                inputs=[
+                    eval_audio,
+                    reference_id,
+                    reference_textbox,
+                    reference_PPM,
+                    predict_mos,
+                    hyp,
+                    wer,
+                    ppm,
+                    msg,
+                ],
+                outputs=None,
+                preprocess=False,
+                api_name="flagging",
+            )
+        with gr.Row():
+            b3 = gr.ClearButton(
+                [
+                    ref_audio,
+                    eval_audio,
+                    reference_id,
+                    reference_textbox,
+                    reference_PPM,
+                    predict_mos,
+                    hyp,
+                    wer,
+                    ppm,
+                    msg,
+                ],
+                value="3.Clear All",
+                elem_id="clear",
+            )
+
+demo.launch(share=True)

local/UV.py

@@ -0,0 +1,29 @@
+import numpy as np
+import librosa.display
+import matplotlib.pyplot as plt
+
+# Plot_UV
+
+
+def plot_UV(signal, audio_interv, sr):
+    fig, ax = plt.subplots(nrows=2, sharex=True)
+    librosa.display.waveshow(signal, sr=sr, ax=ax[0])
+    ax[0].set_title("Signal")
+    ax[1].set_title("U/V")
+    uv_flag = np.zeros(len(signal))
+    for i in audio_interv:
+        uv_flag[i[0]: i[1]] = 1
+
+    ax[1].plot(np.arange(len(signal))/sr, uv_flag, "r")
+    ax[1].set_ylim([-0.1,  1.1])
+    return fig
+
+# Get Speech Interval
+
+
+def get_speech_interval(signal, db):
+    audio_interv = librosa.effects.split(signal, top_db=db)
+    pause_end = [x[0] for x in audio_interv[1:]]
+    pause_start = [x[1] for x in audio_interv[0: -1]]
+    pause_interv = [[x, y] for x, y in zip(pause_start, pause_end)]
+    return audio_interv, pause_interv

local/VCTK_preprocessing.py

@@ -0,0 +1,62 @@
+# Kevin @ Laronix Dec. 2022
+# Data processing at Laronix
+import csv
+import soundfile as sf
+import pandas as pd
+from pathlib import Path
+import librosa
+import sys
+import numpy as np
+import pdb
+from rich.progress import track
+
+wavdir = sys.argv[1]
+txtdir = sys.argv[2]
+thre_len = int(sys.argv[3])
+origin_sr = int(sys.argv[4])
+target_sr = int(sys.argv[5])
+
+wavs = sorted(Path(wavdir).glob("**/*.wav"))
+txts = sorted(Path(txtdir).glob("**/*.txt"))
+target_dir = "./data/%s_%d_%d_len%d" % (
+    Path(wavdir).stem,
+    origin_sr,
+    target_sr,
+    thre_len,
+)
+
+Path.mkdir(Path(target_dir), exist_ok=True)
+# pdb.set_trace()
+tables = []
+for x, y in track(
+    zip(wavs, txts), description="Processing...", total=len(wavs)
+):
+    label = 1
+    with open(y, "r") as f:
+        txt = f.readline()
+    if len(txt.split(" ")) <= thre_len:
+        label = 1
+        record = [x, Path(x).stem, txt, len(txt.split(" ")), label]
+        tables.append(record)
+    # Select length <= 10 words sentences for training
+    if len(txt.split(" ")) <= thre_len:
+        wav, sr = librosa.load(x, sr=origin_sr)
+        wav_ = librosa.resample(wav, orig_sr=sr, target_sr=target_sr)
+        sf.write(
+            Path(target_dir) / Path((x).stem + ".wav"),
+            data=wav_,
+            samplerate=target_sr,
+        )
+
+D = pd.DataFrame(
+    tables, columns=["wav_path", "id", "text", "len", "length_label"]
+)
+D.to_csv(target_dir + ".datalog", sep=",")
+print("Check data log at %s" % (target_dir + ".datalog"))
+
+D.get(["id", "text"]).to_csv(
+    target_dir + ".txt", sep="\t", header=False, index=False, quoting=3
+)
+
+print("Generate id_text at %s" % (target_dir + ".txt"))
+print("Finish")

local/data_preparation.py

@@ -0,0 +1,78 @@
+import os
+import pdb
+import shutil
+import pandas as pd
+from datasets import Dataset, load_dataset
+
+audio_dir = "./data/Patient_sil_trim_16k_normed_5_snr_40/"
+# split_files = {"train": "data/Patient_sil_trim_16k_normed_5_snr_40/train.csv", 
+#                "test": "data/Patient_sil_trim_16k_normed_5_snr_40/test.csv",
+#                "dev": "data/Patient_sil_trim_16k_normed_5_snr_40/dev.csv"}
+src_dataset = load_dataset("audiofolder", data_dir=audio_dir, split="train")
+pdb.set_trace()
+def train_dev_test_split(
+    dataset: Dataset, dev_rate=0.1, test_rate=0.1, seed=1, metadata_output=False, root_dir=None
+):
+    """
+    input: dataset
+    dev_rate,
+    test_rate
+    seed
+    -------
+    Output:
+    dataset_dict{"train", "dev", "test"}
+    """
+    train_dev_test = dataset.train_test_split(test_size=test_rate, seed=seed)
+    test = train_dev_test["test"]
+    train_dev = train_dev_test["train"]
+
+    if len(train_dev) <= int(len(dataset) * dev_rate):
+        train = Dataset.from_dict({"audio": [], "transcription": []})
+        dev = train_dev
+    else:
+        train_dev = train_dev.train_test_split(
+            test_size=int(len(dataset) * dev_rate), seed=seed
+        )
+        train = train_dev["train"]
+        dev = train_dev["test"]
+    
+    train_size = len(train)
+    dev_size = len(dev)
+    test_size = len(test)
+    
+    print(f"Train Size: {len(train)}")
+    print(f"Dev Size: {len(dev)}")
+    print(f"Test Size: {len(test)}")
+    import pdb
+    if metadata_output:
+        pdb.set_trace()
+        train_df = pd.DateFrame(train)
+        dev_df = pd.DataFrame(dev)
+        test_df = pd.DataFrame(test)
+
+    try:
+         os.path.exists(root_dir)
+    except:
+        raise FileNotFoundError
+    
+    # Create directories for train, dev, and test data
+    import pdb
+    if not os.path.exists(f'{root_dir}/train'):
+        os.makedirs(f'{root_dir}/train')
+    if not os.path.exists(f'{root_dir}/dev'):
+        os.makedirs(f'{root_dir}/dev')
+    if not os.path.exists(f'{root_dir}/test'):
+        os.makedirs(f'{root_dir}/test')
+
+    pdb.set_trace()
+    train_df.to_csv(f'{root_dir}/train/metadata.csv', index=False)
+
+    dev_df.to_csv(f'{root_dir}/dev/metadata.csv', index=False)
+
+    test_df.to_csv(f'{root_dir}/test/metadata.csv', index=False)
+
+    return train, dev, test
+
+train, dev, test = train_dev_test_split(src_dataset, dev_rate=0.1, test_rate=0.1, seed=1, metadata_output=True, root_dir=audio_dir)
+
+pdb.set_trace()

local/decode.py

@@ -0,0 +1,427 @@
+fine_tuning_dir = "fine_tuned/SSD/model/Negel_79_AVA_script_conv_train_conv_dev/checkpoint-50"
+
+from typing import Any, Dict, List, Union
+from dataclasses import dataclass
+from transformers import Seq2SeqTrainer
+from transformers import WhisperProcessor, WhisperForConditionalGeneration, WhisperTokenizer, WhisperFeatureExtractor, Seq2SeqTrainingArguments, Seq2SeqTrainer, WhisperModel
+import evaluate
+from transformers import Wav2Vec2Processor, Wav2Vec2ForCTC
+from random import sample
+from sys import flags
+import gradio as gr
+import torchaudio
+import torch.nn as nn
+import jiwer
+import numpy as np
+from rich import print as rprint
+from rich.progress import track
+from transformers import pipeline
+import argparse
+import yaml
+import torch
+from pathlib import Path
+from transformers import AutoTokenizer, AutoFeatureExtractor, AutoModelForCTC, AutoProcessor
+from datasets import load_dataset, concatenate_datasets
+from datasets import Dataset, Audio
+import pdb
+import string
+import librosa
+# local import
+import sys
+
+sys.path.append("src")
+import lightning_module
+
+torch.cuda.set_device("cuda:0")
+
+audio_dir = "./data/Patient_sil_trim_16k_normed_5_snr_40"
+healthy_dir = "./data/Healthy"
+Fary_PAL_30 = "./data/Fary_PAL_p326_20230110_30"
+John_p326 = "./data/John_p326/output"
+John_video = "./data/20230103_video"
+negel_79 = "./data/4_negel_79"
+
+patient_T = "data/Patient_T/Patient_T"
+patient_L = "data/Patient_L/Patient_L"
+# Get Transcription, WER and PPM
+"""
+TODO:
+    [DONE]: Automatic generating Config
+"""
+
+
+sys.path.append("./src")
+
+
+wer = evaluate.load("wer")
+
+# root_path = Path(__file__).parents[1]
+
+
+class ChangeSampleRate(nn.Module):
+    def __init__(self, input_rate: int, output_rate: int):
+        super().__init__()
+        self.output_rate = output_rate
+        self.input_rate = input_rate
+
+    def forward(self, wav: torch.tensor) -> torch.tensor:
+        # Only accepts 1-channel waveform input
+        wav = wav.view(wav.size(0), -1)
+        new_length = wav.size(-1) * self.output_rate // self.input_rate
+        indices = torch.arange(new_length) * (
+            self.input_rate / self.output_rate
+        )
+        round_down = wav[:, indices.long()]
+        round_up = wav[:, (indices.long() + 1).clamp(max=wav.size(-1) - 1)]
+        output = round_down * (1.0 - indices.fmod(1.0)).unsqueeze(
+            0
+        ) + round_up * indices.fmod(1.0).unsqueeze(0)
+        return output
+
+# resample and clean text data
+
+
+def dataclean(example):
+    # pdb.set_trace()
+    if example['audio']['sampling_rate'] != 16000:
+        resampled_audio = librosa.resample(y=example['audio']['array'],
+                                           orig_sr=example['audio']['sampling_rate'],
+                                           target_sr=16000)
+        # torchaudio.transforms.Resample(example['audio']['sampling_rate'], 16000)
+        # resampled_audio = resampler(example['audio']['array'])
+
+        return {"audio": {"path": example['audio']['path'], "array": resampled_audio, "sampling_rate": 16000},
+                "transcription": example["transcription"].upper().translate(str.maketrans('', '', string.punctuation))}
+    else:
+        return {"transcription": example["transcription"].upper().translate(str.maketrans('', '', string.punctuation))}
+
+processor = AutoFeatureExtractor.from_pretrained(
+    "facebook/wav2vec2-base-960h"
+)
+
+def prepare_dataset(batch):
+    audio = batch["audio"]
+    batch = processor(
+        audio["array"], sampling_rate=audio["sampling_rate"], text=batch['transcription'])
+    batch["input_length"] = len(batch["input_values"][0])
+    return batch
+
+
+negel_79_dataset = load_dataset("audiofolder", data_dir=negel_79, split="train")
+negel_79_dataset = negel_79_dataset.map(dataclean)
+
+def train_dev_test_split(dataset: Dataset, dev_rate=0.1, test_rate=0.1, seed=1):
+    """
+    input: dataset
+    dev_rate,
+    test_rate
+    seed
+    -------
+    Output:
+    dataset_dict{"train", "dev", "test"}
+    """
+    train_dev_test = dataset.train_test_split(test_size=test_rate, seed=seed)
+    test = train_dev_test["test"]
+    train_dev = train_dev_test['train']
+    
+    # pdb.set_trace()
+    if len(train_dev) <= int(len(dataset)*dev_rate):
+        train = Dataset.from_dict({"audio": [], "transcription": []})
+        dev = train_dev
+    else:
+        train_dev = train_dev.train_test_split(test_size=int(len(dataset)*dev_rate), seed=seed)
+        train = train_dev['train']
+        dev = train_dev['test']
+    return train, dev, test
+
+# pdb.set_trace()
+# P1tony_train, P1tony_dev, P1tony_test = train_dev_test_split(P1tony_dataset, dev_rate=0.5, test_rate=0.5, seed=1)
+# P1tony_train_ = concatenate_datasets([P1tony_train,P1tony_scripted]) 
+# pdb.set_trace()
+
+Negel_79_train, Negel_79_dev, Negel_79_test = train_dev_test_split(negel_79_dataset, dev_rate=0.1, test_rate=0.1, seed=1)
+ 
+# src_dataset = load_dataset("audiofolder", data_dir=audio_dir, split="train")
+# src_dataset = src_dataset.map(dataclean)
+
+# healthy_test_dataset = load_dataset(
+#     "audiofolder", data_dir=healthy_dir, split='train')
+# healthy_test_dataset = healthy_test_dataset.map(dataclean)
+
+# Fary_PAL_test_dataset = load_dataset(
+#     "audiofolder", data_dir=Fary_PAL_30, split='train')
+# Fary_PAL_test_dataset = Fary_PAL_test_dataset.map(dataclean)
+
+# John_p326_test_dataset = load_dataset(
+#     "audiofolder", data_dir=John_p326, split='train')
+# John_p326_test_dataset = John_p326_test_dataset.map(dataclean)
+
+# John_video_test_dataset = load_dataset(
+#     "audiofolder", data_dir=John_video, split='train')
+# John_video_test_dataset = John_video_test_dataset.map(dataclean)
+
+# patient_T_test_dataset = load_dataset("audiofolder", data_dir=patient_T, split='train')
+# patient_T_test_dataset = patient_T_test_dataset.map(dataclean)
+
+# patient_L_test_dataset = load_dataset("audiofolder", data_dir=patient_L, split='train')
+# patient_L_test_dataset = patient_L_test_dataset.map(dataclean)
+# pdb.set_trace()
+
+# train_dev / test
+# ds = src_dataset.train_test_split(test_size=0.1, seed=1)
+
+# dataset_libri = load_dataset(
+#     "hf-internal-testing/librispeech_asr_dummy", "clean", split="validation")
+
+# train_dev = ds['train']
+# # train / dev
+# train_dev = train_dev.train_test_split(
+#     test_size=int(len(src_dataset)*0.1), seed=1)
+# # train/dev/test
+# train = train_dev['train']
+# test = ds['test']
+# dev = train_dev['test']
+
+# # pdb.set_trace()
+# encoded_train = train.map(prepare_dataset, num_proc=4)
+# encoded_dev = dev.map(prepare_dataset, num_proc=4)
+# encoded_test = test.map(prepare_dataset, num_proc=4)
+
+# encoded_healthy = healthy_test_dataset.map(prepare_dataset, num_proc=4)
+# encoded_Fary = Fary_PAL_test_dataset.map(prepare_dataset, num_proc=4)
+# encoded_John_p326 = John_p326_test_dataset.map(prepare_dataset, num_proc=4)
+# encoded_John_video = John_video_test_dataset.map(prepare_dataset, num_proc=4)
+# pdb.set_trace()
+
+WER = evaluate.load("wer")
+
+# Whisper decoding
+
+processor = WhisperProcessor.from_pretrained("openai/whisper-medium")
+model = WhisperForConditionalGeneration.from_pretrained(
+    "openai/whisper-medium").to("cuda:0")
+tokenizer = WhisperTokenizer.from_pretrained(
+    "openai/whisper-medium", language="English", task="transcribe")
+
+# Need to push tokenizer to hugginface/model to activate online API
+
+# tokenizer.push_to_hub("KevinGeng/whipser_medium_en_PAL300_step25")
+# import pdb
+# pdb.set_trace()
+
+feature_extractor = WhisperFeatureExtractor.from_pretrained(
+    "openai/whisper-medium")
+
+
+def whisper_prepare_dataset(batch):
+    # load and resample audio data from 48 to 16kHz
+    audio = batch["audio"]
+
+    # compute log-Mel input features from input audio array
+    batch["input_features"] = feature_extractor(
+        audio["array"], sampling_rate=audio["sampling_rate"]).input_features[0]
+
+    # encode target text to label ids
+    batch["labels"] = tokenizer(batch["transcription"]).input_ids
+    return batch
+
+
+torch.cuda.empty_cache()
+
+training_args = Seq2SeqTrainingArguments(
+    # change to a repo name of your choice
+    output_dir="./whisper-medium-PAL128-25step",
+    per_device_train_batch_size=8,
+    gradient_accumulation_steps=1,  # increase by 2x for every 2x decrease in batch size
+    learning_rate=1e-5,
+    warmup_steps=100,
+    max_steps=1000,
+    gradient_checkpointing=True,
+    fp16=True,
+    evaluation_strategy="steps",
+    per_device_eval_batch_size=8,
+    predict_with_generate=True,
+    generation_max_length=512,
+    save_steps=100,
+    eval_steps=25,
+    logging_steps=100,
+    report_to=["tensorboard"],
+    load_best_model_at_end=True,
+    metric_for_best_model="wer",
+    greater_is_better=False,
+    push_to_hub=True,
+)
+
+
+def my_map_to_pred(batch):
+    # pdb.set_trace()
+    audio = batch["audio"]
+    input_features = processor(
+        audio["array"], sampling_rate=audio["sampling_rate"], return_tensors="pt").input_features
+    # batch["reference"] = whisper_processor.tokenizer._normalize(batch['text'])
+    batch["reference"] = processor.tokenizer._normalize(batch['transcription'])
+
+    with torch.no_grad():
+        # predicted_ids = whisper_model.generate(input_features.to("cuda"))[0]
+        predicted_ids = model.generate(input_features.to("cuda"))[0]
+    transcription = model.decode(predicted_ids)
+    batch["prediction"] = model.tokenizer._normalize(transcription)
+    return batch
+
+
+@dataclass
+class DataCollatorSpeechSeq2SeqWithPadding:
+    processor: Any
+
+    def __call__(self, features: List[Dict[str, Union[List[int], torch.Tensor]]]) -> Dict[str, torch.Tensor]:
+        # split inputs and labels since they have to be of different lengths and need different padding methods
+        # first treat the audio inputs by simply returning torch tensors
+        input_features = [{"input_features": feature["input_features"]}
+                          for feature in features]
+        batch = self.processor.feature_extractor.pad(
+            input_features, return_tensors="pt")
+
+        # get the tokenized label sequences
+        label_features = [{"input_ids": feature["labels"]}
+                          for feature in features]
+        # pad the labels to max length
+        labels_batch = self.processor.tokenizer.pad(
+            label_features, return_tensors="pt")
+
+        # replace padding with -100 to ignore loss correctly
+        labels = labels_batch["input_ids"].masked_fill(
+            labels_batch.attention_mask.ne(1), -100)
+
+        # if bos token is appended in previous tokenization step,
+        # cut bos token here as it's append later anyways
+        if (labels[:, 0] == self.processor.tokenizer.bos_token_id).all().cpu().item():
+            labels = labels[:, 1:]
+
+        batch["labels"] = labels
+
+        return batch
+
+
+data_collator = DataCollatorSpeechSeq2SeqWithPadding(processor=processor)
+
+
+def compute_metrics(pred):
+    pdb.set_trace()
+    pred_ids = pred.predictions
+    label_ids = pred.label_ids
+
+    # replace -100 with the pad_token_id
+    label_ids[label_ids == -100] = tokenizer.pad_token_id
+
+    # we do not want to group tokens when computing the metrics
+    pred_str = tokenizer.batch_decode(pred_ids, skip_special_tokens=True)
+    label_str = tokenizer.batch_decode(label_ids, skip_special_tokens=True)
+
+    wer = 100 * WER.compute(predictions=pred_str, references=label_str)
+
+    return {"wer": wer}
+
+encode_negel_79_train = Negel_79_train.map(whisper_prepare_dataset, num_proc=4)
+encode_negel_79_dev = Negel_79_dev.map(whisper_prepare_dataset, num_proc=4)
+encode_negel_79_test = Negel_79_test.map(whisper_prepare_dataset, num_proc=4)
+pdb.set_trace()
+torch.cuda.empty_cache()
+
+torch.cuda.empty_cache()
+
+fine_tuned_model = WhisperForConditionalGeneration.from_pretrained(
+    fine_tuning_dir
+).to("cuda")
+    # "fine_tuned/SSD/model/whipser_medium_TEP_patient_T"
+    # "./fine_tuned/whipser_medium_en_PAL300_step25_step2_VCTK/checkpoint-400"
+    #"./fine_tuned/whipser_medium_en_PAL300_step25_step2_VCTK/checkpoint-200"
+
+
+def fine_tuned_map_to_pred(batch):
+    # pdb.set_trace()
+    audio = batch["audio"]
+    input_features = processor(
+        audio["array"], sampling_rate=audio["sampling_rate"], return_tensors="pt").input_features
+    # batch["reference"] = whisper_processor.tokenizer._normalize(batch['text'])
+    batch["reference"] = processor.tokenizer._normalize(batch['transcription'])
+
+    with torch.no_grad():
+        # predicted_ids = whisper_model.generate(input_features.to("cuda"))[0]
+        predicted_ids = fine_tuned_model.generate(input_features.to("cuda"))[0]
+    transcription = tokenizer.decode(predicted_ids)
+    batch["prediction"] = tokenizer._normalize(transcription)
+    return batch
+
+
+    # output_dir="./fine_tuned/whipser_medium_en_PAL300_step25_step2_VCTK/checkpoint-400",
+testing_args = Seq2SeqTrainingArguments(
+    # change to a repo name of your choice
+    output_dir="fine_tuned/SSD/model/whipser_medium_TEP_patient_TL_TL",
+    per_device_train_batch_size=8,
+    gradient_accumulation_steps=1,  # increase by 2x for every 2x decrease in batch size
+    learning_rate=1e-5,
+    warmup_steps=100,
+    max_steps=1000,
+    gradient_checkpointing=True,
+    fp16=True,
+    evaluation_strategy="steps",
+    per_device_eval_batch_size=8,
+    predict_with_generate=True,
+    generation_max_length=512,
+    save_steps=100,
+    eval_steps=25,
+    logging_steps=100,
+    report_to=["tensorboard"],
+    load_best_model_at_end=True,
+    metric_for_best_model="wer",
+    greater_is_better=False,
+    push_to_hub=False,
+)
+
+predict_trainer = Seq2SeqTrainer(
+    args=testing_args,
+    model=fine_tuned_model,
+    data_collator=data_collator,
+    compute_metrics=compute_metrics,
+    tokenizer=processor.feature_extractor,
+)
+
+# trainer.train()
+# fine tuned
+# z_result = encoded_test.map(fine_tuned_map_to_pred)
+pdb.set_trace()
+z_result= encode_negel_79_test.map(fine_tuned_map_to_pred)
+# 0.4692737430167598
+z = WER.compute(references=z_result['reference'], predictions=z_result['prediction'])
+# pdb.set_trace()
+# z_hel_result = encoded_healthy.map(fine_tuned_map_to_pred)
+# z_hel = WER.compute(references=z_hel_result['reference'], predictions=z_hel_result['prediction'])
+# # 0.1591610117211598
+
+# # pdb.set_trace()
+# # z_fary_result = encoded_Fary.map(fine_tuned_map_to_pred)
+# # z_far = WER.compute(references=z_fary_result['reference'], predictions=z_fary_result['prediction'])
+# # 0.1791044776119403
+# z_patient_LT = encoded_patient_TL_test.map(fine_tuned_map_to_pred)
+# z_patient_LT_result = WER.compute(references=z_patient_LT['reference'], predictions=z_patient_LT['prediction'])
+# z_patient_L = encoded_patient_L_test.map(fine_tuned_map_to_pred)
+# z_patient_L_result = WER.compute(references=z_patient_L['reference'], predictions=z_patient_L['prediction'])
+# z_patient_T = encoded_patient_T_test.map(fine_tuned_map_to_pred)
+# z_patient_T_result = WER.compute(references=z_patient_T['reference'], predictions=z_patient_T['prediction'])
+
+# # z_john_p326_result = encoded_John_p326.map(fine_tuned_map_to_pred)
+# # pdb.set_trace() 
+
+# # z_john_p326 = WER.compute(references=z_john_p326_result['reference'], predictions=z_john_p326_result['prediction'])
+# # 0.4648241206030151
+pdb.set_trace() 
+
+# # y_John_video= fine_tuned_trainer.predict(encoded_John_video)
+# # metrics={'test_loss': 2.665189743041992, 'test_wer': 0.7222222222222222, 'test_runtime': 0.1633, 'test_samples_per_second': 48.979, 'test_steps_per_second': 6.122})
+# pdb.set_trace()
+
+# p326 training
+# metrics={'test_loss': 0.4804028868675232, 'test_wer': 0.21787709497206703, 'test_runtime': 0.3594, 'test_samples_per_second': 44.517, 'test_steps_per_second': 5.565})
+# hel metrics={'test_loss': 1.6363693475723267, 'test_wer': 0.17951881554595928, 'test_runtime': 3.8451, 'test_samples_per_second': 41.611, 'test_steps_per_second': 5.201})
+# Fary: metrics={'t est_loss': 1.4633615016937256, 'test_wer': 0.5572139303482587, 'test_runtime': 0.6627, 'test_samples_per_second': 45.27, 'test_steps_per_second': 6.036})
+# p326 large: metrics={'test_loss': 0.6568527817726135, 'test_wer': 0.2889447236180904, 'test_runtime': 0.7169, 'test_samples_per_second': 51.613, 'test_steps_per_second': 6.975})

local/duration_calcutator.py

@@ -0,0 +1,19 @@
+import os
+import librosa
+
+folder_path = "/home/kevingeng/Disk2/laronix/laronix_automos/data/Patient_sil_trim_16k_normed_5_snr_40/Sentences"  # Replace with the path to your folder
+total_duration = 0.0
+
+# Iterate through all files in the folder
+for filename in os.listdir(folder_path):
+    file_path = os.path.join(folder_path, filename)
+    if os.path.isfile(file_path):
+        try:
+            # Load the audio file and get its duration
+            audio_data, _ = librosa.load(file_path)
+            duration = librosa.get_duration(audio_data)
+            total_duration += duration
+        except Exception as e:
+            print(f"Error processing file '{filename}': {e}")
+
+print(f"Total duration of audio files in the folder: {total_duration} seconds.")

local/fine-tuning.py

@@ -0,0 +1,346 @@
+"""
+TODO:
+    + [x] Load Configuration
+    + [ ] Multi ASR Engine
+    + [ ] Batch / Real Time support
+"""
+from pathlib import Path
+from transformers import AutoTokenizer, AutoFeatureExtractor, AutoModelForCTC, AutoProcessor
+from datasets import load_dataset
+from datasets import Dataset, Audio
+import pdb
+import string
+# local import
+import sys
+
+sys.path.append("src")
+
+# token_model = AutoModelForCTC.from_pretrained(
+#     "facebook/wav2vec2-base-960h"
+# )
+
+# ds = load_dataset("patrickvonplaten/librispeech_asr_dummy", "clean", split="validation")
+
+audio_path = "/Users/kevingeng/Laronix/Laronix_PAL_ASR_Offline_Plot/data/samples/3_Healthy1.wav"
+
+audio_dir= "/Users/kevingeng/Laronix/laronix_automos/data/Patient_sil_trim_16k_normed_5_snr_40/"
+# tgt_audio_dir= "/Users/kevingeng/Laronix/Dataset/Pneumatic/automos"
+
+# src_audio_list = sorted(Path(src_audio_dir).glob("**/*.wav"))
+# src_audio_list = [str(x) for x in src_audio_list]
+# src_audio_dict = {"audio": src_audio_list}
+# src_dataset = Dataset.from_dict(src_audio_dict).cast_column("audio", Audio())
+
+# tgt_audio_list = sorted(Path(tgt_audio_dir).glob("**/*.wav"))
+# tgt_audio_list = [str(x) for x in tgt_audio_list]
+# tgt_audio_dict = {"audio": tgt_audio_list}
+# tgt_dataset = Dataset.from_dict(tgt_audio_dict).cast_column("audio", Audio())
+
+# Get Transcription, WER and PPM 
+"""
+TODO:
+    [DONE]: Automatic generating Config
+"""
+
+import yaml
+import argparse
+import sys
+from pathlib import Path
+
+sys.path.append("./src")
+import lightning_module
+from UV import plot_UV, get_speech_interval
+from transformers import pipeline
+from rich.progress import track
+from rich import print as rprint
+import numpy as np
+import jiwer
+import pdb
+import torch.nn as nn
+import torch
+import torchaudio
+import gradio as gr
+from sys import flags
+from random import sample
+from transformers import Wav2Vec2Processor, Wav2Vec2ForCTC
+
+# root_path = Path(__file__).parents[1]
+
+class ChangeSampleRate(nn.Module):
+    def __init__(self, input_rate: int, output_rate: int):
+        super().__init__()
+        self.output_rate = output_rate
+        self.input_rate = input_rate
+
+    def forward(self, wav: torch.tensor) -> torch.tensor:
+        # Only accepts 1-channel waveform input
+        wav = wav.view(wav.size(0), -1) 
+        new_length = wav.size(-1) * self.output_rate // self.input_rate
+        indices = torch.arange(new_length) * (
+            self.input_rate / self.output_rate
+        )
+        round_down = wav[:, indices.long()]
+        round_up = wav[:, (indices.long() + 1).clamp(max=wav.size(-1) - 1)]
+        output = round_down * (1.0 - indices.fmod(1.0)).unsqueeze(
+            0
+        ) + round_up * indices.fmod(1.0).unsqueeze(0)
+        return output
+
+
+model = lightning_module.BaselineLightningModule.load_from_checkpoint(
+    "./src/epoch=3-step=7459.ckpt"
+).eval()
+
+
+def calc_wer(audio_path, ref):
+    wav, sr = torchaudio.load(audio_path)
+    osr = 16_000
+    batch = wav.unsqueeze(0).repeat(10, 1, 1)
+    csr = ChangeSampleRate(sr, osr)
+    out_wavs = csr(wav)
+    # ASR
+    trans = p(audio_path)["text"]
+    # WER
+    wer = jiwer.wer(
+        ref,
+        trans,
+        truth_transform=transformation,
+        hypothesis_transform=transformation,
+    )
+    return trans, wer
+
+# if __name__ == "__main__":
+#     # Argparse
+#     parser = argparse.ArgumentParser(
+#         prog="get_ref_PPM",
+#         description="Generate Phoneme per Minute (and Voice/Unvoice plot)",
+#         epilog="",
+#     )
+#     parser.add_argument(
+#         "--tag",
+#         type=str,
+#         default=None,
+#         required=False,
+#         help="ID tag for output *.csv",
+#     )
+
+#     parser.add_argument("--ref_txt", type=str, required=True, help="Reference TXT")
+#     parser.add_argument(
+#         "--ref_wavs", type=str, required=True, help="Reference WAVs"
+#     )
+
+#     parser.add_argument(
+#         "--output_dir",
+#         type=str,
+#         required=True,
+#         help="Output Directory for *.csv",
+#     )
+#     parser.add_argument(
+#         "--to_config",
+#         choices=["True", "False"],
+#         default="False",
+#         help="Generating Config from .txt and wavs/*wav",
+#     )
+
+
+#     args = parser.parse_args()
+
+#     refs = np.loadtxt(args.ref_txt, delimiter="\n", dtype="str")
+#     refs_ids = [x.split()[0] for x in refs]
+#     refs_txt = [" ".join(x.split()[1:]) for x in refs]
+#     ref_wavs = [str(x) for x in sorted(Path(args.ref_wavs).glob("**/*.wav"))]
+#     # pdb.set_trace()
+#     try:
+#         len(refs) == len(ref_wavs)
+#     except ValueError:
+#         print("Error: Text and Wavs don't match")
+#         exit()
+
+#     # ASR part
+#     p = pipeline("automatic-speech-recognition")
+
+#     # WER part
+#     transformation = jiwer.Compose(
+#         [
+#             jiwer.ToLowerCase(),
+#             jiwer.RemoveWhiteSpace(replace_by_space=True),
+#             jiwer.RemoveMultipleSpaces(),
+#             jiwer.ReduceToListOfListOfWords(word_delimiter=" "),
+#         ]
+#     )
+
+#     # WPM part
+#     processor = Wav2Vec2Processor.from_pretrained(
+#         "facebook/wav2vec2-xlsr-53-espeak-cv-ft"
+#     )
+#     phoneme_model = Wav2Vec2ForCTC.from_pretrained(
+#         "facebook/wav2vec2-xlsr-53-espeak-cv-ft"
+#     )
+#     # phoneme_model =  pipeline(model="facebook/wav2vec2-xlsr-53-espeak-cv-ft")
+
+#     description = """
+#     MOS prediction demo using UTMOS-strong w/o phoneme encoder model, \
+#         which is trained on the main track dataset.
+#         This demo only accepts .wav format. Best at 16 kHz sampling rate.
+
+#     Paper is available [here](https://arxiv.org/abs/2204.02152)
+
+#     Add ASR based on wav2vec-960, currently only English available.
+#     Add WER interface.
+#     """
+
+#     referance_id = gr.Textbox(
+#         value="ID", placeholder="Utter ID", label="Reference_ID"
+#     )
+#     referance_textbox = gr.Textbox(
+#         value="", placeholder="Input reference here", label="Reference"
+#     )
+#     # Set up interface
+#     result = []
+#     result.append("id, trans, wer")
+
+    
+#     for id, x, y in track(
+#         zip(refs_ids, ref_wavs, refs_txt),
+#         total=len(refs_ids),
+#         description="Loading references information",
+#     ):
+#         trans, wer = calc_wer(x, y)
+#         record = ",".join(
+#             [
+#                 id,
+#                 str(trans),
+#                 str(wer)
+#             ]
+#         )
+#         result.append(record)
+
+#     # Output
+#     if args.tag == None:
+#         args.tag = Path(args.ref_wavs).stem
+#     # Make output_dir
+#     # pdb.set_trace()
+#     Path.mkdir(Path(args.output_dir), exist_ok=True)
+#     # pdb.set_trace()
+#     with open("%s/%s.csv" % (args.output_dir, args.tag), "w") as f:
+#         print("\n".join(result), file=f)
+
+#     # Generating config
+#     if args.to_config == "True":
+#         config_dict = {
+#             "exp_id": args.tag,
+#             "ref_txt": args.ref_txt,
+#             "ref_feature": "%s/%s.csv" % (args.output_dir, args.tag),
+#             "ref_wavs": args.ref_wavs,
+#             "thre": {
+#                 "minppm": 100,
+#                 "maxppm": 100,
+#                 "WER": 0.1,
+#                 "AUTOMOS": 4.0,
+#             },
+#             "auth": {"username": None, "password": None},
+#         }
+#         with open("./config/%s.yaml" % args.tag, "w") as config_f:
+#             rprint("Dumping as config ./config/%s.yaml" % args.tag)
+#             rprint(config_dict)
+#             yaml.dump(config_dict, stream=config_f)
+#             rprint("Change parameter ./config/%s.yaml if necessary" % args.tag)
+#     print("Reference Dumping Finished")
+def dataclean(example):
+    return {"transcription": example["transcription"].upper().translate(str.maketrans('', '', string.punctuation))}
+
+# processor = AutoFeatureExtractor.from_pretrained(
+#     "facebook/wav2vec2-base-960h"
+# )
+processor = AutoProcessor.from_pretrained("facebook/wav2vec2-base")
+
+def prepare_dataset(batch):
+    audio = batch["audio"]
+    batch = processor(audio["array"], sampling_rate = audio["sampling_rate"], text=batch['transcription'])
+    batch["input_length"] = len(batch["input_values"][0])
+    return batch
+
+src_dataset = load_dataset("audiofolder", data_dir=audio_dir, split="train")
+src_dataset = src_dataset.map(dataclean)
+# train_dev / test
+ds = src_dataset.train_test_split(test_size=0.1)
+
+train_dev = ds['train']
+# train / dev
+train_dev = train_dev.train_test_split(test_size=int(len(src_dataset)*0.1))
+# train/dev/test
+train = train_dev['train']
+test = ds['test']
+dev = train_dev['test']
+
+    # pdb.set_trace()
+import numpy as np
+
+
+def compute_metrics(pred):
+    pred_logits = pred.predictions
+    pred_ids = np.argmax(pred_logits, axis=-1)
+
+    pred.label_ids[pred.label_ids == -100] = processor.tokenizer.pad_token_id
+
+    pred_str = processor.batch_decode(pred_ids)
+    label_str = processor.batch_decode(pred.label_ids, group_tokens=False)
+
+    wer = wer.compute(predictions=pred_str, references=label_str)
+
+    return {"wer": wer}
+
+
+pdb.set_trace()
+# TOKENLIZER("data/samples/5_Laronix1.wav")
+# pdb.set_trace()
+# tokenizer 
+tokenizer = AutoTokenizer.from_pretrained("facebook/wav2vec2-base-960h")
+
+encoded_train = train.map(prepare_dataset, num_proc=4)
+
+from transformers import AutoModelForCTC, TrainingArguments, Trainer
+
+model = AutoModelForCTC.from_pretrained(
+    "facebook/wav2vec2-base",
+    ctc_loss_reduction="mean",
+    pad_token_id=processor.tokenizer.pad_token_id,
+)
+pdb.set_trace()
+
+training_args = TrainingArguments(
+    output_dir="my_awesome_asr_mind_model",
+    per_device_train_batch_size=8,
+    gradient_accumulation_steps=2,
+    learning_rate=1e-5,
+    warmup_steps=500,
+    max_steps=2000,
+    gradient_checkpointing=True,
+    fp16=True,
+    group_by_length=True,
+    evaluation_strategy="steps",
+    per_device_eval_batch_size=8,
+    save_steps=1000,
+    eval_steps=1000,
+    logging_steps=25,
+    load_best_model_at_end=True,
+    metric_for_best_model="wer",
+    greater_is_better=False,
+    push_to_hub=True,
+)
+
+pdb.set_trace()
+trainer = Trainer(
+    model=model,
+    args=training_args,
+    train_dataset=encoded_train["train"],
+    eval_dataset=encoded_train["test"],
+    tokenizer=processor.feature_extractor,
+    compute_metrics=compute_metrics,
+)
+pdb.set_trace()
+    # data_collator=data_collator,
+
+trainer.train()
+# x = tokenizer(test['transcription'][0])
+

local/get_ASR.py

@@ -0,0 +1,232 @@
+# Get Transcription, WER and PPM 
+"""
+TODO:
+    [DONE]: Automatic generating Config
+"""
+
+import yaml
+import argparse
+import sys
+from pathlib import Path
+
+sys.path.append("./src")
+import lightning_module
+from UV import plot_UV, get_speech_interval
+from transformers import pipeline
+from rich.progress import track
+from rich import print as rprint
+import numpy as np
+import jiwer
+import pdb
+import torch.nn as nn
+import torch
+import torchaudio
+import gradio as gr
+from sys import flags
+from random import sample
+from transformers import Wav2Vec2Processor, Wav2Vec2ForCTC
+
+# root_path = Path(__file__).parents[1]
+
+class ChangeSampleRate(nn.Module):
+    def __init__(self, input_rate: int, output_rate: int):
+        super().__init__()
+        self.output_rate = output_rate
+        self.input_rate = input_rate
+
+    def forward(self, wav: torch.tensor) -> torch.tensor:
+        # Only accepts 1-channel waveform input
+        wav = wav.view(wav.size(0), -1) 
+        new_length = wav.size(-1) * self.output_rate // self.input_rate
+        indices = torch.arange(new_length) * (
+            self.input_rate / self.output_rate
+        )
+        round_down = wav[:, indices.long()]
+        round_up = wav[:, (indices.long() + 1).clamp(max=wav.size(-1) - 1)]
+        output = round_down * (1.0 - indices.fmod(1.0)).unsqueeze(
+            0
+        ) + round_up * indices.fmod(1.0).unsqueeze(0)
+        return output
+
+
+model = lightning_module.BaselineLightningModule.load_from_checkpoint(
+    "./src/epoch=3-step=7459.ckpt"
+).eval()
+
+
+def calc_wer(audio_path, ref, ASR_pipeline):
+    wav, sr = torchaudio.load(audio_path)
+    osr = 16_000
+    batch = wav.unsqueeze(0).repeat(10, 1, 1)
+    csr = ChangeSampleRate(sr, osr)
+    out_wavs = csr(wav)
+    # ASR
+    trans = ASR_pipeline(audio_path)["text"]
+    # WER
+    wer = jiwer.wer(
+        ref,
+        trans,
+        truth_transform=transformation,
+        hypothesis_transform=transformation,
+    )
+    return trans, wer
+
+if __name__ == "__main__":
+    # Argparse
+    parser = argparse.ArgumentParser(
+        prog="get_ref_PPM",
+        description="Generate Phoneme per Minute (and Voice/Unvoice plot)",
+        epilog="",
+    )
+    parser.add_argument(
+        "--tag",
+        type=str,
+        default=None,
+        required=False,
+        help="ID tag for output *.csv",
+    )
+
+    parser.add_argument("--ref_txt", type=str, required=True, help="Reference TXT")
+    parser.add_argument(
+        "--ref_wavs", type=str, required=True, help="Reference WAVs"
+    )
+    parser.add_argument(
+        "--metadata",
+        type=str,
+        required=False,
+        help="metadata.csv including wav_id and reference",
+    )
+    
+    parser.add_argument(
+        "--model",
+        type=str,
+        default='whisper-medium-FT',
+        choices=['wav2vec+ctc', 'whipser-medium-FT', 'whipser-large-v2'],
+        help="ASR engine for evaluation:\n ver1: wav2vec+ctc \n ver2: whipser-medium(Fined-tuned)\n ver3: whipser-large-v2",
+    )
+
+    parser.add_argument(
+        "--output_dir",
+        type=str,
+        required=True,
+        help="Output Directory for *.csv",
+    )
+
+    parser.add_argument(
+        "--to_config",
+        choices=["True", "False"],
+        default="False",
+        help="Generating Config from .txt and wavs/*wav",
+    )
+
+
+    args = parser.parse_args()
+
+    refs = np.loadtxt(args.ref_txt, delimiter="\n", dtype="str")
+    refs_ids = [x.split()[0] for x in refs]
+    refs_txt = [" ".join(x.split()[1:]) for x in refs]
+    ref_wavs = [str(x) for x in sorted(Path(args.ref_wavs).glob("**/*.wav"))]
+    # pdb.set_trace()
+    try:
+        len(refs) == len(ref_wavs)
+    except ValueError:
+        print("Error: Text and Wavs don't match")
+        exit()
+
+    # ASR part
+    if args.model== "whisper-medium-FT":
+        ASR_pipeline = pipeline("automatic-speech-recognition", model="KevinGeng/whipser_medium_en_PAL300_step25")
+    elif args.model == "wav2vec+ctc":
+        ASR_pipeline = pipeline("automatic-speech-recognition")
+    elif args.model == "whisper-large-v2":
+        ASR_pipeline = pipeline("automatic-speech-recognition", model="openai/whisper-large-v2")
+    
+    # pdb.set_trace()
+    # WER part
+    transformation = jiwer.Compose(
+        [
+            jiwer.ToLowerCase(),
+            jiwer.RemoveWhiteSpace(replace_by_space=True),
+            jiwer.RemoveMultipleSpaces(),
+            jiwer.ReduceToListOfListOfWords(word_delimiter=" "),
+        ]
+    )
+
+    # WPM part
+    processor = Wav2Vec2Processor.from_pretrained(
+        "facebook/wav2vec2-xlsr-53-espeak-cv-ft"
+    )
+    phoneme_model = Wav2Vec2ForCTC.from_pretrained(
+        "facebook/wav2vec2-xlsr-53-espeak-cv-ft"
+    )
+    # phoneme_model =  pipeline(model="facebook/wav2vec2-xlsr-53-espeak-cv-ft")
+
+    description = """
+    MOS prediction demo using UTMOS-strong w/o phoneme encoder model, \
+        which is trained on the main track dataset.
+        This demo only accepts .wav format. Best at 16 kHz sampling rate.
+
+    Paper is available [here](https://arxiv.org/abs/2204.02152)
+
+    Add ASR based on wav2vec-960, currently only English available.
+    Add WER interface.
+    """
+
+    referance_id = gr.Textbox(
+        value="ID", placeholder="Utter ID", label="Reference_ID"
+    )
+    referance_textbox = gr.Textbox(
+        value="", placeholder="Input reference here", label="Reference"
+    )
+    # Set up interface
+    result = []
+    result.append("id,ref,hyp,wer")
+
+    
+    for id, x, y in track(
+        zip(refs_ids, ref_wavs, refs_txt),
+        total=len(refs_ids),
+        description="Loading references information",
+    ):
+        trans, wer = calc_wer(x, y, ASR_pipeline=ASR_pipeline)
+        record = ",".join(
+            [
+                id,
+                str(y),
+                str(trans),
+                str(wer)
+            ]
+        )
+        result.append(record)
+
+    # Output
+    if args.tag == None:
+        args.tag = Path(args.ref_wavs).stem
+    # Make output_dir
+    # pdb.set_trace()
+    Path.mkdir(Path(args.output_dir), exist_ok=True)
+    # pdb.set_trace()
+    with open("%s/%s.csv" % (args.output_dir, args.tag), "w") as f:
+        print("\n".join(result), file=f)
+
+    # Generating config
+    if args.to_config == "True":
+        config_dict = {
+            "exp_id": args.tag,
+            "ref_txt": args.ref_txt,
+            "ref_feature": "%s/%s.csv" % (args.output_dir, args.tag),
+            "ref_wavs": args.ref_wavs,
+            "thre": {
+                "minppm": 100,
+                "maxppm": 100,
+                "WER": 0.1,
+                "AUTOMOS": 4.0,
+            },
+            "auth": {"username": None, "password": None},
+        }
+        with open("./config/%s.yaml" % args.tag, "w") as config_f:
+            rprint("Dumping as config ./config/%s.yaml" % args.tag)
+            rprint(config_dict)
+            yaml.dump(config_dict, stream=config_f)
+            rprint("Change parameter ./config/%s.yaml if necessary" % args.tag)
+    print("Reference Dumping Finished")

local/get_ref_PPM.py

@@ -0,0 +1,269 @@
+# Get Transcription, WER and PPM 
+"""
+TODO:
+    [DONE]: Automatic generating Config
+"""
+
+import yaml
+import argparse
+import sys
+from pathlib import Path
+
+sys.path.append("./src")
+import lightning_module
+from UV import plot_UV, get_speech_interval
+from transformers import pipeline
+from rich.progress import track
+from rich import print as rprint
+import numpy as np
+import jiwer
+import pdb
+import torch.nn as nn
+import torch
+import torchaudio
+import gradio as gr
+from sys import flags
+from random import sample
+from transformers import Wav2Vec2Processor, Wav2Vec2ForCTC
+
+# root_path = Path(__file__).parents[1]
+
+class ChangeSampleRate(nn.Module):
+    def __init__(self, input_rate: int, output_rate: int):
+        super().__init__()
+        self.output_rate = output_rate
+        self.input_rate = input_rate
+
+    def forward(self, wav: torch.tensor) -> torch.tensor:
+        # Only accepts 1-channel waveform input
+        wav = wav.view(wav.size(0), -1)
+        new_length = wav.size(-1) * self.output_rate // self.input_rate
+        indices = torch.arange(new_length) * (
+            self.input_rate / self.output_rate
+        )
+        round_down = wav[:, indices.long()]
+        round_up = wav[:, (indices.long() + 1).clamp(max=wav.size(-1) - 1)]
+        output = round_down * (1.0 - indices.fmod(1.0)).unsqueeze(
+            0
+        ) + round_up * indices.fmod(1.0).unsqueeze(0)
+        return output
+
+
+model = lightning_module.BaselineLightningModule.load_from_checkpoint(
+    "./src/epoch=3-step=7459.ckpt"
+).eval()
+
+
+def calc_mos(audio_path, ref):
+    wav, sr = torchaudio.load(audio_path)
+    osr = 16_000
+    batch = wav.unsqueeze(0).repeat(10, 1, 1)
+    csr = ChangeSampleRate(sr, osr)
+    out_wavs = csr(wav)
+    # ASR
+    trans = p(audio_path)["text"]
+    # WER
+    wer = jiwer.wer(
+        ref,
+        trans,
+        truth_transform=transformation,
+        hypothesis_transform=transformation,
+    )
+    # MOS
+    batch = {
+        "wav": out_wavs,
+        "domains": torch.tensor([0]),
+        "judge_id": torch.tensor([288]),
+    }
+    with torch.no_grad():
+        output = model(batch)
+    predic_mos = output.mean(dim=1).squeeze().detach().numpy() * 2 + 3
+    # Phonemes per minute (PPM)
+    with torch.no_grad():
+        logits = phoneme_model(out_wavs).logits
+    phone_predicted_ids = torch.argmax(logits, dim=-1)
+    phone_transcription = processor.batch_decode(phone_predicted_ids)
+    lst_phonemes = phone_transcription[0].split(" ")
+    wav_vad = torchaudio.functional.vad(wav, sample_rate=sr)
+    ppm = len(lst_phonemes) / (wav_vad.shape[-1] / sr) * 60
+    # if float(predic_mos) >= 3.0:
+    #     torchaudio.save("good.wav", wav,sr)
+
+    return predic_mos, trans, wer, phone_transcription, ppm
+
+if __name__ == "__main__":
+    # Argparse
+    parser = argparse.ArgumentParser(
+        prog="get_ref_PPM",
+        description="Generate Phoneme per Minute (and Voice/Unvoice plot)",
+        epilog="",
+    )
+    parser.add_argument(
+        "--tag",
+        type=str,
+        default=None,
+        required=False,
+        help="ID tag for output *.csv",
+    )
+
+    parser.add_argument("--ref_txt", type=str, required=True, help="Reference TXT")
+    parser.add_argument(
+        "--ref_wavs", type=str, required=True, help="Reference WAVs"
+    )
+
+    parser.add_argument(
+        "--output_dir",
+        type=str,
+        required=True,
+        help="Output Directory for *.csv",
+    )
+    parser.add_argument(
+        "--to_config",
+        choices=["True", "False"],
+        default="False",
+        help="Generating Config from .txt and wavs/*wav",
+    )
+
+    parser.add_argument(
+        "--UV_flag",
+        choices=["True", "False"],
+        default="False",
+        help="Toggle for U/V plot",
+    )
+    parser.add_argument(
+        "--UV_thre", type=float, default=40, help="U/V threshold dB"
+    )
+    args = parser.parse_args()
+
+    refs = np.loadtxt(args.ref_txt, delimiter="\n", dtype="str")
+    refs_ids = [x.split()[0] for x in refs]
+    refs_txt = [" ".join(x.split()[1:]) for x in refs]
+    ref_wavs = [str(x) for x in sorted(Path(args.ref_wavs).glob("**/*.wav"))]
+    # pdb.set_trace()
+    try:
+        len(refs) == len(ref_wavs)
+    except ValueError:
+        print("Error: Text and Wavs don't match")
+        exit()
+
+    # ASR part
+    p = pipeline("automatic-speech-recognition")
+
+    # WER part
+    transformation = jiwer.Compose(
+        [
+            jiwer.ToLowerCase(),
+            jiwer.RemoveWhiteSpace(replace_by_space=True),
+            jiwer.RemoveMultipleSpaces(),
+            jiwer.ReduceToListOfListOfWords(word_delimiter=" "),
+        ]
+    )
+
+    # WPM part
+    processor = Wav2Vec2Processor.from_pretrained(
+        "facebook/wav2vec2-xlsr-53-espeak-cv-ft"
+    )
+    phoneme_model = Wav2Vec2ForCTC.from_pretrained(
+        "facebook/wav2vec2-xlsr-53-espeak-cv-ft"
+    )
+    # phoneme_model =  pipeline(model="facebook/wav2vec2-xlsr-53-espeak-cv-ft")
+
+    description = """
+    MOS prediction demo using UTMOS-strong w/o phoneme encoder model, \
+        which is trained on the main track dataset.
+        This demo only accepts .wav format. Best at 16 kHz sampling rate.
+
+    Paper is available [here](https://arxiv.org/abs/2204.02152)
+
+    Add ASR based on wav2vec-960, currently only English available.
+    Add WER interface.
+    """
+
+    referance_id = gr.Textbox(
+        value="ID", placeholder="Utter ID", label="Reference_ID"
+    )
+    referance_textbox = gr.Textbox(
+        value="", placeholder="Input reference here", label="Reference"
+    )
+    # Set up interface
+    result = []
+    result.append("id, pred_mos, trans, wer, pred_phone, ppm")
+
+    if args.UV_flag == "False":
+        for id, x, y in track(
+            zip(refs_ids, ref_wavs, refs_txt),
+            total=len(refs_ids),
+            description="Loading references information",
+        ):
+            predic_mos, trans, wer, phone_transcription, ppm = calc_mos(x, y)
+            record = ",".join(
+                [
+                    id,
+                    str(predic_mos),
+                    str(trans),
+                    str(wer),
+                    str(phone_transcription),
+                    str(ppm),
+                ]
+            )
+            result.append(record)
+
+    elif args.UV_flag == "True":
+        fig_tardir = Path(args.ref_wavs) / Path("PPM_figs")
+        Path.mkdir(Path(args.ref_wavs) / Path("PPM_figs"), exist_ok=True)
+
+        for id, x, y in track(
+            zip(refs_ids, ref_wavs, refs_txt),
+            total=len(refs_ids),
+            description="Loading references information",
+        ):
+            # UV ploting
+            wav, sr = torchaudio.load(x)
+            wav_vad = torchaudio.functional.vad(wav, sample_rate=sr)
+            a_h, p_h = get_speech_interval(wav_vad.numpy(), db=args.UV_thre)
+            fig_h = plot_UV(wav_vad.numpy().squeeze(), a_h, sr=sr)
+            fig_h.savefig(Path(fig_tardir) / Path(id + ".png"), dpi=200)
+            # Acoustic calculation
+            predic_mos, trans, wer, phone_transcription, ppm = calc_mos(x, y)
+            record = ",".join(
+                [
+                    id,
+                    str(predic_mos),
+                    str(trans),
+                    str(wer),
+                    str(phone_transcription),
+                    str(ppm),
+                ]
+            )
+            result.append(record)
+    # Output
+    if args.tag == None:
+        args.tag = Path(args.ref_wavs).stem
+    # Make output_dir
+    # pdb.set_trace()
+    Path.mkdir(Path(args.output_dir), exist_ok=True)
+    # pdb.set_trace()
+    with open("%s/%s.csv" % (args.output_dir, args.tag), "w") as f:
+        print("\n".join(result), file=f)
+
+    # Generating config
+    if args.to_config == "True":
+        config_dict = {
+            "exp_id": args.tag,
+            "ref_txt": args.ref_txt,
+            "ref_feature": "%s/%s.csv" % (args.output_dir, args.tag),
+            "ref_wavs": args.ref_wavs,
+            "thre": {
+                "minppm": 100,
+                "maxppm": 100,
+                "WER": 0.1,
+                "AUTOMOS": 4.0,
+            },
+            "auth": {"username": None, "password": None},
+        }
+        with open("./config/%s.yaml" % args.tag, "w") as config_f:
+            rprint("Dumping as config ./config/%s.yaml" % args.tag)
+            rprint(config_dict)
+            yaml.dump(config_dict, stream=config_f)
+            rprint("Change parameter ./config/%s.yaml if necessary" % args.tag)
+    print("Reference Dumping Finished")

local/new_whisper_fine_tuning.py

@@ -0,0 +1,481 @@
+fine_tuning_dir = "/fine_tuned/whipser_medium_en_PAL300_step25_step2_VCTK/checkpoint-400"
+
+"""
+TODO:
+    + [ ] Data load
+    + [ ] Train / Test / Dev spilt 
+    + [ ] Train / Test Phase
+    + [ ] Logging with Train / Dev / Test Loss
+    + [ ] Evalutation metrics
+"""
+import pdb
+import string
+from pathlib import Path
+
+import evaluate
+import librosa
+import torch
+import torch.nn as nn
+from datasets import Dataset, concatenate_datasets, load_dataset
+from transformers import AutoProcessor
+
+wer = evaluate.load("wer")
+torch.cuda.set_device("cuda:0")
+
+audio_dir = "./data/Patient_sil_trim_16k_normed_5_snr_40"
+healthy_dir = "./data/Healthy"
+Fary_PAL_30 = "./data/Fary_PAL_p326_20230110_30"
+John_p326 = "./data/John_p326/output"
+John_video = "./data/20230103_video"
+
+## train
+p326_300_dir = "./data/John_p326_large"
+P1tony_arthur = "data/Participant1_Tony_Recording/CLEAN_SENTENCES/SCRIPTED/Arthur_the_Rat/PAL"
+P1tony_rainbow = "data/Participant1_Tony_Recording/CLEAN_SENTENCES/SCRIPTED/Rainbow_Passage/Laronix"
+
+P1tony = "data/Participant1_Tony_Recording/CLEAN_SENTENCES/CONVERSATIONAL/PAL"
+
+P4Negel = 'data/4_negal_152_clean_all'
+
+def dataclean(example):
+    if example["audio"]["sampling_rate"] != 16000:
+        resampled_audio = librosa.resample(
+            y=example["audio"]["array"],
+            orig_sr=example["audio"]["sampling_rate"],
+            target_sr=16000,
+        )
+
+        return {
+            "audio": {
+                "path": example["audio"]["path"],
+                "array": resampled_audio,
+                "sampling_rate": 16000,
+            },
+            "transcription": example["transcription"]
+            .upper()
+            .translate(str.maketrans("", "", string.punctuation)),
+        }
+    else:
+        return {
+            "transcription": example["transcription"]
+            .upper()
+            .translate(str.maketrans("", "", string.punctuation))
+        }
+
+
+
+P1tony_dataset = load_dataset("audiofolder", data_dir=P1tony, split="train")
+P1tony_dataset = P1tony_dataset.map(dataclean)
+
+P1tony_scripted1 = load_dataset(
+    "audiofolder", data_dir=P1tony_rainbow, split="train"
+)
+P1tony_scripted2 = load_dataset(
+    "audiofolder", data_dir=P1tony_arthur, split="train"
+)
+P1tony_scripted1 = P1tony_scripted1.map(dataclean)
+P1tony_scripted2 = P1tony_scripted2.map(dataclean)
+P1tony_scripted = concatenate_datasets([P1tony_scripted1, P1tony_scripted2])
+
+class ChangeSampleRate(nn.Module):
+    def __init__(self, input_rate: int, output_rate: int):
+        super().__init__()
+        self.output_rate = output_rate
+        self.input_rate = input_rate
+
+    def forward(self, wav: torch.tensor) -> torch.tensor:
+        # Only accepts 1-channel waveform input
+        wav = wav.view(wav.size(0), -1)
+        new_length = wav.size(-1) * self.output_rate // self.input_rate
+        indices = torch.arange(new_length) * (
+            self.input_rate / self.output_rate
+        )
+        round_down = wav[:, indices.long()]
+        round_up = wav[:, (indices.long() + 1).clamp(max=wav.size(-1) - 1)]
+        output = round_down * (1.0 - indices.fmod(1.0)).unsqueeze(
+            0
+        ) + round_up * indices.fmod(1.0).unsqueeze(0)
+        return output
+
+# resample and clean text data
+def dataclean(example):
+    # pdb.set_trace()
+    if example["audio"]["sampling_rate"] != 16000:
+        resampled_audio = librosa.resample(
+            y=example["audio"]["array"],
+            orig_sr=example["audio"]["sampling_rate"],
+            target_sr=16000,
+        )
+        
+        return {
+            "audio": {
+                "path": example["audio"]["path"],
+                "array": resampled_audio,
+                "sampling_rate": 16000,
+            },
+            "transcription": example["transcription"]
+            .upper()
+            .translate(str.maketrans("", "", string.punctuation)),
+        }
+    else:
+        return {
+            "transcription": example["transcription"]
+            .upper()
+            .translate(str.maketrans("", "", string.punctuation))
+        }
+
+
+# processor = AutoFeatureExtractor.from_pretrained(
+#     "facebook/wav2vec2-base-960h"
+# )
+processor = AutoProcessor.from_pretrained("facebook/wav2vec2-base-960h")
+
+def prepare_dataset(batch):
+    audio = batch["audio"]
+    batch = processor(
+        audio["array"],
+        sampling_rate=audio["sampling_rate"],
+        text=batch["transcription"],
+    )
+    batch["input_length"] = len(batch["input_values"][0])
+    return batch
+
+src_dataset = load_dataset("audiofolder", data_dir=audio_dir, split="train")
+src_dataset = src_dataset.map(dataclean)
+p326_300_dataset = load_dataset(
+    "audiofolder", data_dir=p326_300_dir, split="train"
+)
+p326_300_dataset = p326_300_dataset.map(dataclean)
+
+P4Negel_dataset = load_dataset("audiofolder", data_dir=P4Negel, split="train")
+P4Negel_dataset = P4Negel_dataset.map(dataclean)
+
+healthy_test_dataset = load_dataset(
+    "audiofolder", data_dir=healthy_dir, split="train"
+)
+healthy_test_dataset = healthy_test_dataset.map(dataclean)
+
+Fary_PAL_test_dataset = load_dataset(
+    "audiofolder", data_dir=Fary_PAL_30, split="train"
+)
+Fary_PAL_test_dataset = Fary_PAL_test_dataset.map(dataclean)
+
+John_p326_test_dataset = load_dataset(
+    "audiofolder", data_dir=John_p326, split="train"
+)
+John_p326_test_dataset = John_p326_test_dataset.map(dataclean)
+
+John_video_test_dataset = load_dataset(
+    "audiofolder", data_dir=John_video, split="train"
+)
+John_video_test_dataset = John_video_test_dataset.map(dataclean)
+
+
+def train_dev_test_split(
+    dataset: Dataset, dev_rate=0.1, test_rate=0.1, seed=1
+):
+    """
+    input: dataset
+    dev_rate,
+    test_rate
+    seed
+    -------
+    Output:
+    dataset_dict{"train", "dev", "test"}
+    """
+    train_dev_test = dataset.train_test_split(test_size=test_rate, seed=seed)
+    test = train_dev_test["test"]
+    train_dev = train_dev_test["train"]
+
+    # pdb.set_trace()
+    if len(train_dev) <= int(len(dataset) * dev_rate):
+        train = Dataset.from_dict({"audio": [], "transcription": []})
+        dev = train_dev
+    else:
+        train_dev = train_dev.train_test_split(
+            test_size=int(len(dataset) * dev_rate), seed=seed
+        )
+        train = train_dev["train"]
+        dev = train_dev["test"]
+    return train, dev, test
+
+P1tony_train, P1tony_dev, P1tony_test = train_dev_test_split(
+    P1tony_dataset, dev_rate=0.5, test_rate=0.5, seed=1
+)
+P1tony_train_ = concatenate_datasets([P1tony_train, P1tony_scripted])
+
+# train_dev / test
+ds = src_dataset.train_test_split(test_size=0.1, seed=1)
+
+# dataset_libri = load_dataset(
+#     "hf-internal-testing/librispeech_asr_dummy", "clean", split="validation"
+# )
+
+train_dev = ds["train"]
+# train / dev
+train_dev = train_dev.train_test_split(
+    test_size=int(len(src_dataset) * 0.1), seed=1
+)
+
+# Tony
+Tony_train = P1tony_train_
+Tony_dev = P1tony_dev
+Tony_test = P1tony_test
+
+# John
+John_train, John_dev, John_test = train_dev_test_split(p326_300_dataset, dev_rate=0.1, test_rate=0.1)
+# Negel
+Negel_train, Negel_dev, Negel_test = train_dev_test_split(P4Negel_dataset, dev_rate=0.1, test_rate=0.1)
+
+# train/dev/test
+train = train_dev["train"]
+test = ds["test"]
+dev = train_dev["test"]
+
+# combined
+combine_train = concatenate_datasets([train, Tony_train, John_train, Negel_train])
+conbine_dev = concatenate_datasets([dev, Tony_dev, John_dev, Negel_dev])
+conbine_test = concatenate_datasets([test, Tony_test, John_test, Negel_test])
+
+# encoded_train = combine_train.map(prepare_dataset, num_proc=4)
+# encoded_dev = conbine_dev.map(prepare_dataset, num_proc=4)
+# encoded_test = conbine_test.map(prepare_dataset, num_proc=4)
+
+# # extra_test
+# encoded_Fary = Fary_PAL_test_dataset.map(prepare_dataset, num_proc=4)
+# encoded_healthy = healthy_test_dataset.map(prepare_dataset, num_proc=4)
+
+# encoded_ori_test = test.map(prepare_dataset, num_proc=4)
+# encoded_Tony_test = Tony_test.map(prepare_dataset, num_proc=4)
+# encoded_John_test = John_test.map(prepare_dataset, num_proc=4)
+# encoded_Negel_test = Negel_test.map(prepare_dataset, num_proc=4)
+
+# encoded_train = train.map(prepare_dataset, num_proc=4)
+# encoded_dev = dev.map(prepare_dataset, num_proc=4)
+# p326_encoded_train = p326_300_dataset.map(prepare_dataset, num_proc=4)
+
+# combine large p326 in to training set
+# encoded_train = concatenate_datasets([encoded_train, p326_encoded_train])
+
+# encoded_John_p326 = John_p326_test_dataset.map(prepare_dataset, num_proc=4)
+# encoded_John_video = John_video_test_dataset.map(prepare_dataset, num_proc=4)
+
+# pdb.set_trace()
+import numpy as np
+
+WER = evaluate.load("wer")
+
+## Whisper decoding
+
+from transformers import (Seq2SeqTrainer, Seq2SeqTrainingArguments,
+                          WhisperFeatureExtractor,
+                          WhisperForConditionalGeneration, WhisperModel,
+                          WhisperProcessor, WhisperTokenizer)
+
+processor = WhisperProcessor.from_pretrained("openai/whisper-medium")
+# model = WhisperForConditionalGeneration.from_pretrained(
+#     "./fine_tuned/whipser_medium_en_PAL300_step25_step2_VCTK/checkpoint-400",
+#     use_auth_token=True,
+# ).to("cuda:0")
+model = WhisperForConditionalGeneration.from_pretrained(
+    "openai/whisper-medium",
+).to("cuda:0")
+tokenizer = WhisperTokenizer.from_pretrained(
+    "openai/whisper-medium", language="English", task="transcribe"
+)
+
+from pathlib import Path
+
+id = Path(fine_tuning_dir).stem
+pdb.set_trace()
+tokenizer.push_to_hub("KevinGeng/%s" % id)
+# import pdb
+feature_extractor = WhisperFeatureExtractor.from_pretrained(
+    "openai/whisper-medium"
+)
+
+def whisper_prepare_dataset(batch):
+    # load and resample audio data from 48 to 16kHz
+    audio = batch["audio"]
+
+    # compute log-Mel input features from input audio array
+    batch["input_features"] = feature_extractor(
+        audio["array"], sampling_rate=audio["sampling_rate"]
+    ).input_features[0]
+
+    # encode target text to label ids
+    batch["labels"] = tokenizer(batch["transcription"]).input_ids
+    return batch
+
+torch.cuda.empty_cache()
+
+
+def my_map_to_pred(batch):
+    # pdb.set_trace()
+    audio = batch["audio"]
+    input_features = processor(
+        audio["array"],
+        sampling_rate=audio["sampling_rate"],
+        return_tensors="pt",
+    ).input_features
+    # batch["reference"] = whisper_processor.tokenizer._normalize(batch['text'])
+    batch["reference"] = processor.tokenizer._normalize(batch["transcription"])
+
+    with torch.no_grad():
+        # predicted_ids = whisper_model.generate(input_features.to("cuda"))[0]
+        predicted_ids = model.generate(input_features.to("cuda"))[0]
+    transcription = model.decode(predicted_ids)
+    batch["prediction"] = model.tokenizer._normalize(transcription)
+    return batch
+
+
+from dataclasses import dataclass
+from typing import Any, Dict, List, Union
+
+import torch
+
+
+@dataclass
+class DataCollatorSpeechSeq2SeqWithPadding:
+    processor: Any
+
+    def __call__(
+        self, features: List[Dict[str, Union[List[int], torch.Tensor]]]
+    ) -> Dict[str, torch.Tensor]:
+        # split inputs and labels since they have to be of different lengths and need different padding methods
+        # first treat the audio inputs by simply returning torch tensors
+        input_features = [
+            {"input_features": feature["input_features"]}
+            for feature in features
+        ]
+        batch = self.processor.feature_extractor.pad(
+            input_features, return_tensors="pt"
+        )
+
+        # get the tokenized label sequences
+        label_features = [
+            {"input_ids": feature["labels"]} for feature in features
+        ]
+        # pad the labels to max length
+        labels_batch = self.processor.tokenizer.pad(
+            label_features, return_tensors="pt"
+        )
+
+        # replace padding with -100 to ignore loss correctly
+        labels = labels_batch["input_ids"].masked_fill(
+            labels_batch.attention_mask.ne(1), -100
+        )
+
+        # if bos token is appended in previous tokenization step,
+        # cut bos token here as it's append later anyways
+        if (
+            (labels[:, 0] == self.processor.tokenizer.bos_token_id)
+            .all()
+            .cpu()
+            .item()
+        ):
+            labels = labels[:, 1:]
+
+        batch["labels"] = labels
+
+        return batch
+
+
+data_collator = DataCollatorSpeechSeq2SeqWithPadding(processor=processor)
+
+
+def compute_metrics(pred):
+    pred_ids = pred.predictions
+    label_ids = pred.label_ids
+
+    # replace -100 with the pad_token_id
+    label_ids[label_ids == -100] = tokenizer.pad_token_id
+
+    # we do not want to group tokens when computing the metrics
+    pred_str = tokenizer.batch_decode(pred_ids, skip_special_tokens=True)
+    label_str = tokenizer.batch_decode(label_ids, skip_special_tokens=True)
+
+    wer = 100 * WER.compute(predictions=pred_str, references=label_str)
+
+    return {"wer": wer}
+
+encoded_train = combine_train.map(whisper_prepare_dataset, num_proc=4)
+encoded_dev = conbine_dev.map(whisper_prepare_dataset, num_proc=4)
+encoded_test = conbine_test.map(whisper_prepare_dataset, num_proc=4)
+
+# extra_test
+
+encoded_ori_test = test.map(whisper_prepare_dataset, num_proc=4)
+encoded_Tony_test = Tony_test.map(whisper_prepare_dataset, num_proc=4)
+encoded_John_test = John_test.map(whisper_prepare_dataset, num_proc=4)
+encoded_Negel_test = Negel_test.map(whisper_prepare_dataset, num_proc=4)
+
+encoded_Fary = Fary_PAL_test_dataset.map(whisper_prepare_dataset, num_proc=4)
+encoded_healthy = healthy_test_dataset.map(whisper_prepare_dataset, num_proc=4)
+
+torch.cuda.empty_cache()
+
+training_args = Seq2SeqTrainingArguments(
+    output_dir=fine_tuning_dir,  # change to a repo name of your choice
+    per_device_train_batch_size=8,
+    gradient_accumulation_steps=1,  # increase by 2x for every 2x decrease in batch size
+    learning_rate=1e-5,
+    warmup_steps=50,
+    max_steps=1000,
+    gradient_checkpointing=True,
+    fp16=True,
+    evaluation_strategy="steps",
+    save_strategy="steps",
+    per_device_eval_batch_size=8,
+    predict_with_generate=True,
+    generation_max_length=512,
+    save_steps=20,
+    eval_steps=20,
+    logging_steps=10,
+    report_to=["tensorboard"],
+    load_best_model_at_end=True,
+    metric_for_best_model="wer",
+    greater_is_better=False,
+    save_total_limit=5,
+    push_to_hub=False,
+)
+from transformers import EarlyStoppingCallback, Trainer, TrainingArguments
+
+trainer = Seq2SeqTrainer(
+    args=training_args,
+    model=model,
+    train_dataset=Negel_train,
+    eval_dataset=Negel_dev,
+    data_collator=data_collator,
+    compute_metrics=compute_metrics,
+    tokenizer=processor.feature_extractor,
+    callbacks=[EarlyStoppingCallback(early_stopping_patience=10)],
+)
+# callbacks=[EvalLoggingCallback()]
+pdb.set_trace()
+
+before_result_dict = {
+    "Ori_Test": trainer.evaluate(encoded_ori_test),
+    "Tony_Test": trainer.evaluate(encoded_Tony_test),
+    "John_Test": trainer.evaluate(encoded_John_test),
+    "Negel_Test": trainer.evaluate(encoded_Negel_test),
+    "Zeroshot_Fary_Test": trainer.evaluate(encoded_Fary),
+    "Healthy_Test": trainer.evaluate(encoded_healthy),
+}
+
+print(before_result_dict)
+trainer.train()
+
+pdb.set_trace()
+result_dict = {
+    "Ori_Test": trainer.evaluate(encoded_ori_test),
+    "Tony_Test": trainer.evaluate(encoded_Tony_test),
+    "John_Test": trainer.evaluate(encoded_John_test),
+    "Negel_Test": trainer.evaluate(encoded_Negel_test),
+    "Zeroshot_Fary_Test": trainer.evaluate(encoded_Fary),
+    "Healthy_Test": trainer.evaluate(encoded_healthy),
+}
+
+pdb.set_trace()
+# Evaluation
+model.push_to_hub("KevinGeng/%s" % id)

local/new_whisper_fine_tuning_decode.py

@@ -0,0 +1,509 @@
+fine_tuning_dir = "/home/kevingeng/Disk2/laronix/laronix_automos/fine_tuned/SSD/model/Org_Tony_John_Negel_Train_557_Dev_79_Test_81/checkpoint-160"
+"""
+TODO:
+    + [ ] Data load
+    + [ ] Train / Test / Dev spilt 
+    + [ ] Train / Test Phase
+    + [ ] Logging with Train / Dev / Test Loss
+    + [ ] Evalutation metrics
+"""
+import pdb
+import string
+from pathlib import Path
+
+import evaluate
+import librosa
+import torch
+import torch.nn as nn
+from datasets import Dataset, concatenate_datasets, load_dataset
+from transformers import AutoProcessor
+
+wer = evaluate.load("wer")
+torch.cuda.set_device("cuda:0")
+
+audio_dir = "./data/Patient_sil_trim_16k_normed_5_snr_40"
+healthy_dir = "./data/Healthy"
+Fary_PAL_30 = "./data/Fary_PAL_p326_20230110_30"
+John_p326 = "./data/John_p326/output"
+John_video = "./data/20230103_video"
+
+## train
+p326_300_dir = "./data/John_p326_large"
+P1tony_arthur = "data/Participant1_Tony_Recording/CLEAN_SENTENCES/SCRIPTED/Arthur_the_Rat/PAL"
+P1tony_rainbow = "data/Participant1_Tony_Recording/CLEAN_SENTENCES/SCRIPTED/Rainbow_Passage/Laronix"
+
+P1tony = "data/Participant1_Tony_Recording/CLEAN_SENTENCES/CONVERSATIONAL/PAL"
+
+P4Negel = 'data/4_negal_152_clean_all'
+
+def dataclean(example):
+    if example["audio"]["sampling_rate"] != 16000:
+        resampled_audio = librosa.resample(
+            y=example["audio"]["array"],
+            orig_sr=example["audio"]["sampling_rate"],
+            target_sr=16000,
+        )
+
+        return {
+            "audio": {
+                "path": example["audio"]["path"],
+                "array": resampled_audio,
+                "sampling_rate": 16000,
+            },
+            "transcription": example["transcription"]
+            .upper()
+            .translate(str.maketrans("", "", string.punctuation)),
+        }
+    else:
+        return {
+            "transcription": example["transcription"]
+            .upper()
+            .translate(str.maketrans("", "", string.punctuation))
+        }
+
+P1tony_dataset = load_dataset("audiofolder", data_dir=P1tony, split="train")
+P1tony_dataset = P1tony_dataset.map(dataclean)
+
+P1tony_scripted1 = load_dataset(
+    "audiofolder", data_dir=P1tony_rainbow, split="train"
+)
+P1tony_scripted2 = load_dataset(
+    "audiofolder", data_dir=P1tony_arthur, split="train"
+)
+P1tony_scripted1 = P1tony_scripted1.map(dataclean)
+P1tony_scripted2 = P1tony_scripted2.map(dataclean)
+P1tony_scripted = concatenate_datasets([P1tony_scripted1, P1tony_scripted2])
+
+class ChangeSampleRate(nn.Module):
+    def __init__(self, input_rate: int, output_rate: int):
+        super().__init__()
+        self.output_rate = output_rate
+        self.input_rate = input_rate
+
+    def forward(self, wav: torch.tensor) -> torch.tensor:
+        # Only accepts 1-channel waveform input
+        wav = wav.view(wav.size(0), -1)
+        new_length = wav.size(-1) * self.output_rate // self.input_rate
+        indices = torch.arange(new_length) * (
+            self.input_rate / self.output_rate
+        )
+        round_down = wav[:, indices.long()]
+        round_up = wav[:, (indices.long() + 1).clamp(max=wav.size(-1) - 1)]
+        output = round_down * (1.0 - indices.fmod(1.0)).unsqueeze(
+            0
+        ) + round_up * indices.fmod(1.0).unsqueeze(0)
+        return output
+
+# resample and clean text data
+def dataclean(example):
+    # pdb.set_trace()
+    if example["audio"]["sampling_rate"] != 16000:
+        resampled_audio = librosa.resample(
+            y=example["audio"]["array"],
+            orig_sr=example["audio"]["sampling_rate"],
+            target_sr=16000,
+        )
+        
+        return {
+            "audio": {
+                "path": example["audio"]["path"],
+                "array": resampled_audio,
+                "sampling_rate": 16000,
+            },
+            "transcription": example["transcription"]
+            .upper()
+            .translate(str.maketrans("", "", string.punctuation)),
+        }
+    else:
+        return {
+            "transcription": example["transcription"]
+            .upper()
+            .translate(str.maketrans("", "", string.punctuation))
+        }
+
+
+# processor = AutoFeatureExtractor.from_pretrained(
+#     "facebook/wav2vec2-base-960h"
+# )
+processor = AutoProcessor.from_pretrained("facebook/wav2vec2-base-960h")
+
+def prepare_dataset(batch):
+    audio = batch["audio"]
+    batch = processor(
+        audio["array"],
+        sampling_rate=audio["sampling_rate"],
+        text=batch["transcription"],
+    )
+    batch["input_length"] = len(batch["input_values"][0])
+    return batch
+
+src_dataset = load_dataset("audiofolder", data_dir=audio_dir, split="train")
+src_dataset = src_dataset.map(dataclean)
+p326_300_dataset = load_dataset(
+    "audiofolder", data_dir=p326_300_dir, split="train"
+)
+p326_300_dataset = p326_300_dataset.map(dataclean)
+
+P4Negel_dataset = load_dataset("audiofolder", data_dir=P4Negel, split="train")
+P4Negel_dataset = P4Negel_dataset.map(dataclean)
+
+healthy_test_dataset = load_dataset(
+    "audiofolder", data_dir=healthy_dir, split="train"
+)
+healthy_test_dataset = healthy_test_dataset.map(dataclean)
+
+Fary_PAL_test_dataset = load_dataset(
+    "audiofolder", data_dir=Fary_PAL_30, split="train"
+)
+Fary_PAL_test_dataset = Fary_PAL_test_dataset.map(dataclean)
+
+John_p326_test_dataset = load_dataset(
+    "audiofolder", data_dir=John_p326, split="train"
+)
+John_p326_test_dataset = John_p326_test_dataset.map(dataclean)
+
+John_video_test_dataset = load_dataset(
+    "audiofolder", data_dir=John_video, split="train"
+)
+John_video_test_dataset = John_video_test_dataset.map(dataclean)
+
+
+def train_dev_test_split(
+    dataset: Dataset, dev_rate=0.1, test_rate=0.1, seed=1
+):
+    """
+    input: dataset
+    dev_rate,
+    test_rate
+    seed
+    -------
+    Output:
+    dataset_dict{"train", "dev", "test"}
+    """
+    train_dev_test = dataset.train_test_split(test_size=test_rate, seed=seed)
+    test = train_dev_test["test"]
+    train_dev = train_dev_test["train"]
+
+    # pdb.set_trace()
+    if len(train_dev) <= int(len(dataset) * dev_rate):
+        train = Dataset.from_dict({"audio": [], "transcription": []})
+        dev = train_dev
+    else:
+        train_dev = train_dev.train_test_split(
+            test_size=int(len(dataset) * dev_rate), seed=seed
+        )
+        train = train_dev["train"]
+        dev = train_dev["test"]
+    return train, dev, test
+
+P1tony_train, P1tony_dev, P1tony_test = train_dev_test_split(
+    P1tony_dataset, dev_rate=0.5, test_rate=0.5, seed=1
+)
+P1tony_train_ = concatenate_datasets([P1tony_train, P1tony_scripted])
+
+# train_dev / test
+ds = src_dataset.train_test_split(test_size=0.1, seed=1)
+
+# dataset_libri = load_dataset(
+#     "hf-internal-testing/librispeech_asr_dummy", "clean", split="validation"
+# )
+
+train_dev = ds["train"]
+# train / dev
+train_dev = train_dev.train_test_split(
+    test_size=int(len(src_dataset) * 0.1), seed=1
+)
+
+# Tony
+Tony_train = P1tony_train_
+Tony_dev = P1tony_dev
+Tony_test = P1tony_test
+
+# John
+John_train, John_dev, John_test = train_dev_test_split(p326_300_dataset, dev_rate=0.1, test_rate=0.1)
+# Negel
+Negel_train, Negel_dev, Negel_test = train_dev_test_split(P4Negel_dataset, dev_rate=0.1, test_rate=0.1)
+
+# train/dev/test
+train = train_dev["train"]
+test = ds["test"]
+dev = train_dev["test"]
+
+# combined
+combine_train = concatenate_datasets([train, Tony_train, John_train, Negel_train])
+conbine_dev = concatenate_datasets([dev, Tony_dev, John_dev, Negel_dev])
+conbine_test = concatenate_datasets([test, Tony_test, John_test, Negel_test])
+
+# encoded_train = combine_train.map(prepare_dataset, num_proc=4)
+# encoded_dev = conbine_dev.map(prepare_dataset, num_proc=4)
+# encoded_test = conbine_test.map(prepare_dataset, num_proc=4)
+
+# # extra_test
+# encoded_Fary = Fary_PAL_test_dataset.map(prepare_dataset, num_proc=4)
+# encoded_healthy = healthy_test_dataset.map(prepare_dataset, num_proc=4)
+
+# encoded_ori_test = test.map(prepare_dataset, num_proc=4)
+# encoded_Tony_test = Tony_test.map(prepare_dataset, num_proc=4)
+# encoded_John_test = John_test.map(prepare_dataset, num_proc=4)
+# encoded_Negel_test = Negel_test.map(prepare_dataset, num_proc=4)
+
+# encoded_train = train.map(prepare_dataset, num_proc=4)
+# encoded_dev = dev.map(prepare_dataset, num_proc=4)
+# p326_encoded_train = p326_300_dataset.map(prepare_dataset, num_proc=4)
+
+# combine large p326 in to training set
+# encoded_train = concatenate_datasets([encoded_train, p326_encoded_train])
+
+# encoded_John_p326 = John_p326_test_dataset.map(prepare_dataset, num_proc=4)
+# encoded_John_video = John_video_test_dataset.map(prepare_dataset, num_proc=4)
+
+# pdb.set_trace()
+import numpy as np
+
+WER = evaluate.load("wer")
+
+## Whisper decoding
+
+from transformers import (Seq2SeqTrainer, Seq2SeqTrainingArguments,
+                          WhisperFeatureExtractor,
+                          WhisperForConditionalGeneration, WhisperModel,
+                          WhisperProcessor, WhisperTokenizer)
+
+processor = WhisperProcessor.from_pretrained("openai/whisper-medium")
+model = WhisperForConditionalGeneration.from_pretrained(
+    fine_tuning_dir,
+).to("cuda:0")
+# model = WhisperForConditionalGeneration.from_pretrained(
+#     "openai/whisper-medium",
+# ).to("cuda:0")
+tokenizer = WhisperTokenizer.from_pretrained(
+    "openai/whisper-medium", language="English", task="transcribe"
+)
+
+from pathlib import Path
+
+id = Path(fine_tuning_dir).stem
+pdb.set_trace()
+# tokenizer.push_to_hub("KevinGeng/%s" % id)
+# import pdb
+feature_extractor = WhisperFeatureExtractor.from_pretrained(
+    "openai/whisper-medium"
+)
+
+def whisper_prepare_dataset(batch):
+    # load and resample audio data from 48 to 16kHz
+    audio = batch["audio"]
+
+    # compute log-Mel input features from input audio array
+    batch["input_features"] = feature_extractor(
+        audio["array"], sampling_rate=audio["sampling_rate"]
+    ).input_features[0]
+
+    # encode target text to label ids
+    batch["labels"] = tokenizer(batch["transcription"]).input_ids
+    return batch
+
+torch.cuda.empty_cache()
+
+
+def my_map_to_pred(batch):
+    # pdb.set_trace()
+    audio = batch["audio"]
+    input_features = processor(
+        audio["array"],
+        sampling_rate=audio["sampling_rate"],
+        return_tensors="pt",
+    ).input_features
+    # batch["reference"] = whisper_processor.tokenizer._normalize(batch['text'])
+    batch["reference"] = processor.tokenizer._normalize(batch["transcription"])
+
+    with torch.no_grad():
+        # predicted_ids = whisper_model.generate(input_features.to("cuda"))[0]
+        predicted_ids = model.generate(input_features.to("cuda"))[0]
+    transcription = model.decode(predicted_ids)
+    batch["prediction"] = model.tokenizer._normalize(transcription)
+    return batch
+
+
+from dataclasses import dataclass
+from typing import Any, Dict, List, Union
+
+import torch
+
+
+@dataclass
+class DataCollatorSpeechSeq2SeqWithPadding:
+    processor: Any
+
+    def __call__(
+        self, features: List[Dict[str, Union[List[int], torch.Tensor]]]
+    ) -> Dict[str, torch.Tensor]:
+        # split inputs and labels since they have to be of different lengths and need different padding methods
+        # first treat the audio inputs by simply returning torch tensors
+        input_features = [
+            {"input_features": feature["input_features"]}
+            for feature in features
+        ]
+        batch = self.processor.feature_extractor.pad(
+            input_features, return_tensors="pt"
+        )
+
+        # get the tokenized label sequences
+        label_features = [
+            {"input_ids": feature["labels"]} for feature in features
+        ]
+        # pad the labels to max length
+        labels_batch = self.processor.tokenizer.pad(
+            label_features, return_tensors="pt"
+        )
+
+        # replace padding with -100 to ignore loss correctly
+        labels = labels_batch["input_ids"].masked_fill(
+            labels_batch.attention_mask.ne(1), -100
+        )
+
+        # if bos token is appended in previous tokenization step,
+        # cut bos token here as it's append later anyways
+        if (
+            (labels[:, 0] == self.processor.tokenizer.bos_token_id)
+            .all()
+            .cpu()
+            .item()
+        ):
+            labels = labels[:, 1:]
+
+        batch["labels"] = labels
+
+        return batch
+
+
+data_collator = DataCollatorSpeechSeq2SeqWithPadding(processor=processor)
+
+
+def compute_metrics(pred):
+    pred_ids = pred.predictions
+    label_ids = pred.label_ids
+
+    # replace -100 with the pad_token_id
+    label_ids[label_ids == -100] = tokenizer.pad_token_id
+
+    # we do not want to group tokens when computing the metrics
+    pred_str = tokenizer.batch_decode(pred_ids, skip_special_tokens=True)
+    label_str = tokenizer.batch_decode(label_ids, skip_special_tokens=True)
+
+    wer = 100 * WER.compute(predictions=pred_str, references=label_str)
+
+    return {"wer": wer}
+
+encoded_train = combine_train.map(whisper_prepare_dataset, num_proc=4)
+encoded_dev = conbine_dev.map(whisper_prepare_dataset, num_proc=4)
+encoded_test = conbine_test.map(whisper_prepare_dataset, num_proc=4)
+
+# extra_test
+
+encoded_ori_test = test.map(whisper_prepare_dataset, num_proc=4) # 7 / 16 
+
+encoded_Tony_test = Tony_test.map(whisper_prepare_dataset, num_proc=4) # 0 / 19
+
+encoded_John_test = John_test.map(whisper_prepare_dataset, num_proc=4) # 0 / 30
+# [False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False]
+encoded_Negel_test = Negel_test.map(whisper_prepare_dataset, num_proc=4) # 12 / 16
+# [False, True, True, True, True, True, True, False, True, True, True, True, True, True, False, False]
+encoded_Fary = Fary_PAL_test_dataset.map(whisper_prepare_dataset, num_proc=4) # 12 / 30
+# [True, True, True, True, True, False, False, False, True, False, False, True, False, False, True, False, True, True, False, True, True, False, False, False, False, False, False, False, False, False]
+encoded_healthy = healthy_test_dataset.map(whisper_prepare_dataset, num_proc=4) # 5 / 160
+# [False, False, False, False, False, False, False, False, False, True, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, True, False, True, True, False, False, False, True, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False]
+
+# Make sure the content Variablity
+train_tuple = tuple(encoded_train['transcription'])
+dev_tuple = tuple(encoded_dev['transcription'])
+train_dev_tuple = (train_tuple + dev_tuple)
+
+pdb.set_trace()
+new_encoded_test = encoded_test.select(np.where(np.array([False if x in train_dev_tuple else True for x in encoded_test['transcription']]))[0])
+new_encoded_ori_test = encoded_ori_test.select(np.where(np.array([False if x in train_dev_tuple else True for x in encoded_ori_test['transcription']]))[0])
+new_encoded_Tony_test = encoded_Tony_test.select(np.where(np.array([False if x in train_dev_tuple else True for x in encoded_Tony_test['transcription']]))[0])
+new_encoded_John_test = encoded_John_test.select(np.where(np.array([False if x in train_dev_tuple else True for x in encoded_John_test['transcription']]))[0])
+new_encoded_Negel_test = encoded_Negel_test.select(np.where(np.array([False if x in train_dev_tuple else True for x in encoded_Negel_test['transcription']]))[0])
+new_encoded_Fary = encoded_Fary.select(np.where(np.array([False if x in train_dev_tuple else True for x in encoded_Fary['transcription']]))[0])
+new_encoded_healthy = encoded_healthy.select(np.where(np.array([False if x in train_dev_tuple else True for x in encoded_healthy['transcription']]))[0])
+pdb.set_trace()
+torch.cuda.empty_cache()
+
+training_args = Seq2SeqTrainingArguments(
+    output_dir=fine_tuning_dir,  # change to a repo name of your choice
+    per_device_train_batch_size=8,
+    gradient_accumulation_steps=1,  # increase by 2x for every 2x decrease in batch size
+    learning_rate=1e-5,
+    warmup_steps=50,
+    max_steps=1000,
+    gradient_checkpointing=True,
+    fp16=True,
+    evaluation_strategy="steps",
+    save_strategy="steps",
+    per_device_eval_batch_size=8,
+    predict_with_generate=True,
+    generation_max_length=512,
+    save_steps=20,
+    eval_steps=20,
+    logging_steps=10,
+    report_to=["tensorboard"],
+    load_best_model_at_end=True,
+    metric_for_best_model="wer",
+    greater_is_better=False,
+    save_total_limit=10,
+    push_to_hub=False,
+)
+from transformers import EarlyStoppingCallback, Trainer, TrainingArguments
+
+trainer = Seq2SeqTrainer(
+    args=training_args,
+    model=model,
+    train_dataset=encoded_train,
+    eval_dataset=encoded_dev,
+    data_collator=data_collator,
+    compute_metrics=compute_metrics,
+    tokenizer=processor.feature_extractor,
+    callbacks=[EarlyStoppingCallback(early_stopping_patience=10)],
+)
+# callbacks=[EvalLoggingCallback()]
+pdb.set_trace()
+
+result_dict = {
+    "Ori_Test": trainer.evaluate(encoded_ori_test),
+    "Tony_Test": trainer.evaluate(encoded_Tony_test),
+    "John_Test": trainer.evaluate(encoded_John_test),
+    "Negel_Test": trainer.evaluate(encoded_Negel_test),
+    "Zeroshot_Fary_Test": trainer.evaluate(encoded_Fary),
+    "Healthy_Test": trainer.evaluate(encoded_healthy),
+}
+
+# print(result_dict)
+
+pdb.set_trace()
+trainer.evaluate(encoded_test)
+trainer.evaluate(new_encoded_test)
+new_result_dict = {
+    "Ori_Test": trainer.evaluate(new_encoded_ori_test), # 'eval_wer': 12.345679012345679,
+    "Tony_Test": trainer.evaluate(new_encoded_Tony_test), #  'eval_wer': 25.0,
+    "John_Test": trainer.evaluate(new_encoded_John_test), 
+    "Negel_Test": trainer.evaluate(new_encoded_Negel_test), # 2.08
+    "Zeroshot_Fary_Test": trainer.evaluate(new_encoded_Fary), ## 11.49
+    "Healthy_Test": trainer.evaluate(new_encoded_healthy),
+}
+
+print(new_result_dict)
+
+# pdb.set_trace()
+# result_dict = {
+#     "Ori_Test": trainer.evaluate(encoded_ori_test),
+#     "Tony_Test": trainer.evaluate(encoded_Tony_test),
+#     "John_Test": trainer.evaluate(encoded_John_test),
+#     "Negel_Test": trainer.evaluate(encoded_Negel_test),
+#     "Zeroshot_Fary_Test": trainer.evaluate(encoded_Fary),
+#     "Healthy_Test": trainer.evaluate(encoded_healthy),
+# }
+
+# pdb.set_trace()
+# # Evaluation
+# model.push_to_hub("KevinGeng/%s" % id)

local/post_processing.py

@@ -0,0 +1,56 @@
+'''
+# Post processing module for data recording
+# Author: Kevin Geng @Laronix, Sep. 2022
+
+# Load log.csv, generate standard wav files with selected samplerate, and calculate stastitical features
+'''
+
+from random import sample
+import librosa
+import soundfile as sf
+import numpy as np
+import pdb
+from pathlib import Path
+import sys
+import pandas as pd
+indir = Path(sys.argv[1])
+assert indir.exists() == True 
+wavs = Path(indir/Path("Audio_to_Evaluate")).glob("**/*.wav")
+log = Path(indir/Path("log.csv"))
+
+# x = np.loadtxt(log, dtype=str, delimiter=",")
+x = pd.read_csv(log, header=0)
+
+# y, sr = librosa.load("/home/kevingeng/laronix_automos/Julianna/Audio_to_evaluate/tmp0kgcdpi2.wav", sr=48000)
+outdir = indir/Path("output")
+# pdb.set_trace()
+# outdir_clean = indir/Path("output_clean")
+Path.mkdir(outdir, exist_ok=True)
+# Path.mkdir(outdir_clean, exist_ok=True)
+## Capitalize E valuate
+# for i, j in zip(x["Audio_to_Evaluate"], x["Reference_ID"]):
+#     y, sr = librosa.load(i, sr=48000)
+#     # kevin 1017 John's trial with original data.
+#     y_ = librosa.util.normalize(y, norm=5)
+#     y_cut, index = librosa.effects.trim(y_, top_db=30)
+#     # normalized and cut
+#     # pdb.set_trace()
+#     # sf.write(outdir/Path(str(indir)+"_"+ j +".wav"), y_cut, samplerate=sr)
+#     sf.write(outdir/Path(Path(indir).stem+"_"+ j +".wav"), y_cut, samplerate=sr)
+
+def process_audio(file_path, ref_id, sr=48000, norm=5, top_db=30):
+    y, _ = librosa.load(file_path, sr=sr)
+    y_norm = librosa.util.normalize(y, norm=norm)
+    y_cut, _ = librosa.effects.trim(y_norm, top_db=top_db)
+    return y_cut
+
+def save_audio(y_cut, ref_id, outdir, indir, sr=48000):
+    out_path = outdir / f"{Path(indir).stem}_{ref_id}.wav"
+    sf.write(out_path, y_cut, samplerate=sr)
+
+def main(audio_files, ref_ids, outdir, indir):
+    for file_path, ref_id in zip(audio_files, ref_ids):
+        y_cut = process_audio(file_path, ref_id)
+        save_audio(y_cut, ref_id, outdir, indir)
+
+main(x["Audio_to_Evaluate"], x["Reference_ID"], outdir, indir)

local/wer_plot_report.py

@@ -0,0 +1,45 @@
+from pathlib import Path
+import pandas as pd
+
+import numpy as np
+
+import matplotlib.pyplot as plt
+import sys
+import pdb
+
+threshold = 0.3
+if __name__ == "__main__":
+    wer_csv = sys.argv[1] 
+    df = pd.read_csv(wer_csv)
+    fig, ax = plt.subplots(nrows=2, ncols=1, figsize=(25, 15))
+
+    # Hist for distribution
+    ax[0].set_xlabel("Word Error Rate")
+    ax[0].set_ylabel("Counts")
+    ax[0].set_xlim(left=0.0, right=df['wer'].max())
+    ax[0].hist(df['wer'], bins=50)
+    ax[0].axvline(x=threshold, color="r")
+    # plt.savefig("hist.png")
+    
+    # Line curve for each sentences
+    colors = ['green' if x < threshold else 'red' for x in df['wer']]
+
+    new_ids = [str(x).split('.')[0] for x in df['id']]    
+    ax[1].set_xlabel("IDs")
+    ax[1].set_ylabel("Word Error Rate")
+    ax[1].scatter(new_ids, df['wer'], c=colors, marker='o')
+    ax[1].vlines(new_ids, ymin=0, ymax=df['wer'], colors='grey', linestyle='dotted', label='Vertical Lines')
+    ax[1].axhline(y=threshold, xmin=0, xmax=len(new_ids), color='r')
+
+    # ax[0].axhline(y=threshold, color="black")
+
+    # for i, v in enumerate(df['wer']):
+        # plt.text(str(df['id'][i]).split('.')[0], -2, str(df['id'][i]), ha='center', fontsize=3)
+
+    ax[1].set_xticklabels(new_ids, rotation=90, fontsize=10)
+    ax[1].tick_params(axis='x', width=20)
+    # ax[1].set_xlim(10, len(df['id']) + 10)
+    plt.tight_layout()
+    pdb.set_trace()
+    # fig.savefig("%s/%s.png"%(Path(sys.argv[1]).parent, sys.argv[1].split('/')[-1]), format='png')
+    fig.savefig("%s.png"%(sys.argv[1]), format='png')

local/whisper_fine_tuning_large_with_negel.py

@@ -0,0 +1,442 @@
+fine_tuning_dir = "/home/kevingeng/Disk2/laronix/laronix_automos/fine_tuned/SSD/model/Michael_52_with_Large_AVA_script_conv_train_conv_dev/checkpoint-60"
+"""
+TODO:
+    + [x] Load Configuration
+    + [ ] Multi ASR Engine
+    + [ ] Batch / Real Time support
+"""
+from pathlib import Path
+from transformers import AutoTokenizer, AutoFeatureExtractor, AutoModelForCTC, AutoProcessor
+from datasets import load_dataset, concatenate_datasets
+from datasets import Dataset, Audio
+import pdb
+import string
+import librosa
+# local import
+import sys
+
+sys.path.append("src")
+import torch
+torch.cuda.set_device("cuda:0")
+# token_model = AutoModelForCTC.from_pretrained(
+#     "facebook/wav2vec2-base-960h"
+# )
+
+# audio_dir= "/Users/kevingeng/Laronix/laronix_automos/data/Patient_sil_trim_16k_normed_5_snr_40/"
+audio_dir ="./data/Patient_sil_trim_16k_normed_5_snr_40"
+healthy_dir="./data/Healthy"
+Fary_PAL_30="./data/Fary_PAL_p326_20230110_30"
+John_p326 = "./data/John_p326/output"
+John_video = "./data/20230103_video"
+p326_300_dir ="./data/John_p326_large"
+negel_152 = "./data/4_negal_152_clean_all"
+
+michael3_52 = "data/3_michael_20230619_52"
+
+patient_T = "data/Patient_T/Patient_T"
+patient_L = "data/Patient_L/Patient_L"
+P1tony = "data/Participant1_Tony_Recording/CLEAN_SENTENCES/CONVERSATIONAL/PAL"
+P1tony_arthur = "data/Participant1_Tony_Recording/CLEAN_SENTENCES/SCRIPTED/Arthur_the_Rat/PAL"
+P1tony_rainbow = "data/Participant1_Tony_Recording/CLEAN_SENTENCES/SCRIPTED/Rainbow_Passage/Laronix"
+
+def dataclean(example):
+    # pdb.set_trace()
+    if example['audio']['sampling_rate'] != 16000:
+        resampled_audio = librosa.resample(y=example['audio']['array'],
+                                     orig_sr= example['audio']['sampling_rate'],
+                                     target_sr=16000)
+        # torchaudio.transforms.Resample(example['audio']['sampling_rate'], 16000)
+        # resampled_audio = resampler(example['audio']['array'])
+
+        return {"audio": {"path": example['audio']['path'], "array": resampled_audio, "sampling_rate": 16000},
+                "transcription": example["transcription"].upper().translate(str.maketrans('', '', string.punctuation))}
+    else:
+        return {"transcription": example["transcription"].upper().translate(str.maketrans('', '', string.punctuation))}
+
+# patient_L_test_dataset = load_dataset("audiofolder", data_dir=patient_L, split="train")
+# patient_L_test_dataset = patient_L_test_dataset.map(dataclean)
+
+# patient_T_test_dataset = load_dataset("audiofolder", data_dir=patient_T, split="train")
+# patient_T_test_dataset = patient_T_test_dataset.map(dataclean)
+
+P1tony_dataset = load_dataset("audiofolder", data_dir=P1tony, split="train")
+P1tony_dataset = P1tony_dataset.map(dataclean)
+
+P3Micheal_dataset_52 = load_dataset("audiofolder", data_dir=michael3_52, split="train")
+P3Micheal_dataset_52 = P3Micheal_dataset_52.map(dataclean)
+
+# negel_152_dataset = load_dataset("audiofolder", data_dir=negel_152, split="train")
+# negel_152_dataset = negel_152_dataset.map(dataclean)
+
+
+# pdb.set_trace()
+# P1tony_scripted1 = load_dataset("audiofolder", data_dir=P1tony_rainbow, split="train")
+# P1tony_scripted2 = load_dataset("audiofolder", data_dir=P1tony_arthur, split="train")
+# P1tony_scripted1 = P1tony_scripted1.map(dataclean)
+# P1tony_scripted2 = P1tony_scripted2.map(dataclean)
+# P1tony_scripted = concatenate_datasets([P1tony_scripted1, P1tony_scripted2])
+
+# audio_dir ="/home/kevingeng/laronix/laronix_automos/data/Healthy"
+# tgt_audio_dir= "/Users/kevingeng/Laronix/Dataset/Pneumatic/automos"
+
+# Get Transcription, WER and PPM 
+"""
+TODO:
+    [DONE]: Automatic generating Config
+"""
+
+import yaml
+import argparse
+import sys
+from pathlib import Path
+
+sys.path.append("./src")
+import lightning_module
+from UV import plot_UV, get_speech_interval
+from transformers import pipeline
+from rich.progress import track
+from rich import print as rprint
+import numpy as np
+import jiwer
+import pdb
+import torch.nn as nn
+import torch
+import torchaudio
+import gradio as gr
+from sys import flags
+from random import sample
+from transformers import Wav2Vec2Processor, Wav2Vec2ForCTC
+
+import evaluate
+
+wer = evaluate.load("wer")
+
+# root_path = Path(__file__).parents[1]
+
+class ChangeSampleRate(nn.Module):
+    def __init__(self, input_rate: int, output_rate: int):
+        super().__init__()
+        self.output_rate = output_rate
+        self.input_rate = input_rate
+
+    def forward(self, wav: torch.tensor) -> torch.tensor:
+        # Only accepts 1-channel waveform input
+        wav = wav.view(wav.size(0), -1) 
+        new_length = wav.size(-1) * self.output_rate // self.input_rate
+        indices = torch.arange(new_length) * (
+            self.input_rate / self.output_rate
+        )
+        round_down = wav[:, indices.long()]
+        round_up = wav[:, (indices.long() + 1).clamp(max=wav.size(-1) - 1)]
+        output = round_down * (1.0 - indices.fmod(1.0)).unsqueeze(
+            0
+        ) + round_up * indices.fmod(1.0).unsqueeze(0)
+        return output
+
+# resample and clean text data
+def dataclean(example):
+    # pdb.set_trace()
+    if example['audio']['sampling_rate'] != 16000:
+        resampled_audio = librosa.resample(y=example['audio']['array'],
+                                     orig_sr= example['audio']['sampling_rate'],
+                                     target_sr=16000)
+        # torchaudio.transforms.Resample(example['audio']['sampling_rate'], 16000)
+        # resampled_audio = resampler(example['audio']['array'])
+
+        return {"audio": {"path": example['audio']['path'], "array": resampled_audio, "sampling_rate": 16000},
+                "transcription": example["transcription"].upper().translate(str.maketrans('', '', string.punctuation))}
+    else:
+        return {"transcription": example["transcription"].upper().translate(str.maketrans('', '', string.punctuation))}
+
+# processor = AutoFeatureExtractor.from_pretrained(
+#     "facebook/wav2vec2-base-960h"
+# )
+processor = AutoProcessor.from_pretrained("facebook/wav2vec2-base-960h")
+
+def prepare_dataset(batch):
+    audio = batch["audio"]
+    batch = processor(audio["array"], sampling_rate = audio["sampling_rate"], text=batch['transcription'])
+    batch["input_length"] = len(batch["input_values"][0])
+    return batch
+
+src_dataset = load_dataset("audiofolder", data_dir=audio_dir, split="train")
+src_dataset = src_dataset.map(dataclean)
+p326_300_dataset = load_dataset("audiofolder", data_dir=p326_300_dir, split="train")
+p326_300_dataset = p326_300_dataset.map(dataclean)
+
+# healthy_test_dataset = load_dataset("audiofolder", data_dir=healthy_dir, split='train')
+# healthy_test_dataset = healthy_test_dataset.map(dataclean)
+
+# Fary_PAL_test_dataset = load_dataset("audiofolder", data_dir=Fary_PAL_30, split='train')
+# Fary_PAL_test_dataset = Fary_PAL_test_dataset.map(dataclean)
+
+# John_p326_test_dataset = load_dataset("audiofolder", data_dir=John_p326, split='train')
+# John_p326_test_dataset = John_p326_test_dataset.map(dataclean)
+
+# John_video_test_dataset = load_dataset("audiofolder", data_dir=John_video, split='train')
+# John_video_test_dataset = John_video_test_dataset.map(dataclean)
+
+# pdb.set_trace()
+
+def train_dev_test_split(dataset: Dataset, dev_rate=0.1, test_rate=0.1, seed=1):
+    """
+    input: dataset
+    dev_rate,
+    test_rate
+    seed
+    -------
+    Output:
+    dataset_dict{"train", "dev", "test"}
+    """
+    train_dev_test = dataset.train_test_split(test_size=test_rate, seed=seed)
+    test = train_dev_test["test"]
+    train_dev = train_dev_test['train']
+    
+    # pdb.set_trace()
+    if len(train_dev) <= int(len(dataset)*dev_rate):
+        train = Dataset.from_dict({"audio": [], "transcription": []})
+        dev = train_dev
+    else:
+        train_dev = train_dev.train_test_split(test_size=int(len(dataset)*dev_rate), seed=seed)
+        train = train_dev['train']
+        dev = train_dev['test']
+    return train, dev, test
+
+# pdb.set_trace()
+# P1tony_train, P1tony_dev, P1tony_test = train_dev_test_split(P1tony_dataset, dev_rate=0.5, test_rate=0.5, seed=1)
+# P1tony_train_ = concatenate_datasets([P1tony_train,P1tony_scripted]) 
+# pdb.set_trace()
+
+Michael_52_train, Michael_52_dev,  Michael_52_test = train_dev_test_split(P3Micheal_dataset_52, dev_rate=0.1.5, test_rate=0.15, seed=1)
+ 
+# train_dev / test
+ds = src_dataset.train_test_split(test_size=0.1, seed=1)
+
+# dataset_libri = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation")
+
+train_dev = ds['train']
+# train / dev
+train_dev = train_dev.train_test_split(test_size=int(len(src_dataset)*0.1), seed=1)
+# train/dev/test
+train = train_dev['train']
+test = ds['test']
+dev = train_dev['test']
+
+encoded_train = train.map(prepare_dataset, num_proc=4)
+encoded_dev = dev.map(prepare_dataset, num_proc=4)
+encoded_test = test.map(prepare_dataset, num_proc=4)
+p326_encoded_train = p326_300_dataset.map(prepare_dataset, num_proc=4)
+
+# combine large p326 in to training set
+encoded_train = concatenate_datasets([encoded_train, p326_encoded_train])
+
+# encoded_healthy = healthy_test_dataset.map(prepare_dataset, num_proc=4)
+# encoded_Fary = Fary_PAL_test_dataset.map(prepare_dataset, num_proc=4)
+# encoded_John_p326 = John_p326_test_dataset.map(prepare_dataset, num_proc=4)
+# encoded_John_video = John_video_test_dataset.map(prepare_dataset, num_proc=4)
+
+# encoded_P1tony_train = P1tony_train.map(prepare_dataset, num_proc=4)
+# encoded_P1tony_dev = P1tony_dev.map(prepare_dataset, num_proc=4)
+# encoded_P1tony_test = P1tony_test.map(prepare_dataset, num_proc=4)
+
+# pdb.set_trace()
+import numpy as np
+
+WER = evaluate.load("wer")
+
+## Whisper decoding
+
+from transformers import WhisperProcessor, WhisperForConditionalGeneration, WhisperTokenizer, WhisperFeatureExtractor, Seq2SeqTrainingArguments, Seq2SeqTrainer, WhisperModel
+processor = WhisperProcessor.from_pretrained("openai/whisper-medium")
+# model = WhisperForConditionalGeneration.from_pretrained("openai/whisper-medium").to("cuda:0")
+model = WhisperForConditionalGeneration.from_pretrained("./fine_tuned/whipser_medium_en_PAL300_step25_step2_VCTK/checkpoint-400", use_auth_token=True).to("cuda:0")
+tokenizer = WhisperTokenizer.from_pretrained("openai/whisper-medium", language="English", task="transcribe")
+
+from pathlib import Path
+id = Path(fine_tuning_dir).stem
+pdb.set_trace()
+tokenizer.push_to_hub("KevinGeng/%s"%id)
+# import pdb
+feature_extractor = WhisperFeatureExtractor.from_pretrained("openai/whisper-medium")
+
+def whisper_prepare_dataset(batch):
+    # load and resample audio data from 48 to 16kHz
+    audio = batch["audio"]
+
+    # compute log-Mel input features from input audio array 
+    batch["input_features"] = feature_extractor(audio["array"], sampling_rate=audio["sampling_rate"]).input_features[0]
+
+    # encode target text to label ids 
+    batch["labels"] = tokenizer(batch["transcription"]).input_ids
+    return batch
+
+torch.cuda.empty_cache()
+
+def my_map_to_pred(batch):
+    # pdb.set_trace()
+    audio = batch["audio"]
+    input_features = processor(audio["array"], sampling_rate=audio["sampling_rate"], return_tensors="pt").input_features
+    # batch["reference"] = whisper_processor.tokenizer._normalize(batch['text'])
+    batch["reference"] = processor.tokenizer._normalize(batch['transcription'])
+
+    with torch.no_grad():
+        # predicted_ids = whisper_model.generate(input_features.to("cuda"))[0]
+        predicted_ids = model.generate(input_features.to("cuda"))[0]
+    transcription = model.decode(predicted_ids)
+    batch["prediction"] = model.tokenizer._normalize(transcription)
+    return batch
+
+import torch
+
+from dataclasses import dataclass
+from typing import Any, Dict, List, Union
+
+@dataclass
+class DataCollatorSpeechSeq2SeqWithPadding:
+    processor: Any
+
+    def __call__(self, features: List[Dict[str, Union[List[int], torch.Tensor]]]) -> Dict[str, torch.Tensor]:
+        # split inputs and labels since they have to be of different lengths and need different padding methods
+        # first treat the audio inputs by simply returning torch tensors
+        input_features = [{"input_features": feature["input_features"]} for feature in features]
+        batch = self.processor.feature_extractor.pad(input_features, return_tensors="pt")
+
+        # get the tokenized label sequences
+        label_features = [{"input_ids": feature["labels"]} for feature in features]
+        # pad the labels to max length
+        labels_batch = self.processor.tokenizer.pad(label_features, return_tensors="pt")
+
+        # replace padding with -100 to ignore loss correctly
+        labels = labels_batch["input_ids"].masked_fill(labels_batch.attention_mask.ne(1), -100)
+
+        # if bos token is appended in previous tokenization step,
+        # cut bos token here as it's append later anyways
+        if (labels[:, 0] == self.processor.tokenizer.bos_token_id).all().cpu().item():
+            labels = labels[:, 1:]
+
+        batch["labels"] = labels
+
+        return batch
+    
+data_collator = DataCollatorSpeechSeq2SeqWithPadding(processor=processor)
+
+def compute_metrics(pred):
+    pred_ids = pred.predictions
+    label_ids = pred.label_ids
+
+    # replace -100 with the pad_token_id
+    label_ids[label_ids == -100] = tokenizer.pad_token_id
+
+    # we do not want to group tokens when computing the metrics
+    pred_str = tokenizer.batch_decode(pred_ids, skip_special_tokens=True)
+    label_str = tokenizer.batch_decode(label_ids, skip_special_tokens=True)
+
+    wer = 100 * WER.compute(predictions=pred_str, references=label_str)
+
+    return {"wer": wer}
+
+# whisper_train = train.map(whisper_prepare_dataset, num_proc=4)
+# pdb.set_trace()
+whisper_train_large = encoded_train.map(whisper_prepare_dataset, num_proc=4)
+whisper_dev = dev.map(whisper_prepare_dataset, num_proc=4)
+whisper_test = test.map(whisper_prepare_dataset, num_proc=4)
+
+encoded_Michael_52_train = Michael_52_train.map(whisper_prepare_dataset, num_proc=4)
+encoded_Michael_52_dev = Michael_52_dev.map(whisper_prepare_dataset, num_proc=4)
+encoded_Michael_52_test = Michael_52_test.map(whisper_prepare_dataset, num_proc=4)
+# pdb.set_trace()
+# # Add scirtped tony
+# encoded_P1tony_train = P1tony_train_.map(whisper_prepare_dataset, num_proc=4)
+# encoded_P1tony_dev = P1tony_dev.map(whisper_prepare_dataset, num_proc=4)
+# encoded_P1tony_test = P1tony_test.map(whisper_prepare_dataset, num_proc=4)
+
+# encode_negel_152_train = negel_152_train.map(whisper_prepare_dataset, num_proc=4)
+# encode_negel_152_dev = negel_152_dev.map(whisper_prepare_dataset, num_proc=4)
+# encode_negel_152_test = negel_152_test.map(whisper_prepare_dataset, num_proc=4)
+
+# encoded_train_large = concatenate_datasets([whisper_train_large, encode_negel_152_train])
+# encoded_dev_large = concatenate_datasets([whisper_dev, encode_negel_152_dev])
+
+pdb.set_trace()
+torch.cuda.empty_cache()
+
+training_args = Seq2SeqTrainingArguments(
+    output_dir=fine_tuning_dir,  # change to a repo name of your choice
+    per_device_train_batch_size=8,
+    gradient_accumulation_steps=1,  # increase by 2x for every 2x decrease in batch size
+    learning_rate=1e-5,
+    warmup_steps=50,
+    max_steps=1000,
+    gradient_checkpointing=True,
+    fp16=True,
+    evaluation_strategy="steps",
+    save_strategy="steps",
+    per_device_eval_batch_size=8,
+    predict_with_generate=True,
+    generation_max_length=512,
+    save_steps=10,
+    eval_steps=10,
+    logging_steps=10,
+    report_to=["tensorboard"],
+    load_best_model_at_end=True,
+    metric_for_best_model="wer",
+    greater_is_better=False,
+    save_total_limit=5,
+    push_to_hub=False,
+)
+from transformers import Trainer, TrainingArguments, EarlyStoppingCallback
+
+# pdb.set_trace()
+# # from transformers.trainer.callbacks import TensorBoardCallback
+# class EvalLoggingCallback(TrainerCallback):
+#     def on_evaluate(self, args, state, control, metrics, **kwargs):
+#         print(f"Eval loss: {metrics['eval_loss']:.4f}, Accuracy: {metrics['eval_wer']:.4f}")
+        
+# pdb.set_trace()
+
+trainer = Seq2SeqTrainer(
+    args=training_args,
+    model=model,
+    train_dataset=encoded_Michael_52_train,
+    eval_dataset=encoded_Michael_52_dev,
+    data_collator=data_collator,
+    compute_metrics=compute_metrics,
+    tokenizer=processor.feature_extractor,
+    callbacks=[EarlyStoppingCallback(early_stopping_patience=5)],
+
+)
+    # callbacks=[EvalLoggingCallback()]
+trainer.train()
+# trainer.evaluate(encoded_P1tony_test, metrix_key_prefix="test")
+# trainer.callback_handler.on_test_end(trainer, datasets=encoded_P1tony_test)
+
+
+# ## Not fine tuned
+# z_result = encoded_test.map(my_map_to_pred)·
+# # pdb.set_trace()
+# # 0.4692737430167598
+# z = WER.compute(references=z_result['reference'], predictions=z_result['prediction'])
+
+# z_hel_result = encoded_healthy.map(my_map_to_pred)
+# # 
+# z_hel = WER.compute(references=z_hel_result['reference'], predictions=z_hel_result['prediction'])
+# # 0.1591610117211598
+
+# z_fary_result = encoded_Fary.map(my_map_to_pred)
+# z_far = WER.compute(references=z_fary_result['reference'], predictions=z_fary_result['prediction'])
+# # 0.1791044776119403
+
+
+# z_john_p326_result = encoded_John_p326.map(my_map_to_pred)
+# z_john_p326 = WER.compute(references=z_john_p326_result['reference'], predictions=z_john_p326_result['prediction'])
+# # 0.4648241206030151
+
+# # y_John_video= fine_tuned_trainer.predict(encoded_John_video)
+# # metrics={'test_loss': 2.665189743041992, 'test_wer': 0.7222222222222222, 'test_runtime': 0.1633, 'test_samples_per_second': 48.979, 'test_steps_per_second': 6.122})
+# pdb.set_trace()
+
+# p326 training
+# metrics={'test_loss': 0.4804028868675232, 'test_wer': 0.21787709497206703, 'test_runtime': 0.3594, 'test_samples_per_second': 44.517, 'test_steps_per_second': 5.565})
+# hel metrics={'test_loss': 1.6363693475723267, 'test_wer': 0.17951881554595928, 'test_runtime': 3.8451, 'test_samples_per_second': 41.611, 'test_steps_per_second': 5.201})
+# Fary: metrics={'test_loss': 1.4633615016937256, 'test_wer': 0.5572139303482587, 'test_runtime': 0.6627, 'test_samples_per_second': 45.27, 'test_steps_per_second': 6.036})
+# p326 large: metrics={'test_loss': 0.6568527817726135, 'test_wer': 0.2889447236180904, 'test_runtime': 0.7169, 'test_samples_per_second': 51.613, 'test_steps_per_second': 6.975}) 

local/whisper_fine_tuning_michael_100.py

@@ -0,0 +1,442 @@
+fine_tuning_dir = "fine_tuned/SSD/model/Michael_100_with_Large_AVA_script_conv_train_conv_dev"
+"""
+TODO:
+    + [x] Load Configuration
+    + [ ] Multi ASR Engine
+    + [ ] Batch / Real Time support
+"""
+from pathlib import Path
+from transformers import AutoTokenizer, AutoFeatureExtractor, AutoModelForCTC, AutoProcessor
+from datasets import load_dataset, concatenate_datasets
+from datasets import Dataset, Audio
+import pdb
+import string
+import librosa
+# local import
+import sys
+
+sys.path.append("src")
+import torch
+torch.cuda.set_device("cuda:0")
+# token_model = AutoModelForCTC.from_pretrained(
+#     "facebook/wav2vec2-base-960h"
+# )
+
+# audio_dir= "/Users/kevingeng/Laronix/laronix_automos/data/Patient_sil_trim_16k_normed_5_snr_40/"
+audio_dir ="./data/Patient_sil_trim_16k_normed_5_snr_40"
+healthy_dir="./data/Healthy"
+Fary_PAL_30="./data/Fary_PAL_p326_20230110_30"
+John_p326 = "./data/John_p326/output"
+John_video = "./data/20230103_video"
+p326_300_dir ="./data/John_p326_large"
+negel_152 = "./data/4_negal_152_clean_all"
+
+michael3_52 = "data/3_michael_20230619_100"
+
+patient_T = "data/Patient_T/Patient_T"
+patient_L = "data/Patient_L/Patient_L"
+P1tony = "data/Participant1_Tony_Recording/CLEAN_SENTENCES/CONVERSATIONAL/PAL"
+P1tony_arthur = "data/Participant1_Tony_Recording/CLEAN_SENTENCES/SCRIPTED/Arthur_the_Rat/PAL"
+P1tony_rainbow = "data/Participant1_Tony_Recording/CLEAN_SENTENCES/SCRIPTED/Rainbow_Passage/Laronix"
+
+def dataclean(example):
+    # pdb.set_trace()
+    if example['audio']['sampling_rate'] != 16000:
+        resampled_audio = librosa.resample(y=example['audio']['array'],
+                                     orig_sr= example['audio']['sampling_rate'],
+                                     target_sr=16000)
+        # torchaudio.transforms.Resample(example['audio']['sampling_rate'], 16000)
+        # resampled_audio = resampler(example['audio']['array'])
+
+        return {"audio": {"path": example['audio']['path'], "array": resampled_audio, "sampling_rate": 16000},
+                "transcription": example["transcription"].upper().translate(str.maketrans('', '', string.punctuation))}
+    else:
+        return {"transcription": example["transcription"].upper().translate(str.maketrans('', '', string.punctuation))}
+
+# patient_L_test_dataset = load_dataset("audiofolder", data_dir=patient_L, split="train")
+# patient_L_test_dataset = patient_L_test_dataset.map(dataclean)
+
+# patient_T_test_dataset = load_dataset("audiofolder", data_dir=patient_T, split="train")
+# patient_T_test_dataset = patient_T_test_dataset.map(dataclean)
+
+P1tony_dataset = load_dataset("audiofolder", data_dir=P1tony, split="train")
+P1tony_dataset = P1tony_dataset.map(dataclean)
+
+P3Micheal_dataset_52 = load_dataset("audiofolder", data_dir=michael3_52, split="train")
+P3Micheal_dataset_52 = P3Micheal_dataset_52.map(dataclean)
+
+# negel_152_dataset = load_dataset("audiofolder", data_dir=negel_152, split="train")
+# negel_152_dataset = negel_152_dataset.map(dataclean)
+
+
+# pdb.set_trace()
+# P1tony_scripted1 = load_dataset("audiofolder", data_dir=P1tony_rainbow, split="train")
+# P1tony_scripted2 = load_dataset("audiofolder", data_dir=P1tony_arthur, split="train")
+# P1tony_scripted1 = P1tony_scripted1.map(dataclean)
+# P1tony_scripted2 = P1tony_scripted2.map(dataclean)
+# P1tony_scripted = concatenate_datasets([P1tony_scripted1, P1tony_scripted2])
+
+# audio_dir ="/home/kevingeng/laronix/laronix_automos/data/Healthy"
+# tgt_audio_dir= "/Users/kevingeng/Laronix/Dataset/Pneumatic/automos"
+
+# Get Transcription, WER and PPM 
+"""
+TODO:
+    [DONE]: Automatic generating Config
+"""
+
+import yaml
+import argparse
+import sys
+from pathlib import Path
+
+sys.path.append("./src")
+import lightning_module
+# from UV import plot_UV, get_speech_interval
+from transformers import pipeline
+from rich.progress import track
+from rich import print as rprint
+import numpy as np
+import jiwer
+import pdb
+import torch.nn as nn
+import torch
+import torchaudio
+import gradio as gr
+from sys import flags
+from random import sample
+from transformers import Wav2Vec2Processor, Wav2Vec2ForCTC
+
+import evaluate
+
+wer = evaluate.load("wer")
+
+# root_path = Path(__file__).parents[1]
+
+class ChangeSampleRate(nn.Module):
+    def __init__(self, input_rate: int, output_rate: int):
+        super().__init__()
+        self.output_rate = output_rate
+        self.input_rate = input_rate
+
+    def forward(self, wav: torch.tensor) -> torch.tensor:
+        # Only accepts 1-channel waveform input
+        wav = wav.view(wav.size(0), -1) 
+        new_length = wav.size(-1) * self.output_rate // self.input_rate
+        indices = torch.arange(new_length) * (
+            self.input_rate / self.output_rate
+        )
+        round_down = wav[:, indices.long()]
+        round_up = wav[:, (indices.long() + 1).clamp(max=wav.size(-1) - 1)]
+        output = round_down * (1.0 - indices.fmod(1.0)).unsqueeze(
+            0
+        ) + round_up * indices.fmod(1.0).unsqueeze(0)
+        return output
+
+# resample and clean text data
+def dataclean(example):
+    # pdb.set_trace()
+    if example['audio']['sampling_rate'] != 16000:
+        resampled_audio = librosa.resample(y=example['audio']['array'],
+                                     orig_sr= example['audio']['sampling_rate'],
+                                     target_sr=16000)
+        # torchaudio.transforms.Resample(example['audio']['sampling_rate'], 16000)
+        # resampled_audio = resampler(example['audio']['array'])
+
+        return {"audio": {"path": example['audio']['path'], "array": resampled_audio, "sampling_rate": 16000},
+                "transcription": example["transcription"].upper().translate(str.maketrans('', '', string.punctuation))}
+    else:
+        return {"transcription": example["transcription"].upper().translate(str.maketrans('', '', string.punctuation))}
+
+# processor = AutoFeatureExtractor.from_pretrained(
+#     "facebook/wav2vec2-base-960h"
+# )
+processor = AutoProcessor.from_pretrained("facebook/wav2vec2-base-960h")
+
+def prepare_dataset(batch):
+    audio = batch["audio"]
+    batch = processor(audio["array"], sampling_rate = audio["sampling_rate"], text=batch['transcription'])
+    batch["input_length"] = len(batch["input_values"][0])
+    return batch
+
+src_dataset = load_dataset("audiofolder", data_dir=audio_dir, split="train")
+src_dataset = src_dataset.map(dataclean)
+p326_300_dataset = load_dataset("audiofolder", data_dir=p326_300_dir, split="train")
+p326_300_dataset = p326_300_dataset.map(dataclean)
+
+# healthy_test_dataset = load_dataset("audiofolder", data_dir=healthy_dir, split='train')
+# healthy_test_dataset = healthy_test_dataset.map(dataclean)
+
+# Fary_PAL_test_dataset = load_dataset("audiofolder", data_dir=Fary_PAL_30, split='train')
+# Fary_PAL_test_dataset = Fary_PAL_test_dataset.map(dataclean)
+
+# John_p326_test_dataset = load_dataset("audiofolder", data_dir=John_p326, split='train')
+# John_p326_test_dataset = John_p326_test_dataset.map(dataclean)
+
+# John_video_test_dataset = load_dataset("audiofolder", data_dir=John_video, split='train')
+# John_video_test_dataset = John_video_test_dataset.map(dataclean)
+
+# pdb.set_trace()
+
+def train_dev_test_split(dataset: Dataset, dev_rate=0.1, test_rate=0.1, seed=1):
+    """
+    input: dataset
+    dev_rate,
+    test_rate
+    seed
+    -------
+    Output:
+    dataset_dict{"train", "dev", "test"}
+    """
+    train_dev_test = dataset.train_test_split(test_size=test_rate, seed=seed)
+    test = train_dev_test["test"]
+    train_dev = train_dev_test['train']
+    
+    # pdb.set_trace()
+    if len(train_dev) <= int(len(dataset)*dev_rate):
+        train = Dataset.from_dict({"audio": [], "transcription": []})
+        dev = train_dev
+    else:
+        train_dev = train_dev.train_test_split(test_size=int(len(dataset)*dev_rate), seed=seed)
+        train = train_dev['train']
+        dev = train_dev['test']
+    return train, dev, test
+
+# pdb.set_trace()
+# P1tony_train, P1tony_dev, P1tony_test = train_dev_test_split(P1tony_dataset, dev_rate=0.5, test_rate=0.5, seed=1)
+# P1tony_train_ = concatenate_datasets([P1tony_train,P1tony_scripted]) 
+# pdb.set_trace()
+
+Michael_52_train, Michael_52_dev,  Michael_52_test = train_dev_test_split(P3Micheal_dataset_52, dev_rate=0.10, test_rate=0.1, seed=1)
+ 
+# train_dev / test
+ds = src_dataset.train_test_split(test_size=0.1, seed=1)
+
+# dataset_libri = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation")
+
+train_dev = ds['train']
+# train / dev
+train_dev = train_dev.train_test_split(test_size=int(len(src_dataset)*0.1), seed=1)
+# train/dev/test
+train = train_dev['train']
+test = ds['test']
+dev = train_dev['test']
+
+encoded_train = train.map(prepare_dataset, num_proc=4)
+encoded_dev = dev.map(prepare_dataset, num_proc=4)
+encoded_test = test.map(prepare_dataset, num_proc=4)
+p326_encoded_train = p326_300_dataset.map(prepare_dataset, num_proc=4)
+
+# combine large p326 in to training set
+encoded_train = concatenate_datasets([encoded_train, p326_encoded_train])
+
+# encoded_healthy = healthy_test_dataset.map(prepare_dataset, num_proc=4)
+# encoded_Fary = Fary_PAL_test_dataset.map(prepare_dataset, num_proc=4)
+# encoded_John_p326 = John_p326_test_dataset.map(prepare_dataset, num_proc=4)
+# encoded_John_video = John_video_test_dataset.map(prepare_dataset, num_proc=4)
+
+# encoded_P1tony_train = P1tony_train.map(prepare_dataset, num_proc=4)
+# encoded_P1tony_dev = P1tony_dev.map(prepare_dataset, num_proc=4)
+# encoded_P1tony_test = P1tony_test.map(prepare_dataset, num_proc=4)
+
+# pdb.set_trace()
+import numpy as np
+
+WER = evaluate.load("wer")
+
+## Whisper decoding
+
+from transformers import WhisperProcessor, WhisperForConditionalGeneration, WhisperTokenizer, WhisperFeatureExtractor, Seq2SeqTrainingArguments, Seq2SeqTrainer, WhisperModel
+processor = WhisperProcessor.from_pretrained("openai/whisper-medium")
+# model = WhisperForConditionalGeneration.from_pretrained("openai/whisper-medium").to("cuda:0")
+model = WhisperForConditionalGeneration.from_pretrained("./fine_tuned/whipser_medium_en_PAL300_step25_step2_VCTK/checkpoint-400", use_auth_token=True).to("cuda:0")
+tokenizer = WhisperTokenizer.from_pretrained("openai/whisper-medium", language="English", task="transcribe")
+
+from pathlib import Path
+id = Path(fine_tuning_dir).stem
+pdb.set_trace()
+tokenizer.push_to_hub("KevinGeng/%s"%id)
+# import pdb
+feature_extractor = WhisperFeatureExtractor.from_pretrained("openai/whisper-medium")
+
+def whisper_prepare_dataset(batch):
+    # load and resample audio data from 48 to 16kHz
+    audio = batch["audio"]
+
+    # compute log-Mel input features from input audio array 
+    batch["input_features"] = feature_extractor(audio["array"], sampling_rate=audio["sampling_rate"]).input_features[0]
+
+    # encode target text to label ids 
+    batch["labels"] = tokenizer(batch["transcription"]).input_ids
+    return batch
+
+torch.cuda.empty_cache()
+
+def my_map_to_pred(batch):
+    # pdb.set_trace()
+    audio = batch["audio"]
+    input_features = processor(audio["array"], sampling_rate=audio["sampling_rate"], return_tensors="pt").input_features
+    # batch["reference"] = whisper_processor.tokenizer._normalize(batch['text'])
+    batch["reference"] = processor.tokenizer._normalize(batch['transcription'])
+
+    with torch.no_grad():
+        # predicted_ids = whisper_model.generate(input_features.to("cuda"))[0]
+        predicted_ids = model.generate(input_features.to("cuda"))[0]
+    transcription = model.decode(predicted_ids)
+    batch["prediction"] = model.tokenizer._normalize(transcription)
+    return batch
+
+import torch
+
+from dataclasses import dataclass
+from typing import Any, Dict, List, Union
+
+@dataclass
+class DataCollatorSpeechSeq2SeqWithPadding:
+    processor: Any
+
+    def __call__(self, features: List[Dict[str, Union[List[int], torch.Tensor]]]) -> Dict[str, torch.Tensor]:
+        # split inputs and labels since they have to be of different lengths and need different padding methods
+        # first treat the audio inputs by simply returning torch tensors
+        input_features = [{"input_features": feature["input_features"]} for feature in features]
+        batch = self.processor.feature_extractor.pad(input_features, return_tensors="pt")
+
+        # get the tokenized label sequences
+        label_features = [{"input_ids": feature["labels"]} for feature in features]
+        # pad the labels to max length
+        labels_batch = self.processor.tokenizer.pad(label_features, return_tensors="pt")
+
+        # replace padding with -100 to ignore loss correctly
+        labels = labels_batch["input_ids"].masked_fill(labels_batch.attention_mask.ne(1), -100)
+
+        # if bos token is appended in previous tokenization step,
+        # cut bos token here as it's append later anyways
+        if (labels[:, 0] == self.processor.tokenizer.bos_token_id).all().cpu().item():
+            labels = labels[:, 1:]
+
+        batch["labels"] = labels
+
+        return batch
+    
+data_collator = DataCollatorSpeechSeq2SeqWithPadding(processor=processor)
+
+def compute_metrics(pred):
+    pred_ids = pred.predictions
+    label_ids = pred.label_ids
+
+    # replace -100 with the pad_token_id
+    label_ids[label_ids == -100] = tokenizer.pad_token_id
+
+    # we do not want to group tokens when computing the metrics
+    pred_str = tokenizer.batch_decode(pred_ids, skip_special_tokens=True)
+    label_str = tokenizer.batch_decode(label_ids, skip_special_tokens=True)
+
+    wer = 100 * WER.compute(predictions=pred_str, references=label_str)
+
+    return {"wer": wer}
+
+# whisper_train = train.map(whisper_prepare_dataset, num_proc=4)
+# pdb.set_trace()
+whisper_train_large = encoded_train.map(whisper_prepare_dataset, num_proc=4)
+whisper_dev = dev.map(whisper_prepare_dataset, num_proc=4)
+whisper_test = test.map(whisper_prepare_dataset, num_proc=4)
+
+encoded_Michael_52_train = Michael_52_train.map(whisper_prepare_dataset, num_proc=4)
+encoded_Michael_52_dev = Michael_52_dev.map(whisper_prepare_dataset, num_proc=4)
+encoded_Michael_52_test = Michael_52_test.map(whisper_prepare_dataset, num_proc=4)
+# pdb.set_trace()
+# # Add scirtped tony
+# encoded_P1tony_train = P1tony_train_.map(whisper_prepare_dataset, num_proc=4)
+# encoded_P1tony_dev = P1tony_dev.map(whisper_prepare_dataset, num_proc=4)
+# encoded_P1tony_test = P1tony_test.map(whisper_prepare_dataset, num_proc=4)
+
+# encode_negel_152_train = negel_152_train.map(whisper_prepare_dataset, num_proc=4)
+# encode_negel_152_dev = negel_152_dev.map(whisper_prepare_dataset, num_proc=4)
+# encode_negel_152_test = negel_152_test.map(whisper_prepare_dataset, num_proc=4)
+
+# encoded_train_large = concatenate_datasets([whisper_train_large, encode_negel_152_train])
+# encoded_dev_large = concatenate_datasets([whisper_dev, encode_negel_152_dev])
+
+pdb.set_trace()
+torch.cuda.empty_cache()
+
+training_args = Seq2SeqTrainingArguments(
+    output_dir=fine_tuning_dir,  # change to a repo name of your choice
+    per_device_train_batch_size=8,
+    gradient_accumulation_steps=1,  # increase by 2x for every 2x decrease in batch size
+    learning_rate=1e-5,
+    warmup_steps=50,
+    max_steps=1000,
+    gradient_checkpointing=True,
+    fp16=True,
+    evaluation_strategy="steps",
+    save_strategy="steps",
+    per_device_eval_batch_size=8,
+    predict_with_generate=True,
+    generation_max_length=512,
+    save_steps=10,
+    eval_steps=10,
+    logging_steps=10,
+    report_to=["tensorboard"],
+    load_best_model_at_end=True,
+    metric_for_best_model="wer",
+    greater_is_better=False,
+    save_total_limit=5,
+    push_to_hub=False,
+)
+from transformers import Trainer, TrainingArguments, EarlyStoppingCallback
+
+# pdb.set_trace()
+# # from transformers.trainer.callbacks import TensorBoardCallback
+# class EvalLoggingCallback(TrainerCallback):
+#     def on_evaluate(self, args, state, control, metrics, **kwargs):
+#         print(f"Eval loss: {metrics['eval_loss']:.4f}, Accuracy: {metrics['eval_wer']:.4f}")
+        
+# pdb.set_trace()
+
+trainer = Seq2SeqTrainer(
+    args=training_args,
+    model=model,
+    train_dataset=encoded_Michael_52_train,
+    eval_dataset=encoded_Michael_52_dev,
+    data_collator=data_collator,
+    compute_metrics=compute_metrics,
+    tokenizer=processor.feature_extractor,
+    callbacks=[EarlyStoppingCallback(early_stopping_patience=10)],
+
+)
+    # callbacks=[EvalLoggingCallback()]
+trainer.train()
+# trainer.evaluate(encoded_P1tony_test, metrix_key_prefix="test")
+# trainer.callback_handler.on_test_end(trainer, datasets=encoded_P1tony_test)
+
+
+# ## Not fine tuned
+# z_result = encoded_test.map(my_map_to_pred)·
+# # pdb.set_trace()
+# # 0.4692737430167598
+# z = WER.compute(references=z_result['reference'], predictions=z_result['prediction'])
+
+# z_hel_result = encoded_healthy.map(my_map_to_pred)
+# # 
+# z_hel = WER.compute(references=z_hel_result['reference'], predictions=z_hel_result['prediction'])
+# # 0.1591610117211598
+
+# z_fary_result = encoded_Fary.map(my_map_to_pred)
+# z_far = WER.compute(references=z_fary_result['reference'], predictions=z_fary_result['prediction'])
+# # 0.1791044776119403
+
+
+# z_john_p326_result = encoded_John_p326.map(my_map_to_pred)
+# z_john_p326 = WER.compute(references=z_john_p326_result['reference'], predictions=z_john_p326_result['prediction'])
+# # 0.4648241206030151
+
+# # y_John_video= fine_tuned_trainer.predict(encoded_John_video)
+# # metrics={'test_loss': 2.665189743041992, 'test_wer': 0.7222222222222222, 'test_runtime': 0.1633, 'test_samples_per_second': 48.979, 'test_steps_per_second': 6.122})
+# pdb.set_trace()
+
+# p326 training
+# metrics={'test_loss': 0.4804028868675232, 'test_wer': 0.21787709497206703, 'test_runtime': 0.3594, 'test_samples_per_second': 44.517, 'test_steps_per_second': 5.565})
+# hel metrics={'test_loss': 1.6363693475723267, 'test_wer': 0.17951881554595928, 'test_runtime': 3.8451, 'test_samples_per_second': 41.611, 'test_steps_per_second': 5.201})
+# Fary: metrics={'test_loss': 1.4633615016937256, 'test_wer': 0.5572139303482587, 'test_runtime': 0.6627, 'test_samples_per_second': 45.27, 'test_steps_per_second': 6.036})
+# p326 large: metrics={'test_loss': 0.6568527817726135, 'test_wer': 0.2889447236180904, 'test_runtime': 0.7169, 'test_samples_per_second': 51.613, 'test_steps_per_second': 6.975}) 

local/whisper_fine_tuning_negel.py

@@ -0,0 +1,431 @@
+fine_tuning_dir = "fine_tuned/SSD/model/Negel_152_AVA_script_conv_train_conv_dev"
+"""
+TODO:
+    + [x] Load Configuration
+    + [ ] Multi ASR Engine
+    + [ ] Batch / Real Time support
+"""
+from pathlib import Path
+from transformers import AutoTokenizer, AutoFeatureExtractor, AutoModelForCTC, AutoProcessor
+from datasets import load_dataset, concatenate_datasets
+from datasets import Dataset, Audio
+import pdb
+import string
+import librosa
+# local import
+import sys
+
+sys.path.append("src")
+import torch
+torch.cuda.set_device("cuda:0")
+# token_model = AutoModelForCTC.from_pretrained(
+#     "facebook/wav2vec2-base-960h"
+# )
+
+# audio_dir= "/Users/kevingeng/Laronix/laronix_automos/data/Patient_sil_trim_16k_normed_5_snr_40/"
+audio_dir ="./data/Patient_sil_trim_16k_normed_5_snr_40"
+healthy_dir="./data/Healthy"
+Fary_PAL_30="./data/Fary_PAL_p326_20230110_30"
+John_p326 = "./data/John_p326/output"
+John_video = "./data/20230103_video"
+p326_300_dir ="./data/John_p326_large"
+
+negel_152 = "./data/4_negal_152_clean_all"
+
+patient_T = "data/Patient_T/Patient_T"
+patient_L = "data/Patient_L/Patient_L"
+P1tony = "data/Participant1_Tony_Recording/CLEAN_SENTENCES/CONVERSATIONAL/PAL"
+P1tony_arthur = "data/Participant1_Tony_Recording/CLEAN_SENTENCES/SCRIPTED/Arthur_the_Rat/PAL"
+P1tony_rainbow = "data/Participant1_Tony_Recording/CLEAN_SENTENCES/SCRIPTED/Rainbow_Passage/Laronix"
+
+def dataclean(example):
+    # pdb.set_trace()
+    if example['audio']['sampling_rate'] != 16000:
+        resampled_audio = librosa.resample(y=example['audio']['array'],
+                                     orig_sr= example['audio']['sampling_rate'],
+                                     target_sr=16000)
+        # torchaudio.transforms.Resample(example['audio']['sampling_rate'], 16000)
+        # resampled_audio = resampler(example['audio']['array'])
+
+        return {"audio": {"path": example['audio']['path'], "array": resampled_audio, "sampling_rate": 16000},
+                "transcription": example["transcription"].upper().translate(str.maketrans('', '', string.punctuation))}
+    else:
+        return {"transcription": example["transcription"].upper().translate(str.maketrans('', '', string.punctuation))}
+
+# patient_L_test_dataset = load_dataset("audiofolder", data_dir=patient_L, split="train")
+# patient_L_test_dataset = patient_L_test_dataset.map(dataclean)
+
+# patient_T_test_dataset = load_dataset("audiofolder", data_dir=patient_T, split="train")
+# patient_T_test_dataset = patient_T_test_dataset.map(dataclean)
+
+P1tony_dataset = load_dataset("audiofolder", data_dir=P1tony, split="train")
+P1tony_dataset = P1tony_dataset.map(dataclean)
+
+negel_152_dataset = load_dataset("audiofolder", data_dir=negel_152, split="train")
+negel_152_dataset = negel_152_dataset.map(dataclean)
+# pdb.set_trace()
+# P1tony_scripted1 = load_dataset("audiofolder", data_dir=P1tony_rainbow, split="train")
+# P1tony_scripted2 = load_dataset("audiofolder", data_dir=P1tony_arthur, split="train")
+# P1tony_scripted1 = P1tony_scripted1.map(dataclean)
+# P1tony_scripted2 = P1tony_scripted2.map(dataclean)
+# P1tony_scripted = concatenate_datasets([P1tony_scripted1, P1tony_scripted2])
+
+# audio_dir ="/home/kevingeng/laronix/laronix_automos/data/Healthy"
+# tgt_audio_dir= "/Users/kevingeng/Laronix/Dataset/Pneumatic/automos"
+
+# Get Transcription, WER and PPM 
+"""
+TODO:
+    [DONE]: Automatic generating Config
+"""
+
+import yaml
+import argparse
+import sys
+from pathlib import Path
+
+sys.path.append("./src")
+import lightning_module
+from UV import plot_UV, get_speech_interval
+from transformers import pipeline
+from rich.progress import track
+from rich import print as rprint
+import numpy as np
+import jiwer
+import pdb
+import torch.nn as nn
+import torch
+import torchaudio
+import gradio as gr
+from sys import flags
+from random import sample
+from transformers import Wav2Vec2Processor, Wav2Vec2ForCTC
+
+import evaluate
+
+wer = evaluate.load("wer")
+
+# root_path = Path(__file__).parents[1]
+
+class ChangeSampleRate(nn.Module):
+    def __init__(self, input_rate: int, output_rate: int):
+        super().__init__()
+        self.output_rate = output_rate
+        self.input_rate = input_rate
+
+    def forward(self, wav: torch.tensor) -> torch.tensor:
+        # Only accepts 1-channel waveform input
+        wav = wav.view(wav.size(0), -1) 
+        new_length = wav.size(-1) * self.output_rate // self.input_rate
+        indices = torch.arange(new_length) * (
+            self.input_rate / self.output_rate
+        )
+        round_down = wav[:, indices.long()]
+        round_up = wav[:, (indices.long() + 1).clamp(max=wav.size(-1) - 1)]
+        output = round_down * (1.0 - indices.fmod(1.0)).unsqueeze(
+            0
+        ) + round_up * indices.fmod(1.0).unsqueeze(0)
+        return output
+
+# resample and clean text data
+def dataclean(example):
+    # pdb.set_trace()
+    if example['audio']['sampling_rate'] != 16000:
+        resampled_audio = librosa.resample(y=example['audio']['array'],
+                                     orig_sr= example['audio']['sampling_rate'],
+                                     target_sr=16000)
+        # torchaudio.transforms.Resample(example['audio']['sampling_rate'], 16000)
+        # resampled_audio = resampler(example['audio']['array'])
+
+        return {"audio": {"path": example['audio']['path'], "array": resampled_audio, "sampling_rate": 16000},
+                "transcription": example["transcription"].upper().translate(str.maketrans('', '', string.punctuation))}
+    else:
+        return {"transcription": example["transcription"].upper().translate(str.maketrans('', '', string.punctuation))}
+
+# processor = AutoFeatureExtractor.from_pretrained(
+#     "facebook/wav2vec2-base-960h"
+# )
+processor = AutoProcessor.from_pretrained("facebook/wav2vec2-base-960h")
+
+def prepare_dataset(batch):
+    audio = batch["audio"]
+    batch = processor(audio["array"], sampling_rate = audio["sampling_rate"], text=batch['transcription'])
+    batch["input_length"] = len(batch["input_values"][0])
+    return batch
+
+# src_dataset = load_dataset("audiofolder", data_dir=audio_dir, split="train")
+# src_dataset = src_dataset.map(dataclean)
+# p326_300_dataset = load_dataset("audiofolder", data_dir=p326_300_dir, split="train")
+# p326_300_dataset = p326_300_dataset.map(dataclean)
+
+
+# healthy_test_dataset = load_dataset("audiofolder", data_dir=healthy_dir, split='train')
+# healthy_test_dataset = healthy_test_dataset.map(dataclean)
+
+# Fary_PAL_test_dataset = load_dataset("audiofolder", data_dir=Fary_PAL_30, split='train')
+# Fary_PAL_test_dataset = Fary_PAL_test_dataset.map(dataclean)
+
+# John_p326_test_dataset = load_dataset("audiofolder", data_dir=John_p326, split='train')
+# John_p326_test_dataset = John_p326_test_dataset.map(dataclean)
+
+# John_video_test_dataset = load_dataset("audiofolder", data_dir=John_video, split='train')
+# John_video_test_dataset = John_video_test_dataset.map(dataclean)
+
+
+
+# pdb.set_trace()
+
+def train_dev_test_split(dataset: Dataset, dev_rate=0.1, test_rate=0.1, seed=1):
+    """
+    input: dataset
+    dev_rate,
+    test_rate
+    seed
+    -------
+    Output:
+    dataset_dict{"train", "dev", "test"}
+    """
+    train_dev_test = dataset.train_test_split(test_size=test_rate, seed=seed)
+    test = train_dev_test["test"]
+    train_dev = train_dev_test['train']
+    
+    # pdb.set_trace()
+    if len(train_dev) <= int(len(dataset)*dev_rate):
+        train = Dataset.from_dict({"audio": [], "transcription": []})
+        dev = train_dev
+    else:
+        train_dev = train_dev.train_test_split(test_size=int(len(dataset)*dev_rate), seed=seed)
+        train = train_dev['train']
+        dev = train_dev['test']
+    return train, dev, test
+
+# pdb.set_trace()
+# P1tony_train, P1tony_dev, P1tony_test = train_dev_test_split(P1tony_dataset, dev_rate=0.5, test_rate=0.5, seed=1)
+# P1tony_train_ = concatenate_datasets([P1tony_train,P1tony_scripted]) 
+# pdb.set_trace()
+
+negel_152_train, negel_152_dev, negel_152_test = train_dev_test_split(negel_152_dataset, dev_rate=0.1, test_rate=0.1, seed=1)
+ 
+# train_dev / test
+# ds = src_dataset.train_test_split(test_size=0.1, seed=1)
+
+# dataset_libri = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation")
+
+# train_dev = ds['train']
+# # train / dev
+# train_dev = train_dev.train_test_split(test_size=int(len(src_dataset)*0.1), seed=1)
+# # train/dev/test
+# train = train_dev['train']
+# test = ds['test']
+# dev = train_dev['test']
+
+# encoded_train = train.map(prepare_dataset, num_proc=4)
+# encoded_dev = dev.map(prepare_dataset, num_proc=4)
+# encoded_test = test.map(prepare_dataset, num_proc=4)
+# p326_encoded_train = p326_300_dataset.map(prepare_dataset, num_proc=4)
+
+# # combine large p326 in to training set
+# # encoded_train = concatenate_datasets([encoded_train, p326_encoded_train])
+
+# encoded_healthy = healthy_test_dataset.map(prepare_dataset, num_proc=4)
+# encoded_Fary = Fary_PAL_test_dataset.map(prepare_dataset, num_proc=4)
+# encoded_John_p326 = John_p326_test_dataset.map(prepare_dataset, num_proc=4)
+# encoded_John_video = John_video_test_dataset.map(prepare_dataset, num_proc=4)
+
+# encoded_P1tony_train = P1tony_train.map(prepare_dataset, num_proc=4)
+# encoded_P1tony_dev = P1tony_dev.map(prepare_dataset, num_proc=4)
+# encoded_P1tony_test = P1tony_test.map(prepare_dataset, num_proc=4)
+
+# pdb.set_trace()
+import numpy as np
+
+WER = evaluate.load("wer")
+
+## Whisper decoding
+
+from transformers import WhisperProcessor, WhisperForConditionalGeneration, WhisperTokenizer, WhisperFeatureExtractor, Seq2SeqTrainingArguments, Seq2SeqTrainer, WhisperModel
+processor = WhisperProcessor.from_pretrained("openai/whisper-medium")
+# model = WhisperForConditionalGeneration.from_pretrained("openai/whisper-medium").to("cuda:0")
+model = WhisperForConditionalGeneration.from_pretrained("./fine_tuned/whipser_medium_en_PAL300_step25_step2_VCTK/checkpoint-400", use_auth_token=True).to("cuda:0")
+tokenizer = WhisperTokenizer.from_pretrained("openai/whisper-medium", language="English", task="transcribe")
+
+from pathlib import Path
+id = Path(fine_tuning_dir).stem
+# pdb.set_trace()
+tokenizer.push_to_hub("KevinGeng/%s"%id)
+# import pdb
+feature_extractor = WhisperFeatureExtractor.from_pretrained("openai/whisper-medium")
+
+def whisper_prepare_dataset(batch):
+    # load and resample audio data from 48 to 16kHz
+    audio = batch["audio"]
+
+    # compute log-Mel input features from input audio array 
+    batch["input_features"] = feature_extractor(audio["array"], sampling_rate=audio["sampling_rate"]).input_features[0]
+
+    # encode target text to label ids 
+    batch["labels"] = tokenizer(batch["transcription"]).input_ids
+    return batch
+
+torch.cuda.empty_cache()
+
+def my_map_to_pred(batch):
+    # pdb.set_trace()
+    audio = batch["audio"]
+    input_features = processor(audio["array"], sampling_rate=audio["sampling_rate"], return_tensors="pt").input_features
+    # batch["reference"] = whisper_processor.tokenizer._normalize(batch['text'])
+    batch["reference"] = processor.tokenizer._normalize(batch['transcription'])
+
+    with torch.no_grad():
+        # predicted_ids = whisper_model.generate(input_features.to("cuda"))[0]
+        predicted_ids = model.generate(input_features.to("cuda"))[0]
+    transcription = model.decode(predicted_ids)
+    batch["prediction"] = model.tokenizer._normalize(transcription)
+    return batch
+
+import torch
+
+from dataclasses import dataclass
+from typing import Any, Dict, List, Union
+
+@dataclass
+class DataCollatorSpeechSeq2SeqWithPadding:
+    processor: Any
+
+    def __call__(self, features: List[Dict[str, Union[List[int], torch.Tensor]]]) -> Dict[str, torch.Tensor]:
+        # split inputs and labels since they have to be of different lengths and need different padding methods
+        # first treat the audio inputs by simply returning torch tensors
+        input_features = [{"input_features": feature["input_features"]} for feature in features]
+        batch = self.processor.feature_extractor.pad(input_features, return_tensors="pt")
+
+        # get the tokenized label sequences
+        label_features = [{"input_ids": feature["labels"]} for feature in features]
+        # pad the labels to max length
+        labels_batch = self.processor.tokenizer.pad(label_features, return_tensors="pt")
+
+        # replace padding with -100 to ignore loss correctly
+        labels = labels_batch["input_ids"].masked_fill(labels_batch.attention_mask.ne(1), -100)
+
+        # if bos token is appended in previous tokenization step,
+        # cut bos token here as it's append later anyways
+        if (labels[:, 0] == self.processor.tokenizer.bos_token_id).all().cpu().item():
+            labels = labels[:, 1:]
+
+        batch["labels"] = labels
+
+        return batch
+    
+data_collator = DataCollatorSpeechSeq2SeqWithPadding(processor=processor)
+
+def compute_metrics(pred):
+    pred_ids = pred.predictions
+    label_ids = pred.label_ids
+
+    # replace -100 with the pad_token_id
+    label_ids[label_ids == -100] = tokenizer.pad_token_id
+
+    # we do not want to group tokens when computing the metrics
+    pred_str = tokenizer.batch_decode(pred_ids, skip_special_tokens=True)
+    label_str = tokenizer.batch_decode(label_ids, skip_special_tokens=True)
+
+    wer = 100 * WER.compute(predictions=pred_str, references=label_str)
+
+    return {"wer": wer}
+
+# whisper_train = train.map(whisper_prepare_dataset, num_proc=4)
+# pdb.set_trace()
+# whisper_train_large = encoded_train.map(whisper_prepare_dataset, num_proc=4)
+# whisper_dev = dev.map(whisper_prepare_dataset, num_proc=4)
+# whisper_test = test.map(whisper_prepare_dataset, num_proc=4)
+# pdb.set_trace()
+# # Add scirtped tony
+# encoded_P1tony_train = P1tony_train_.map(whisper_prepare_dataset, num_proc=4)
+# encoded_P1tony_dev = P1tony_dev.map(whisper_prepare_dataset, num_proc=4)
+# encoded_P1tony_test = P1tony_test.map(whisper_prepare_dataset, num_proc=4)
+
+encode_negel_152_train = negel_152_train.map(whisper_prepare_dataset, num_proc=4)
+encode_negel_152_dev = negel_152_dev.map(whisper_prepare_dataset, num_proc=4)
+encode_negel_152_test = negel_152_test.map(whisper_prepare_dataset, num_proc=4)
+pdb.set_trace()
+torch.cuda.empty_cache()
+
+training_args = Seq2SeqTrainingArguments(
+    output_dir=fine_tuning_dir,  # change to a repo name of your choice
+    per_device_train_batch_size=8,
+    gradient_accumulation_steps=1,  # increase by 2x for every 2x decrease in batch size
+    learning_rate=1e-5,
+    warmup_steps=50,
+    max_steps=1000,
+    gradient_checkpointing=True,
+    fp16=True,
+    evaluation_strategy="steps",
+    save_strategy="steps",
+    per_device_eval_batch_size=8,
+    predict_with_generate=True,
+    generation_max_length=512,
+    save_steps=10,
+    eval_steps=10,
+    logging_steps=10,
+    report_to=["tensorboard"],
+    load_best_model_at_end=True,
+    metric_for_best_model="wer",
+    greater_is_better=False,
+    save_total_limit=5,
+    push_to_hub=True,
+)
+from transformers import Trainer, TrainingArguments, EarlyStoppingCallback
+
+# pdb.set_trace()
+# # from transformers.trainer.callbacks import TensorBoardCallback
+# class EvalLoggingCallback(TrainerCallback):
+#     def on_evaluate(self, args, state, control, metrics, **kwargs):
+#         print(f"Eval loss: {metrics['eval_loss']:.4f}, Accuracy: {metrics['eval_wer']:.4f}")
+        
+# pdb.set_trace()
+
+trainer = Seq2SeqTrainer(
+    args=training_args,
+    model=model,
+    train_dataset=encode_negel_152_train,
+    eval_dataset=encode_negel_152_dev,
+    data_collator=data_collator,
+    compute_metrics=compute_metrics,
+    tokenizer=processor.feature_extractor,
+    callbacks=[EarlyStoppingCallback(early_stopping_patience=5)],
+
+)
+    # callbacks=[EvalLoggingCallback()]
+trainer.train()
+# trainer.evaluate(encoded_P1tony_test, metrix_key_prefix="test")
+# trainer.callback_handler.on_test_end(trainer, datasets=encoded_P1tony_test)
+
+
+# ## Not fine tuned
+# z_result = encoded_test.map(my_map_to_pred)·
+# # pdb.set_trace()
+# # 0.4692737430167598
+# z = WER.compute(references=z_result['reference'], predictions=z_result['prediction'])
+
+# z_hel_result = encoded_healthy.map(my_map_to_pred)
+# # 
+# z_hel = WER.compute(references=z_hel_result['reference'], predictions=z_hel_result['prediction'])
+# # 0.1591610117211598
+
+# z_fary_result = encoded_Fary.map(my_map_to_pred)
+# z_far = WER.compute(references=z_fary_result['reference'], predictions=z_fary_result['prediction'])
+# # 0.1791044776119403
+
+
+# z_john_p326_result = encoded_John_p326.map(my_map_to_pred)
+# z_john_p326 = WER.compute(references=z_john_p326_result['reference'], predictions=z_john_p326_result['prediction'])
+# # 0.4648241206030151
+
+# # y_John_video= fine_tuned_trainer.predict(encoded_John_video)
+# # metrics={'test_loss': 2.665189743041992, 'test_wer': 0.7222222222222222, 'test_runtime': 0.1633, 'test_samples_per_second': 48.979, 'test_steps_per_second': 6.122})
+# pdb.set_trace()
+
+# p326 training
+# metrics={'test_loss': 0.4804028868675232, 'test_wer': 0.21787709497206703, 'test_runtime': 0.3594, 'test_samples_per_second': 44.517, 'test_steps_per_second': 5.565})
+# hel metrics={'test_loss': 1.6363693475723267, 'test_wer': 0.17951881554595928, 'test_runtime': 3.8451, 'test_samples_per_second': 41.611, 'test_steps_per_second': 5.201})
+# Fary: metrics={'test_loss': 1.4633615016937256, 'test_wer': 0.5572139303482587, 'test_runtime': 0.6627, 'test_samples_per_second': 45.27, 'test_steps_per_second': 6.036})
+# p326 large: metrics={'test_loss': 0.6568527817726135, 'test_wer': 0.2889447236180904, 'test_runtime': 0.7169, 'test_samples_per_second': 51.613, 'test_steps_per_second': 6.975}) 

local/whisper_fine_tuning_negel_decode.py

@@ -0,0 +1,473 @@
+fine_tuning_dir = "fine_tuned/SSD/model/Michael_100_with_Large_AVA_script_conv_train_conv_dev/checkpoint-100"
+# fine_tuning_dir = "fine_tuned/SSD/model/Michael_52_with_Large_AVA_script_conv_train_conv_dev/checkpoint-60"
+# fine_tuning_dir = "fine_tuned/SSD/model/Negel_152_AVA_script_conv_train_conv_dev/checkpoint-100"
+# fine_tuning_dir = "fine_tuned/SSD/model/Tony1_AVA_script_conv_train_conv_dev/checkpoint-160"
+# fine_tuning_dir = "fine_tuned/SSD/model/Negel_with_Large_AVA_script_conv_train_conv_dev/checkpoint-210"
+
+"""
+TODO:
+    + [x] Whipser Fine Tuned Model Evalutation
+    + [ ] 
+    + [ ] Batch / Real Time support
+"""
+from typing import Any, Dict, List, Union
+from dataclasses import dataclass
+from transformers import Seq2SeqTrainer
+from transformers import WhisperProcessor, WhisperForConditionalGeneration, WhisperTokenizer, WhisperFeatureExtractor, Seq2SeqTrainingArguments, Seq2SeqTrainer, WhisperModel
+import evaluate
+from transformers import Wav2Vec2Processor, Wav2Vec2ForCTC
+from random import sample
+from sys import flags
+import gradio as gr
+import torchaudio
+import torch.nn as nn
+import jiwer
+import numpy as np
+from rich import print as rprint
+from rich.progress import track
+from transformers import pipeline
+import argparse
+import yaml
+import torch
+from pathlib import Path
+from transformers import AutoTokenizer, AutoFeatureExtractor, AutoModelForCTC, AutoProcessor
+from datasets import load_dataset, concatenate_datasets
+from datasets import Dataset, Audio
+import pdb
+import string
+import librosa
+# local import
+import sys
+
+sys.path.append("src")
+import lightning_module
+
+torch.cuda.set_device("cuda:0")
+
+audio_dir = "./data/Patient_sil_trim_16k_normed_5_snr_40"
+healthy_dir = "./data/Healthy"
+Fary_PAL_30 = "./data/Fary_PAL_p326_20230110_30"
+John_p326 = "./data/John_p326/output"
+John_video = "./data/20230103_video"
+negel_79 = "./data/4_negal_152_clean_all"
+negel_152 = "./data/4_negal_152_clean_all"
+P1tony = "data/Participant1_Tony_Recording/CLEAN_SENTENCES/CONVERSATIONAL/PAL"
+
+michael3_52 = "data/3_michael_20230619_100"
+
+patient_T = "data/Patient_T/Patient_T"
+patient_L = "data/Patient_L/Patient_L"
+# Get Transcription, WER and PPM
+"""
+TODO:
+    [DONE]: Automatic generating Config
+"""
+
+
+sys.path.append("./src")
+
+wer = evaluate.load("wer")
+
+# root_path = Path(__file__).parents[1]
+
+class ChangeSampleRate(nn.Module):
+    def __init__(self, input_rate: int, output_rate: int):
+        super().__init__()
+        self.output_rate = output_rate
+        self.input_rate = input_rate
+
+    def forward(self, wav: torch.tensor) -> torch.tensor:
+        # Only accepts 1-channel waveform input
+        wav = wav.view(wav.size(0), -1)
+        new_length = wav.size(-1) * self.output_rate // self.input_rate
+        indices = torch.arange(new_length) * (
+            self.input_rate / self.output_rate
+        )
+        round_down = wav[:, indices.long()]
+        round_up = wav[:, (indices.long() + 1).clamp(max=wav.size(-1) - 1)]
+        output = round_down * (1.0 - indices.fmod(1.0)).unsqueeze(
+            0
+        ) + round_up * indices.fmod(1.0).unsqueeze(0)
+        return output
+
+# resample and clean text data
+
+
+def dataclean(example):
+    # pdb.set_trace()
+    if example['audio']['sampling_rate'] != 16000:
+        resampled_audio = librosa.resample(y=example['audio']['array'],
+                                           orig_sr=example['audio']['sampling_rate'],
+                                           target_sr=16000)
+        # torchaudio.transforms.Resample(example['audio']['sampling_rate'], 16000)
+        # resampled_audio = resampler(example['audio']['array'])
+
+        return {"audio": {"path": example['audio']['path'], "array": resampled_audio, "sampling_rate": 16000},
+                "transcription": example["transcription"].upper().translate(str.maketrans('', '', string.punctuation))}
+    else:
+        return {"transcription": example["transcription"].upper().translate(str.maketrans('', '', string.punctuation))}
+
+processor = AutoFeatureExtractor.from_pretrained(
+    "facebook/wav2vec2-base-960h"
+)
+
+def prepare_dataset(batch):
+    audio = batch["audio"]
+    batch = processor(
+        audio["array"], sampling_rate=audio["sampling_rate"], text=batch['transcription'])
+    batch["input_length"] = len(batch["input_values"][0])
+    return batch
+
+
+negel_79_dataset = load_dataset("audiofolder", data_dir=negel_79, split="train")
+negel_79_dataset = negel_79_dataset.map(dataclean)
+
+def train_dev_test_split(dataset: Dataset, dev_rate=0.1, test_rate=0.1, seed=1):
+    """
+    input: dataset
+    dev_rate,
+    test_rate
+    seed
+    -------
+    Output:
+    dataset_dict{"train", "dev", "test"}
+    """
+    train_dev_test = dataset.train_test_split(test_size=test_rate, seed=seed)
+    test = train_dev_test["test"]
+    train_dev = train_dev_test['train']
+    
+    # pdb.set_trace()
+    if len(train_dev) <= int(len(dataset)*dev_rate):
+        train = Dataset.from_dict({"audio": [], "transcription": []})
+        dev = train_dev
+    else:
+        train_dev = train_dev.train_test_split(test_size=int(len(dataset)*dev_rate), seed=seed)
+        train = train_dev['train']
+        dev = train_dev['test']
+    return train, dev, test
+P1tony_dataset = load_dataset("audiofolder", data_dir=P1tony, split="train")
+P1tony_dataset = P1tony_dataset.map(dataclean)
+# pdb.set_trace()
+P1tony_train, P1tony_dev, P1tony_test = train_dev_test_split(P1tony_dataset, dev_rate=0.5, test_rate=0.5, seed=1)
+
+P3Micheal_dataset_52 = load_dataset("audiofolder", data_dir=michael3_52, split="train")
+P3Micheal_dataset_52 = P3Micheal_dataset_52.map(dataclean)
+
+Michael_52_train, Michael_52_dev,  Michael_52_test = train_dev_test_split(P3Micheal_dataset_52, dev_rate=0.1, test_rate=0.1, seed=1)
+
+# P1tony_train_ = concatenate_datasets([P1tony_train,P1tony_scripted]) 
+# pdb.set_trace()
+
+# Negel_79_train, Negel_79_dev, Negel_79_test = train_dev_test_split(negel_79_dataset, dev_rate=0.1, test_rate=0.1, seed=1)
+ 
+src_dataset = load_dataset("audiofolder", data_dir=audio_dir, split="train")
+src_dataset = src_dataset.map(dataclean)
+
+negel_152_dataset = load_dataset("audiofolder", data_dir=negel_152, split="train")
+negel_152_dataset = negel_152_dataset.map(dataclean)
+
+healthy_test_dataset = load_dataset(
+    "audiofolder", data_dir=healthy_dir, split='train')
+healthy_test_dataset = healthy_test_dataset.map(dataclean)
+
+Fary_PAL_test_dataset = load_dataset(
+    "audiofolder", data_dir=Fary_PAL_30, split='train')
+Fary_PAL_test_dataset = Fary_PAL_test_dataset.map(dataclean)
+
+John_p326_test_dataset = load_dataset(
+    "audiofolder", data_dir=John_p326, split='train')
+John_p326_test_dataset = John_p326_test_dataset.map(dataclean)
+
+John_video_test_dataset = load_dataset(
+    "audiofolder", data_dir=John_video, split='train')
+John_video_test_dataset = John_video_test_dataset.map(dataclean)
+
+patient_T_test_dataset = load_dataset("audiofolder", data_dir=patient_T, split='train')
+patient_T_test_dataset = patient_T_test_dataset.map(dataclean)
+
+patient_L_test_dataset = load_dataset("audiofolder", data_dir=patient_L, split='train')
+patient_L_test_dataset = patient_L_test_dataset.map(dataclean)
+pdb.set_trace()
+
+# train_dev / test
+ds = src_dataset.train_test_split(test_size=0.1, seed=1)
+
+dataset_libri = load_dataset(
+    "hf-internal-testing/librispeech_asr_dummy", "clean", split="validation")
+
+train_dev = ds['train']
+# train / dev
+train_dev = train_dev.train_test_split(
+    test_size=int(len(src_dataset)*0.1), seed=1)
+# train/dev/test
+train = train_dev['train']
+test = ds['test']
+dev = train_dev['test']
+
+# # pdb.set_trace()
+encoded_train = train.map(prepare_dataset, num_proc=4)
+encoded_dev = dev.map(prepare_dataset, num_proc=4)
+encoded_test = test.map(prepare_dataset, num_proc=4)
+
+encoded_Tony_test = P1tony_test.map(prepare_dataset, num_proc=4)
+encoded_healthy = healthy_test_dataset.map(prepare_dataset, num_proc=4)
+encoded_Fary = Fary_PAL_test_dataset.map(prepare_dataset, num_proc=4)
+encoded_John_p326 = John_p326_test_dataset.map(prepare_dataset, num_proc=4)
+encoded_John_video = John_video_test_dataset.map(prepare_dataset, num_proc=4)
+
+# pdb.set_trace()
+
+WER = evaluate.load("wer")
+
+# Whisper decoding
+
+processor = WhisperProcessor.from_pretrained("openai/whisper-medium")
+model = WhisperForConditionalGeneration.from_pretrained(
+    "openai/whisper-medium").to("cuda:0")
+tokenizer = WhisperTokenizer.from_pretrained(
+    "openai/whisper-medium", language="English", task="transcribe")
+
+# Need to push tokenizer to hugginface/model to activate online API
+
+# tokenizer.push_to_hub("KevinGeng/whipser_medium_en_PAL300_step25")
+# import pdb
+# pdb.set_trace()
+
+feature_extractor = WhisperFeatureExtractor.from_pretrained(
+    "openai/whisper-medium")
+
+
+def whisper_prepare_dataset(batch):
+    # load and resample audio data from 48 to 16kHz
+    audio = batch["audio"]
+
+    # compute log-Mel input features from input audio array
+    batch["input_features"] = feature_extractor(
+        audio["array"], sampling_rate=audio["sampling_rate"]).input_features[0]
+
+    # encode target text to label ids
+    batch["labels"] = tokenizer(batch["transcription"]).input_ids
+    return batch
+
+
+torch.cuda.empty_cache()
+
+training_args = Seq2SeqTrainingArguments(
+    # change to a repo name of your choice
+    output_dir="./whisper-medium-PAL128-25step",
+    per_device_train_batch_size=8,
+    gradient_accumulation_steps=1,  # increase by 2x for every 2x decrease in batch size
+    learning_rate=1e-5,
+    warmup_steps=100,
+    max_steps=1000,
+    gradient_checkpointing=True,
+    fp16=True,
+    evaluation_strategy="steps",
+    per_device_eval_batch_size=8,
+    predict_with_generate=True,
+    generation_max_length=512,
+    save_steps=100,
+    eval_steps=25,
+    logging_steps=100,
+    report_to=["tensorboard"],
+    load_best_model_at_end=True,
+    metric_for_best_model="wer",
+    greater_is_better=False,
+    push_to_hub=True,
+)
+
+
+def my_map_to_pred(batch):
+    # pdb.set_trace()
+    audio = batch["audio"]
+    input_features = processor(
+        audio["array"], sampling_rate=audio["sampling_rate"], return_tensors="pt").input_features
+    # batch["reference"] = whisper_processor.tokenizer._normalize(batch['text'])
+    batch["reference"] = processor.tokenizer._normalize(batch['transcription'])
+
+    with torch.no_grad():
+        # predicted_ids = whisper_model.generate(input_features.to("cuda"))[0]
+        predicted_ids = model.generate(input_features.to("cuda"))[0]
+    transcription = model.decode(predicted_ids)
+    batch["prediction"] = model.tokenizer._normalize(transcription)
+    return batch
+
+
+@dataclass
+class DataCollatorSpeechSeq2SeqWithPadding:
+    processor: Any
+
+    def __call__(self, features: List[Dict[str, Union[List[int], torch.Tensor]]]) -> Dict[str, torch.Tensor]:
+        # split inputs and labels since they have to be of different lengths and need different padding methods
+        # first treat the audio inputs by simply returning torch tensors
+        input_features = [{"input_features": feature["input_features"]}
+                          for feature in features]
+        batch = self.processor.feature_extractor.pad(
+            input_features, return_tensors="pt")
+
+        # get the tokenized label sequences
+        label_features = [{"input_ids": feature["labels"]}
+                          for feature in features]
+        # pad the labels to max length
+        labels_batch = self.processor.tokenizer.pad(
+            label_features, return_tensors="pt")
+
+        # replace padding with -100 to ignore loss correctly
+        labels = labels_batch["input_ids"].masked_fill(
+            labels_batch.attention_mask.ne(1), -100)
+
+        # if bos token is appended in previous tokenization step,
+        # cut bos token here as it's append later anyways
+        if (labels[:, 0] == self.processor.tokenizer.bos_token_id).all().cpu().item():
+            labels = labels[:, 1:]
+
+        batch["labels"] = labels
+
+        return batch
+
+
+data_collator = DataCollatorSpeechSeq2SeqWithPadding(processor=processor)
+
+
+def compute_metrics(pred):
+    pdb.set_trace()
+    pred_ids = pred.predictions
+    label_ids = pred.label_ids
+
+    # replace -100 with the pad_token_id
+    label_ids[label_ids == -100] = tokenizer.pad_token_id
+
+    # we do not want to group tokens when computing the metrics
+    pred_str = tokenizer.batch_decode(pred_ids, skip_special_tokens=True)
+    label_str = tokenizer.batch_decode(label_ids, skip_special_tokens=True)
+
+    wer = 100 * WER.compute(predictions=pred_str, references=label_str)
+
+    return {"wer": wer}
+
+# encode_negel_79_train = Negel_79_train.map(whisper_prepare_dataset, num_proc=4)
+# encode_negel_79_dev = Negel_79_dev.map(whisper_prepare_dataset, num_proc=4)
+# encode_negel_79_test = Negel_79_test.map(whisper_prepare_dataset, num_proc=4)
+whisper_test = test.map(whisper_prepare_dataset, num_proc=4)
+
+encoded_Michael_52_train = Michael_52_train.map(whisper_prepare_dataset, num_proc=4)
+encoded_Michael_52_dev = Michael_52_dev.map(whisper_prepare_dataset, num_proc=4)
+encoded_Michael_52_test = Michael_52_test.map(whisper_prepare_dataset, num_proc=4)
+# negel_152_train, negel_152_dev, negel_152_test = train_dev_test_split(negel_152_dataset, dev_rate=0.1, test_rate=0.1, seed=1)
+# encoded_negel_152_test = negel_152_test.map(whisper_prepare_dataset, num_proc=4)
+pdb.set_trace()
+torch.cuda.empty_cache()
+
+torch.cuda.empty_cache()
+
+fine_tuned_model = WhisperForConditionalGeneration.from_pretrained(
+    fine_tuning_dir
+).to("cuda")
+    # "fine_tuned/SSD/model/whipser_medium_TEP_patient_T"
+    # "./fine_tuned/whipser_medium_en_PAL300_step25_step2_VCTK/checkpoint-400"
+    #"./fine_tuned/whipser_medium_en_PAL300_step25_step2_VCTK/checkpoint-200"
+
+
+def fine_tuned_map_to_pred(batch):
+    # pdb.set_trace()
+    audio = batch["audio"]
+    input_features = processor(
+        audio["array"], sampling_rate=audio["sampling_rate"], return_tensors="pt").input_features
+    # batch["reference"] = whisper_processor.tokenizer._normalize(batch['text'])
+    batch["reference"] = processor.tokenizer._normalize(batch['transcription'])
+
+    with torch.no_grad():
+        # predicted_ids = whisper_model.generate(input_features.to("cuda"))[0]
+        predicted_ids = fine_tuned_model.generate(input_features.to("cuda"))[0]
+    transcription = tokenizer.decode(predicted_ids)
+    batch["prediction"] = tokenizer._normalize(transcription)
+    return batch
+
+
+    # output_dir="./fine_tuned/whipser_medium_en_PAL300_step25_step2_VCTK/checkpoint-400",
+    # output_dir="fine_tuned/SSD/model/whipser_medium_TEP_patient_TL_TL",
+testing_args = Seq2SeqTrainingArguments(
+    # change to a repo name of your choice
+    output_dir=fine_tuning_dir,
+    per_device_train_batch_size=8,
+    gradient_accumulation_steps=1,  # increase by 2x for every 2x decrease in batch size
+    learning_rate=1e-5,
+    warmup_steps=100,
+    max_steps=1000,
+    gradient_checkpointing=True,
+    fp16=True,
+    evaluation_strategy="steps",
+    per_device_eval_batch_size=8,
+    predict_with_generate=True,
+    generation_max_length=512,
+    save_steps=100,
+    eval_steps=25,
+    logging_steps=100,
+    report_to=["tensorboard"],
+    load_best_model_at_end=True,
+    metric_for_best_model="wer",
+    greater_is_better=False,
+    push_to_hub=False,
+)
+
+predict_trainer = Seq2SeqTrainer(
+    args=testing_args,
+    model=fine_tuned_model,
+    data_collator=data_collator,
+    compute_metrics=compute_metrics,
+    tokenizer=processor.feature_extractor,
+)
+
+# trainer.train()
+# fine tuned
+# z_result = encoded_test.map(fine_tuned_map_to_pred)
+
+pdb.set_trace()
+encoded_Michael_test_result = encoded_Michael_52_test.map(fine_tuned_map_to_pred)
+z_M = WER.compute(references=encoded_Michael_test_result['reference'], predictions=encoded_Michael_test_result['prediction'])
+pdb.set_trace()
+encoded_Tony_test_result = encoded_Tony_test.map(fine_tuned_map_to_pred)
+z = WER.compute(references=encoded_Tony_test_result['reference'], predictions=encoded_Tony_test_result['prediction'])
+pdb.set_trace()
+
+z_result= test.map(fine_tuned_map_to_pred)
+# 0.4692737430167598
+z = WER.compute(references=z_result['reference'], predictions=z_result['prediction'])
+# pdb.set_trace()
+z_hel_result = encoded_healthy.map(fine_tuned_map_to_pred)
+z_hel = WER.compute(references=z_hel_result['reference'], predictions=z_hel_result['prediction'])
+# 0.1591610117211598
+
+# encoded_negel_152_test
+# encoded_negel_test_result = encoded_negel_152_test.map(fine_tuned_map_to_pred)
+# z_negel = WER.compute(references=encoded_negel_test_result['reference'], predictions=encoded_negel_test_result['prediction'])
+
+pdb.set_trace()
+z_fary_result = encoded_Fary.map(fine_tuned_map_to_pred)
+z_far = WER.compute(references=z_fary_result['reference'], predictions=z_fary_result['prediction'])
+# 0.1791044776119403
+# z_patient_LT = encoded_patient_TL_test.map(fine_tuned_map_to_pred)
+# z_patient_LT_result = WER.compute(references=z_patient_LT['reference'], predictions=z_patient_LT['prediction'])
+# z_patient_L = encoded_patient_L_test.map(fine_tuned_map_to_pred)
+# z_patient_L_result = WER.compute(references=z_patient_L['reference'], predictions=z_patient_L['prediction'])
+# z_patient_T = encoded_patient_T_test.map(fine_tuned_map_to_pred)
+# z_patient_T_result = WER.compute(references=z_patient_T['reference'], predictions=z_patient_T['prediction'])
+
+z_john_p326_result = encoded_John_p326.map(fine_tuned_map_to_pred)
+z_john_p326 = WER.compute(references=z_john_p326_result['reference'], predictions=z_john_p326_result['prediction'])
+
+pdb.set_trace() 
+
+# # z_john_p326 = WER.compute(references=z_john_p326_result['reference'], predictions=z_john_p326_result['prediction'])
+# # 0.4648241206030151
+pdb.set_trace() 
+
+# # y_John_video= fine_tuned_trainer.predict(encoded_John_video)
+# # metrics={'test_loss': 2.665189743041992, 'test_wer': 0.7222222222222222, 'test_runtime': 0.1633, 'test_samples_per_second': 48.979, 'test_steps_per_second': 6.122})
+# pdb.set_trace()
+
+# p326 training
+# metrics={'test_loss': 0.4804028868675232, 'test_wer': 0.21787709497206703, 'test_runtime': 0.3594, 'test_samples_per_second': 44.517, 'test_steps_per_second': 5.565})
+# hel metrics={'test_loss': 1.6363693475723267, 'test_wer': 0.17951881554595928, 'test_runtime': 3.8451, 'test_samples_per_second': 41.611, 'test_steps_per_second': 5.201})
+# Fary: metrics={'t est_loss': 1.4633615016937256, 'test_wer': 0.5572139303482587, 'test_runtime': 0.6627, 'test_samples_per_second': 45.27, 'test_steps_per_second': 6.036})
+# p326 large: metrics={'test_loss': 0.6568527817726135, 'test_wer': 0.2889447236180904, 'test_runtime': 0.7169, 'test_samples_per_second': 51.613, 'test_steps_per_second': 6.975})

packages.txt

@@ -0,0 +1,2 @@
+festival 
+espeak # or espeak-ng on Linux

requirements.txt

@@ -0,0 +1,103 @@
+absl-py==1.0.0
+aiohttp==3.8.1
+aiosignal==1.2.0
+analytics-python==1.4.0
+antlr4-python3-runtime==4.8
+anyio==3.5.0
+asgiref==3.5.0
+async-timeout==4.0.2
+attrs==21.4.0
+backoff==1.10.0
+bcrypt==3.2.0
+bitarray==2.4.0
+cachetools==5.0.0
+certifi==2021.10.8
+cffi==1.15.0
+charset-normalizer==2.0.12
+click==8.0.4
+colorama==0.4.4
+cryptography==36.0.1
+cycler==0.11.0
+Cython==0.29.28
+fairseq @ git+https://github.com/pytorch/fairseq.git@d03f4e771484a433f025f47744017c2eb6e9c6bc
+fastapi==0.75.0
+ffmpy==0.3.0
+fonttools==4.30.0
+frozenlist==1.3.0
+fsspec==2022.2.0
+future==0.18.2
+google-auth==2.6.0
+google-auth-oauthlib==0.4.6
+gradio==3.40.0
+grpcio==1.44.0
+h11==0.12.0
+hydra-core==1.0.7
+idna==3.3
+importlib-metadata==4.11.3
+Jinja2==3.0.3
+kiwisolver==1.3.2
+linkify-it-py==1.0.3
+Markdown==3.3.6
+markdown-it-py==2.0.1
+MarkupSafe==2.1.0
+matplotlib==3.5.1
+mdit-py-plugins==0.3.0
+mdurl==0.1.0
+monotonic==1.6
+multidict==6.0.2
+numpy==1.22.3
+oauthlib==3.2.0
+omegaconf==2.0.6
+orjson==3.6.7
+packaging==21.3
+pandas==1.4.1
+paramiko==2.10.1
+Pillow==9.0.1
+portalocker==2.4.0
+protobuf==3.19.4
+pyasn1==0.4.8
+pyasn1-modules==0.2.8
+pycparser==2.21
+pycryptodome==3.14.1
+pydantic==1.9.0
+pyDeprecate==0.3.1
+pydub==0.25.1
+PyNaCl==1.5.0
+pyparsing==3.0.7
+python-dateutil==2.8.2
+python-multipart==0.0.5
+pytorch-lightning==1.5.10
+pytz==2021.3
+PyYAML==6.0
+regex==2022.3.2
+requests==2.27.1
+requests-oauthlib==1.3.1
+rsa==4.8
+sacrebleu==2.0.0
+six==1.16.0
+sniffio==1.2.0
+starlette==0.17.1
+tabulate==0.8.9
+tensorboard==2.8.0
+tensorboard-data-server==0.6.1
+tensorboard-plugin-wit==1.8.1
+torch==1.12.1
+torchaudio==0.12.1
+torchmetrics==0.7.2
+tqdm==4.63.0
+typing-extensions==4.1.1
+uc-micro-py==1.0.1
+urllib3==1.26.8
+uvicorn==0.17.6
+Werkzeug==2.0.3
+yarl==1.7.2
+zipp==3.7.0
+
+transformers
+deepspeech
+tensorboardX
+jiwer
+phonemizer
+librosa
+
+rich

src/description.html

@@ -0,0 +1,13 @@
+<p>This is the experiment page for Laronix Data Recording.<br>
+    <br>
+    1. Select one example from below, a sound file, with its reference transcription and its speaking rate will be loaded as inputs.<br>
+     You can check the sound file first and prepare for reading the transcription at a similar tempo.<br>
+    2. Delete the sound file (click the X button on the right), a recording button will appear.<br>
+    3. Click the recording button to start, click again to stop. Make sure you are not mispronouncing or including any detectable noises.<br>
+    4. Click &quot;Submit&quot; button and wait for the result.<br>
+    5. Please check the message box to see the feedback, if ERROR appears, delete your previous recording and try again :).<br>
+    6. If &quot;GOOD JOB!&quot; message appears, click &quot;Flag as Perfect&quot; and start another recording.<br>
+    7. If you try several times (N >= 10) and still can not clear the mission, you can flag your best recording by clicking &quot;Doubtful Speaking Rate&quot; or &quot;Doubtful Naturalness&quot;. <br>
+     Yet this seldom happens, so please try to meet the system's requirement first!<br>
+    8. If you have any other question, Please contact kevin@laronix.com </p>
+<img src="https://static.wixstatic.com/media/e7e144_93e98148d06147828031797eb4525b80~mv2.png/v1/crop/x_0,y_25,w_2606,h_882/fill/w_396,h_142,al_c,q_85,usm_0.66_1.00_0.01,enc_auto/newlogo.png" align="right" height="20%" width="20%">

src/lightning_module.py

@@ -0,0 +1,43 @@
+import sys
+sys.path.append("src")
+import pytorch_lightning as pl
+import torch
+import torch.nn as nn
+import os
+import numpy as np
+import hydra
+from model import load_ssl_model, PhonemeEncoder, DomainEmbedding, LDConditioner, Projection
+
+
+class BaselineLightningModule(pl.LightningModule):
+    def __init__(self, cfg):
+        super().__init__()
+        self.cfg = cfg
+        self.construct_model()
+        self.save_hyperparameters()
+    
+    def construct_model(self):
+        self.feature_extractors = nn.ModuleList([
+            load_ssl_model(cp_path='src/wav2vec_small.pt'),
+            DomainEmbedding(3,128),
+        ])
+        output_dim = sum([ feature_extractor.get_output_dim() for feature_extractor in self.feature_extractors])
+        output_layers = [
+            LDConditioner(judge_dim=128,num_judges=3000,input_dim=output_dim)
+        ]
+        output_dim = output_layers[-1].get_output_dim()
+        output_layers.append(
+            Projection(hidden_dim=2048,activation=torch.nn.ReLU(),range_clipping=False,input_dim=output_dim)
+
+        )
+
+        self.output_layers = nn.ModuleList(output_layers)
+
+    def forward(self, inputs):
+        outputs = {}
+        for feature_extractor in self.feature_extractors:
+            outputs.update(feature_extractor(inputs))
+        x = outputs
+        for output_layer in self.output_layers:
+            x = output_layer(x,inputs)
+        return x

src/model.py

@@ -0,0 +1,191 @@
+import torch
+import torch.nn as nn
+import fairseq
+import os
+import hydra
+
+def load_ssl_model(cp_path):
+    ssl_model_type = cp_path.split("/")[-1]
+    wavlm =  "WavLM" in ssl_model_type
+    if wavlm:
+        checkpoint = torch.load(cp_path)
+        cfg = WavLMConfig(checkpoint['cfg'])
+        ssl_model = WavLM(cfg)
+        ssl_model.load_state_dict(checkpoint['model'])
+        if 'Large' in ssl_model_type:
+            SSL_OUT_DIM = 1024
+        else:
+            SSL_OUT_DIM = 768
+    else:
+        if ssl_model_type == "wav2vec_small.pt":
+            SSL_OUT_DIM = 768
+        elif ssl_model_type in ["w2v_large_lv_fsh_swbd_cv.pt", "xlsr_53_56k.pt"]:
+            SSL_OUT_DIM = 1024
+        else:
+            print("*** ERROR *** SSL model type " + ssl_model_type + " not supported.")
+            exit()
+        model, cfg, task = fairseq.checkpoint_utils.load_model_ensemble_and_task(
+            [cp_path]
+        )
+        ssl_model = model[0]
+        ssl_model.remove_pretraining_modules()
+    return SSL_model(ssl_model, SSL_OUT_DIM, wavlm)
+
+class SSL_model(nn.Module):
+    def __init__(self,ssl_model,ssl_out_dim,wavlm) -> None:
+        super(SSL_model,self).__init__()
+        self.ssl_model, self.ssl_out_dim = ssl_model, ssl_out_dim
+        self.WavLM = wavlm
+
+    def forward(self,batch):
+        wav = batch['wav'] 
+        wav = wav.squeeze(1) # [batches, audio_len]
+        if self.WavLM:
+            x = self.ssl_model.extract_features(wav)[0]
+        else:
+            res = self.ssl_model(wav, mask=False, features_only=True)
+            x = res["x"]
+        return {"ssl-feature":x}
+    def get_output_dim(self):
+        return self.ssl_out_dim
+
+
+class PhonemeEncoder(nn.Module):
+    '''
+    PhonemeEncoder consists of an embedding layer, an LSTM layer, and a linear layer.
+    Args:
+        vocab_size: the size of the vocabulary
+        hidden_dim: the size of the hidden state of the LSTM
+        emb_dim: the size of the embedding layer
+        out_dim: the size of the output of the linear layer
+        n_lstm_layers: the number of LSTM layers
+    '''
+    def __init__(self, vocab_size, hidden_dim, emb_dim, out_dim,n_lstm_layers,with_reference=True) -> None:
+        super().__init__()
+        self.with_reference = with_reference
+        self.embedding = nn.Embedding(vocab_size, emb_dim)
+        self.encoder = nn.LSTM(emb_dim, hidden_dim,
+                               num_layers=n_lstm_layers, dropout=0.1, bidirectional=True)
+        self.linear = nn.Sequential(
+                nn.Linear(hidden_dim + hidden_dim*self.with_reference, out_dim),
+                nn.ReLU()
+                )
+        self.out_dim = out_dim
+
+    def forward(self,batch):
+        seq = batch['phonemes']
+        lens = batch['phoneme_lens']
+        reference_seq = batch['reference']
+        reference_lens = batch['reference_lens']
+        emb = self.embedding(seq)
+        emb = torch.nn.utils.rnn.pack_padded_sequence(
+            emb, lens, batch_first=True, enforce_sorted=False)
+        _, (ht, _) = self.encoder(emb)
+        feature = ht[-1] + ht[0]
+        if self.with_reference:
+            if reference_seq==None or reference_lens ==None:
+                raise ValueError("reference_batch and reference_lens should not be None when with_reference is True")
+            reference_emb = self.embedding(reference_seq)
+            reference_emb = torch.nn.utils.rnn.pack_padded_sequence(
+                reference_emb, reference_lens, batch_first=True, enforce_sorted=False)
+            _, (ht_ref, _) = self.encoder(emb)
+            reference_feature = ht_ref[-1] + ht_ref[0]
+            feature = self.linear(torch.cat([feature,reference_feature],1))
+        else:
+            feature = self.linear(feature)
+        return {"phoneme-feature": feature}
+    def get_output_dim(self):
+        return self.out_dim
+
+class DomainEmbedding(nn.Module):
+    def __init__(self,n_domains,domain_dim) -> None:
+        super().__init__()
+        self.embedding = nn.Embedding(n_domains,domain_dim)
+        self.output_dim = domain_dim
+    def forward(self, batch):
+        return {"domain-feature": self.embedding(batch['domains'])}
+    def get_output_dim(self):
+        return self.output_dim
+
+
+class LDConditioner(nn.Module):
+    '''
+    Conditions ssl output by listener embedding
+    '''
+    def __init__(self,input_dim, judge_dim, num_judges=None):
+        super().__init__()
+        self.input_dim = input_dim
+        self.judge_dim = judge_dim
+        self.num_judges = num_judges
+        assert num_judges !=None
+        self.judge_embedding = nn.Embedding(num_judges, self.judge_dim)
+        # concat [self.output_layer, phoneme features]
+        
+        self.decoder_rnn = nn.LSTM(
+            input_size = self.input_dim + self.judge_dim,
+            hidden_size = 512,
+            num_layers = 1,
+            batch_first = True,
+            bidirectional = True
+        ) # linear?
+        self.out_dim = self.decoder_rnn.hidden_size*2
+
+    def get_output_dim(self):
+        return self.out_dim
+
+
+    def forward(self, x, batch):
+        judge_ids = batch['judge_id']
+        if 'phoneme-feature' in x.keys():
+            concatenated_feature = torch.cat((x['ssl-feature'], x['phoneme-feature'].unsqueeze(1).expand(-1,x['ssl-feature'].size(1) ,-1)),dim=2)
+        else:
+            concatenated_feature = x['ssl-feature']
+        if 'domain-feature' in x.keys():
+            concatenated_feature = torch.cat(
+                (
+                    concatenated_feature,
+                    x['domain-feature']
+                    .unsqueeze(1)
+                    .expand(-1, concatenated_feature.size(1), -1),
+                ),
+                dim=2,
+            )
+        if judge_ids != None:
+            concatenated_feature = torch.cat(
+                (
+                    concatenated_feature,
+                    self.judge_embedding(judge_ids)
+                    .unsqueeze(1)
+                    .expand(-1, concatenated_feature.size(1), -1),
+                ),
+                dim=2,
+            )
+            decoder_output, (h, c) = self.decoder_rnn(concatenated_feature)
+        return decoder_output
+
+class Projection(nn.Module):
+    def __init__(self, input_dim, hidden_dim, activation, range_clipping=False):
+        super(Projection, self).__init__()
+        self.range_clipping = range_clipping
+        output_dim = 1
+        if range_clipping:
+            self.proj = nn.Tanh()
+        
+        self.net = nn.Sequential(
+            nn.Linear(input_dim, hidden_dim),
+            activation,
+            nn.Dropout(0.3),
+            nn.Linear(hidden_dim, output_dim),
+        )
+        self.output_dim = output_dim
+    
+    def forward(self, x, batch):
+        output = self.net(x)
+
+        # range clipping
+        if self.range_clipping:
+            return self.proj(output) * 2.0 + 3
+        else:
+            return output
+    def get_output_dim(self):
+        return self.output_dim