{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceId":73047,"databundleVersionId":8149390,"sourceType":"competition"},{"sourceId":6707460,"sourceType":"datasetVersion","datasetId":3865741},{"sourceId":8068743,"sourceType":"datasetVersion","datasetId":4760663}],"dockerImageVersionId":30674,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# **0. Import Libraries**","metadata":{}},{"cell_type":"code","source":"import os\n\nimport pandas as pd\n\nimport librosa\nimport librosa.display\n\nimport numpy as np\n\nimport IPython.display as ipd\n\nimport matplotlib.pyplot as plt\n\nimport random\n\nfrom collections import Counter\n\nfrom sklearn.model_selection import train_test_split\n\nimport torch\nimport torchaudio\n\nfrom dataclasses import dataclass\nfrom typing import Any, Dict, List, Union\nfrom datasets import DatasetDict\nfrom datasets import Dataset as DS\n\nfrom transformers import (\n    WhisperFeatureExtractor,\n    WhisperTokenizer,\n    WhisperProcessor,\n    WhisperForConditionalGeneration,\n    GenerationConfig,\n    Seq2SeqTrainingArguments,\n    Seq2SeqTrainer,\n    TrainerCallback,\n    TrainingArguments,\n    TrainerState,\n    TrainerControl,\n    EarlyStoppingCallback,\n    pipeline\n)\nfrom torchmetrics.text import WordErrorRate, CharErrorRate","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# **1.1 Load Dataset**","metadata":{}},{"cell_type":"code","source":"# BASE_DIR = '/kaggle/input/ben10/ben10'\n# train_data_dir = f\"{BASE_DIR}/16_kHz_train_audio/\"\n# test_data_dir = f\"{BASE_DIR}/16_kHz_valid_audio/\"\n# data_path = f\"{BASE_DIR}/train.csv\"\n\n# split2path = {\n#     \"train\": train_data_dir,\n#     \"test\": test_data_dir,\n# }\n# data = pd.read_csv(data_path)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"BASE_DIR = '/kaggle/input/bengali-asr-dialect-dataset/dataset'\ntrain_data_dir = f\"{BASE_DIR}/train/train_audio\"\nvalidation_data_dir = f\"{BASE_DIR}/validation/validation_audio\"\ntest_data_dir = f\"{BASE_DIR}/test/test_audio\"\ntrain_csv_path = f\"{BASE_DIR}/train/train.csv\"\nvalidation_csv_path = f\"{BASE_DIR}/validation/validation.csv\"\n\ntrain_data = pd.read_csv(train_csv_path)\nvalidation_data = pd.read_csv(validation_csv_path)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_data","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# **1.2 Cleaning Dataset + Spliting + Dataset Dictionary Creation**","metadata":{}},{"cell_type":"code","source":"def beautify_dataset(data):\n    splits = []\n    transcripts = []\n    \n    for i in range(len(data)):\n        filename, transcript,district,file_path = data.iloc[i]\n\n        split = filename.split(\"_\")[0]\n        transcript = transcript.replace(\"\\n\", \" \")\n        transcript = \" \".join(transcript.split())\n        \n        splits.append(split)\n        transcripts.append(transcript)\n    \n    data['split'] = splits\n    data['transcripts'] = transcripts\n    \n    return data\n\n# train_data= train_data.head(20)\ntrain_data = beautify_dataset(train_data)\n\ntrain_data.drop(train_data[train_data['transcripts'] == ''].index, inplace=True)\ntrain_data.drop(train_data[train_data['transcripts'] == \"<>\"].index, inplace=True)\ntrain_data.drop(train_data[train_data['transcripts'] == \"..\"].index, inplace=True)\ntrain_data[\"transcripts\"] = train_data[\"transcripts\"].str.strip()\n\ntrain_df, eval_df = train_test_split(train_data, test_size=0.01, shuffle=True)\nds_splits = DatasetDict({\n    'train': DS.from_pandas(train_df),\n    'eval': DS.from_pandas(eval_df)\n})\nds_splits = ds_splits.remove_columns([\"split\"])","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"ds_splits","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# **2. Preparing Feature Extractor, Tokenizer and Processor & pre-trained model**","metadata":{}},{"cell_type":"code","source":"TASK = \"transcribe\"\n# MODEL_NAME = \"openai/whisper-large-v3\"\n# MODEL_NAME = \"openai/whisper-medium\"\n# MODEL_NAME = \"openai/whisper-small\"\nMODEL_NAME=\"/kaggle/input/bengali-ai-asr-submission/bengali-whisper-medium\"\nLANGUAGE='bn'\n\nfeature_extractor = WhisperFeatureExtractor.from_pretrained(MODEL_NAME)\ntokenizer = WhisperTokenizer.from_pretrained(MODEL_NAME, language=LANGUAGE, task=TASK)\nprocessor = WhisperProcessor.from_pretrained(MODEL_NAME, language=LANGUAGE, task=TASK)\nmodel = WhisperForConditionalGeneration.from_pretrained(MODEL_NAME, device_map=\"auto\")\n\nmodel.generation_config=GenerationConfig.from_pretrained(\"openai/whisper-medium\")\nmodel.generation_config.language = LANGUAGE\nmodel.generation_config.task = TASK\nmodel.generation_config.forced_decoder_ids = None\nmodel.config.suppress_tokens = [] # added later","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# **3. Prepare Dataset**","metadata":{}},{"cell_type":"code","source":"def prepare_dataset(example):\n    audio_path = example[\"file_path\"]\n    audio, sr = librosa.load(audio_path, sr=16_000)\n    example[\"input_features\"] = feature_extractor(audio, sampling_rate=sr).input_features[0]\n    example[\"labels\"] = tokenizer(f\"{example['transcripts']}\", max_length=448, padding=True, truncation=True).input_ids\n    return example\n\nnp.object=object\nds_splits = ds_splits.map(prepare_dataset, remove_columns=ds_splits.column_names[\"train\"],num_proc=2)\n# ds_splits.save_to_disk(\"/kaggle/working/bengali-dialect-pre-trained-full-data\")","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# **4. Training**","metadata":{}},{"cell_type":"code","source":"@dataclass\nclass DataCollatorSpeechSeq2SeqWithPadding:\n    processor: Any\n\n    def __call__(self, features: List[Dict[str, Union[List[int], torch.Tensor]]]) -> Dict[str, torch.Tensor]:\n        # split inputs and labels since they have to be of different lengths and need different padding methods\n        # first treat the audio inputs by simply returning torch tensors\n        input_features = [{\"input_features\": feature[\"input_features\"]} for feature in features]\n        batch = self.processor.feature_extractor.pad(input_features, return_tensors=\"pt\")\n\n        # get the tokenized label sequences\n        label_features = [{\"input_ids\": feature[\"labels\"]} for feature in features]\n        # pad the labels to max length\n        labels_batch = self.processor.tokenizer.pad(label_features, return_tensors=\"pt\")\n\n        # replace padding with -100 to ignore loss correctly\n        labels = labels_batch[\"input_ids\"].masked_fill(labels_batch.attention_mask.ne(1), -100)\n\n        # if bos token is appended in previous tokenization step,\n        # cut bos token here as it's append later anyways\n        if (labels[:, 0] == self.processor.tokenizer.bos_token_id).all().cpu().item():\n            labels = labels[:, 1:]\n\n        batch[\"labels\"] = labels\n        \n        torch.cuda.empty_cache()\n\n        return batch\ndata_collator = DataCollatorSpeechSeq2SeqWithPadding(processor=processor)\n\ncer = CharErrorRate()\nwer = WordErrorRate()\n\ndef compute_metrics(pred):\n    pred_ids = pred.predictions\n    label_ids = pred.label_ids\n\n    label_ids[label_ids == -100] = tokenizer.pad_token_id\n\n    pred_str = tokenizer.batch_decode(pred_ids, skip_special_tokens=True)\n    label_str = tokenizer.batch_decode(label_ids, skip_special_tokens=True)\n\n    wer_res = wer(pred_str, label_str)\n    cer_res = cer(pred_str, label_str)\n    \n    print(\"WER:\",wer_res,\"| CER:\", cer_res) # to show up during running logs\n    print(\"Pred:\",pred_str[0])\n    print(\"Label:\",label_str[0])\n    \n    return {\"wer\": wer_res, \"cer\": cer_res}\n\nsteps=1000\nbatch=8\ntraining_args = Seq2SeqTrainingArguments(\n    output_dir=\"trained_pretrained_whisper_medium\",\n    per_device_train_batch_size=batch,\n    per_device_eval_batch_size=batch,\n    gradient_accumulation_steps=1,\n    gradient_checkpointing=True,\n    learning_rate=1e-5,\n    fp16=True,\n    weight_decay=1e-2,\n    warmup_steps=50,\n    num_train_epochs=1,\n    evaluation_strategy=\"steps\", # or \"epochs\"\n    predict_with_generate=True,\n#     generation_max_length=448,\n    save_steps=steps,\n    eval_steps=steps,\n    logging_steps=steps,\n    save_total_limit=1,\n    load_best_model_at_end=True,\n    metric_for_best_model=\"wer\",\n    greater_is_better=False,\n    push_to_hub=False,\n    report_to=\"none\",\n    remove_unused_columns=False,\n)\n\ntrainer = Seq2SeqTrainer(\n    args=training_args,\n    model=model,\n    train_dataset=ds_splits[\"train\"],\n    eval_dataset=ds_splits[\"eval\"],\n    data_collator=data_collator,\n    tokenizer=processor.feature_extractor,\n    compute_metrics=compute_metrics,\n#     callbacks=[EarlyStoppingCallback(2, 1.0)]\n)\ntrainer.train()\ntrainer.save_model(training_args.output_dir)\nprocessor.save_pretrained(training_args.output_dir)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import gc\ngc.collect()\ntorch.cuda.empty_cache()","metadata":{"trusted":true},"execution_count":null,"outputs":[]}]}