eval/evaluate-whisper-lora.ipynb

{
	"cells": [
		{
			"cell_type": "markdown",
			"metadata": {},
			"source": [
				"# evaluate WER of whisper with PEFT LoRA\n",
				"\n",
				"[![](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/phineas-pta/fine-tune-whisper-vi/blob/main/eval/evaluate-whisper-lora.ipynb)\n",
				"\n",
				"*kaggle TPU crash when running inference* ~~can be used on kaggle TPU, but do not enable `XLA_USE_BF16` because of AMP (Automatic Mixed Precision)~~\n",
				"\n",
				"try `transformers.pipeline` but error with `torch.autocast`"
			]
		},
		{
			"cell_type": "code",
			"execution_count": null,
			"metadata": {},
			"outputs": [],
			"source": [
				"from huggingface_hub import notebook_login\n",
				"notebook_login()\n",
				"# !huggingface-cli login --token=███"
			]
		},
		{
			"cell_type": "code",
			"execution_count": null,
			"metadata": {},
			"outputs": [],
			"source": [
				"# workaround for a bug in `datasets` package\n",
				"%pip uninstall -y cudf dask-cuda dask-cudf\n",
				"%pip install -q cudf-cu12 --extra-index-url=https://pypi.nvidia.com\n",
				"%pip install -qU 'datasets[audio]' accelerate transformers jiwer bitsandbytes peft\n",
				"# install then `import evaluate` throw error on kaggle"
			]
		},
		{
			"cell_type": "code",
			"execution_count": null,
			"metadata": {},
			"outputs": [],
			"source": [
				"from tqdm import tqdm\n",
				"import torch\n",
				"# import torch_xla.core.xla_model as xm  # on kaggle TPU\n",
				"from peft import PeftModel, PeftConfig\n",
				"from transformers import AutomaticSpeechRecognitionPipeline, WhisperForConditionalGeneration, WhisperFeatureExtractor, WhisperTokenizer\n",
				"import datasets as hugDS\n",
				"import jiwer\n",
				"\n",
				"JIWER_TRANS = jiwer.Compose([  # DO NOT use `jiwer.RemoveEmptyStrings` it can cause rows count mismatch\n",
				"\tjiwer.ToLowerCase(),\n",
				"\tjiwer.RemoveKaldiNonWords(),\n",
				"\tjiwer.RemoveMultipleSpaces(),\n",
				"\tjiwer.Strip(),\n",
				"\tjiwer.RemovePunctuation(),\n",
				"\tjiwer.ReduceToListOfListOfWords(),\n",
				"])"
			]
		},
		{
			"cell_type": "code",
			"execution_count": null,
			"metadata": {},
			"outputs": [],
			"source": [
				"SAMPLING_RATE = 16_000\n",
				"def load_my_data(**kwargs):\n",
				"\treturn hugDS.load_dataset(**kwargs, split=\"test\", trust_remote_code=True, streaming=True).cast_column(\"audio\", hugDS.Audio(sampling_rate=SAMPLING_RATE))\n",
				"\n",
				"MY_DATA = hugDS.IterableDatasetDict()\n",
				"MY_DATA[\"commonvoice\"] = load_my_data(path=\"mozilla-foundation/common_voice_16_1\", name=\"vi\",  ).select_columns([\"audio\", \"sentence\"])\n",
				"# MY_DATA[\"fleurs\"] # disable FLEURS because error with tensor size mismatch when batching, see bottom for non-batched inference\n",
				"MY_DATA[\"vivos\"]       = load_my_data(path=\"vivos\"                                             ).select_columns([\"audio\", \"sentence\"])\n",
				"MY_DATA[\"bud500\"]      = load_my_data(path=\"linhtran92/viet_bud500\"                            ).rename_column(\"transcription\", \"sentence\")\n",
				"MY_DATA[\"lsvsc\"]       = load_my_data(path=\"doof-ferb/LSVSC\"                                   ).select_columns([\"audio\", \"transcription\"]).rename_column(\"transcription\", \"sentence\")\n",
				"\n",
				"ROWS_COUNT = {\n",
				"\t\"commonvoice\": 1326,\n",
				"\t\"fleurs\":       857,\n",
				"\t\"vivos\":        760,\n",
				"\t\"bud500\":      7500,\n",
				"\t\"lsvsc\":       5683,\n",
				"}"
			]
		},
		{
			"cell_type": "code",
			"execution_count": null,
			"metadata": {},
			"outputs": [],
			"source": [
				"PEFT_MODEL_ID = \"doof-ferb/whisper-large-peft-lora-vi\"  # @param [\"doof-ferb/whisper-large-peft-lora-vi\", \"daila/whisper-large-v3_LoRA_Common-Vi_WER\", \"daila/whisper-large-v3_LoRA_vi\", \"vikas85/whisper-vlsp-peft\", \"vikas85/whisper-vlsp\", \"vikas85/whisper-fosd-peft\", \"vikas85/whisper-fleurs-peft-vi-2\", \"DuyTa/vi-whisper-medium-Lora\", \"vikas85/whisper-cv-fleur-v6\", \"vikas85/fleurs-vn-peft-v2\", \"Yuhthe/openai-whisper-small-vivos-LORA-colab\"]\n",
				"BASE_MODEL_ID = PeftConfig.from_pretrained(PEFT_MODEL_ID).base_model_name_or_path\n",
				"print(\"adapter to\", BASE_MODEL_ID)\n",
				"\n",
				"# declare task & language in extractor & tokenizer have no effect in inference\n",
				"FEATURE_EXTRACTOR = WhisperFeatureExtractor.from_pretrained(BASE_MODEL_ID)\n",
				"TOKENIZER = WhisperTokenizer.from_pretrained(BASE_MODEL_ID)\n",
				"\n",
				"MODEL = PeftModel.from_pretrained(\n",
				"\tWhisperForConditionalGeneration.from_pretrained(\n",
				"\t\tBASE_MODEL_ID,\n",
				"\t\tdevice_map=\"auto\",  # {\"\": \"cuda:0\"}  # {\"\": xm.xla_device()}  # on kaggle TPU\n",
				"\t\ttorch_dtype=torch.float16,  # bfloat16 on TPU or newer GPU\n",
				"\t\t# quantization_config=BitsAndBytesConfig(load_in_8bit=True)  # do not enable because it make inference super slow\n",
				"\t),  # do not use `.to(\"cuda\")` because throw weird error about torch dtype mismatch\n",
				"\tPEFT_MODEL_ID\n",
				").merge_and_unload(progressbar=True)  # reduce latency with LoRA\n",
				"\n",
				"PIPE = AutomaticSpeechRecognitionPipeline(model=MODEL, tokenizer=TOKENIZER, feature_extractor=FEATURE_EXTRACTOR)\n",
				"PIPE_KWARGS = {\"language\": \"vi\", \"task\": \"transcribe\"}\n",
				"\n",
				"BATCH_SIZE =  16 # @param {type: \"integer\"}\n",
				"# @markdown for colab free T4 @ `float16`: 16 for large model, 32 medium, 48 small"
			]
		},
		{
			"cell_type": "code",
			"execution_count": null,
			"metadata": {},
			"outputs": [],
			"source": [
				"# workaround because KeyDataset(MY_DATA[split], \"audio\") raise error with streaming datasets\n",
				"def data(batch):\n",
				"\tfor row in batch:\n",
				"\t\tyield row[\"audio\"]\n",
				"\n",
				"\n",
				"@torch.autocast(device_type=\"cuda\")  # required by PEFT\n",
				"# @torch.autocast(device_type=\"xla\", dtype=torch.bfloat16)  # on kaggle TPU\n",
				"@torch.inference_mode()\n",
				"def predict(split: str) -> tuple[list[str]]:\n",
				"\tbatch = MY_DATA[split]\n",
				"\ty_pred = [out[\"text\"] for out in tqdm(PIPE(data(batch), generate_kwargs=PIPE_KWARGS, batch_size=BATCH_SIZE), total=ROWS_COUNT[split], unit=\"samples\", desc=f\"{split=}\")]\n",
				"\ty_true = [row[\"sentence\"] for row in batch]\n",
				"\tassert len(y_pred) == len(y_true)\n",
				"\treturn y_true, y_pred\n",
				"\n",
				"\n",
				"for split in MY_DATA.keys():\n",
				"\ty_true, y_pred = predict(split)\n",
				"\ttorch.cuda.empty_cache()  # forced clean\n",
				"\twer = 100 * jiwer.wer(\n",
				"\t\treference=y_true,\n",
				"\t\thypothesis=y_pred,\n",
				"\t\treference_transform=JIWER_TRANS,\n",
				"\t\thypothesis_transform=JIWER_TRANS,\n",
				"\t)\n",
				"\tif 0 < wer < 100:\n",
				"\t\tprint(f\"WER on {split} = {wer:.1f}%\", end=\"\\n\\n\")\n",
				"\telse:\n",
				"\t\tprint(\"something wrong, check 5 first & last transcription:\")\n",
				"\t\tprint(y_true[:5], y_true[-5:])\n",
				"\t\tprint(y_pred[:5], y_pred[-5:])\n"
			]
		},
		{
			"cell_type": "code",
			"execution_count": null,
			"metadata": {},
			"outputs": [],
			"source": [
				"data_fleurs = load_my_data(path=\"google/fleurs\", name=\"vi_vn\", streaming=False).select_columns([\"audio\", \"transcription\"])\n",
				"\n",
				"@torch.autocast(device_type=\"cuda\")  # required by PEFT\n",
				"@torch.inference_mode()\n",
				"def predict_fleurs(batch):\n",
				"\tbatch[\"pred\"] = PIPE(batch[\"audio\"], generate_kwargs=PIPE_KWARGS)[\"text\"]\n",
				"\treturn batch\n",
				"data_fleurs = data_fleurs.map(predict_fleurs)  # progress bar included\n",
				"\n",
				"wer = 100 * jiwer.wer(\n",
				"\treference=data_fleurs[\"transcription\"],\n",
				"\thypothesis=data_fleurs[\"pred\"],\n",
				"\treference_transform=JIWER_TRANS,\n",
				"\thypothesis_transform=JIWER_TRANS,\n",
				")\n",
				"print(f\"WER on FLEURS = {wer:.1f}%\", end=\"\\n\\n\")"
			]
		}
	],
	"metadata": {
		"accelerator": "GPU",
		"colab": {
			"gpuType": "T4",
			"private_outputs": true,
			"provenance": []
		},
		"kaggle": {
			"dataSources": [],
			"isGpuEnabled": true,
			"isInternetEnabled": true,
			"language": "python",
			"sourceType": "notebook"
		},
		"kernelspec": {
			"display_name": "Python 3",
			"language": "python",
			"name": "python3"
		},
		"language_info": {
			"name": "python"
		}
	},
	"nbformat": 4,
	"nbformat_minor": 0
}