{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Hugging Face🤗 Wav2Vec2.0 Inference Notebook\n### This Notebook is forked my private notebook. Due to the limitations of Kaggle notebook, so trained shorter epoch.\n\n## [This code was very helpful in creating this notebook.](https://colab.research.google.com/github/patrickvonplaten/notebooks/blob/master/Fine_Tune_XLSR_Wav2Vec2_on_Turkish_ASR_with_%F0%9F%A4%97_Transformers.ipynb)","metadata":{}},{"cell_type":"markdown","source":"## Import Libralies","metadata":{}},{"cell_type":"code","source":"import torch \nimport torch.nn as nn\nimport torchaudio\nimport torchaudio.transforms as tat\nimport numpy as np\nimport pandas as pd\nfrom datasets import load_dataset, load_metric, Audio\nimport os\nimport json\nfrom tqdm import tqdm","metadata":{"execution":{"iopub.status.busy":"2023-08-10T21:43:44.322214Z","iopub.execute_input":"2023-08-10T21:43:44.322483Z","iopub.status.idle":"2023-08-10T21:43:49.802136Z","shell.execute_reply.started":"2023-08-10T21:43:44.322458Z","shell.execute_reply":"2023-08-10T21:43:49.801047Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"torchaudio.set_audio_backend(\"soundfile\")\n# Default is set \"sox\", but it cannnot load mp3 (on experiment)","metadata":{"execution":{"iopub.status.busy":"2023-08-10T21:43:49.807405Z","iopub.execute_input":"2023-08-10T21:43:49.810066Z","iopub.status.idle":"2023-08-10T21:43:49.817880Z","shell.execute_reply.started":"2023-08-10T21:43:49.810028Z","shell.execute_reply":"2023-08-10T21:43:49.816967Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"SR = 16000 # Wav2Vec2.0 requires samplerate 16000.\ntorch.backends.cudnn.benchmark = True","metadata":{"execution":{"iopub.status.busy":"2023-08-10T21:43:49.822674Z","iopub.execute_input":"2023-08-10T21:43:49.825543Z","iopub.status.idle":"2023-08-10T21:43:49.835409Z","shell.execute_reply.started":"2023-08-10T21:43:49.825509Z","shell.execute_reply":"2023-08-10T21:43:49.834481Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model_dir = f'/kaggle/input/hf-wav2vec2-0-train-baseline/ASR_Bengali/checkpoint-15000'","metadata":{"execution":{"iopub.status.busy":"2023-08-10T21:43:49.841518Z","iopub.execute_input":"2023-08-10T21:43:49.842967Z","iopub.status.idle":"2023-08-10T21:43:49.849737Z","shell.execute_reply.started":"2023-08-10T21:43:49.842935Z","shell.execute_reply":"2023-08-10T21:43:49.848926Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Load Dataset","metadata":{}},{"cell_type":"code","source":"from glob import glob\nTEST_DIRECTORY = '/kaggle/input/bengaliai-speech/test_mp3s'\npaths = glob(os.path.join(TEST_DIRECTORY,'*.mp3'))\npaths[:5]","metadata":{"execution":{"iopub.status.busy":"2023-08-10T21:43:49.851359Z","iopub.execute_input":"2023-08-10T21:43:49.852866Z","iopub.status.idle":"2023-08-10T21:43:49.871036Z","shell.execute_reply.started":"2023-08-10T21:43:49.852834Z","shell.execute_reply":"2023-08-10T21:43:49.869829Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Prepare to infer in transformers' Wav2Vec2.0","metadata":{}},{"cell_type":"code","source":"from transformers import Wav2Vec2CTCTokenizer\n\ntokenizer = Wav2Vec2CTCTokenizer(f\"/kaggle/input/hf-wav2vec2-0-preprocess-baseline/vocab_bengali.json\", unk_token=\"[UNK]\", pad_token=\"[PAD]\", word_delimiter_token=\"|\")\ntokenizer","metadata":{"execution":{"iopub.status.busy":"2023-08-10T21:43:49.873571Z","iopub.execute_input":"2023-08-10T21:43:49.874629Z","iopub.status.idle":"2023-08-10T21:43:51.745931Z","shell.execute_reply.started":"2023-08-10T21:43:49.874595Z","shell.execute_reply":"2023-08-10T21:43:51.744941Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"token_list = open(f\"/kaggle/input/hf-wav2vec2-0-preprocess-baseline/vocab_bengali.json\",'r')\ntoken_list = json.load(token_list)\ntoken_list = {k: v for k, v in enumerate(token_list)}\ntoken_list ","metadata":{"execution":{"iopub.status.busy":"2023-08-10T21:43:51.747453Z","iopub.execute_input":"2023-08-10T21:43:51.748114Z","iopub.status.idle":"2023-08-10T21:43:51.761360Z","shell.execute_reply.started":"2023-08-10T21:43:51.748074Z","shell.execute_reply":"2023-08-10T21:43:51.760195Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from transformers import Wav2Vec2FeatureExtractor\n\nfeature_extractor = Wav2Vec2FeatureExtractor(feature_size=1, sampling_rate=SR, padding_value=0.0, do_normalize=True, return_attention_mask=True)\nfeature_extractor","metadata":{"execution":{"iopub.status.busy":"2023-08-10T21:43:51.763127Z","iopub.execute_input":"2023-08-10T21:43:51.763531Z","iopub.status.idle":"2023-08-10T21:43:51.775626Z","shell.execute_reply.started":"2023-08-10T21:43:51.763499Z","shell.execute_reply":"2023-08-10T21:43:51.774224Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from transformers import Wav2Vec2Processor\nprocessor = Wav2Vec2Processor(feature_extractor=feature_extractor, tokenizer=tokenizer)","metadata":{"execution":{"iopub.status.busy":"2023-08-10T21:43:51.776940Z","iopub.execute_input":"2023-08-10T21:43:51.777340Z","iopub.status.idle":"2023-08-10T21:43:51.789603Z","shell.execute_reply.started":"2023-08-10T21:43:51.777309Z","shell.execute_reply":"2023-08-10T21:43:51.788733Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class W2v2Dataset(torch.utils.data.Dataset):\n    def __init__(self, paths):\n        self.paths = paths\n        self.resampler = tat.Resample(32000, SR)\n\n    def __getitem__(self, idx):\n        apath = self.paths[idx]\n        waveform, sample_rate = torchaudio.load(apath, format=\"mp3\")\n        waveform = self.resampler(waveform)\n        audio = processor(waveform.reshape(-1), sampling_rate=SR).input_values[0]   \n        \n        id_name = self.paths[idx].split(os.sep)[-1].replace('.mp3','')\n        \n        return audio, id_name\n\n    def __len__(self):\n        return len(self.paths)\n\ntest_dataset = W2v2Dataset(paths)","metadata":{"execution":{"iopub.status.busy":"2023-08-10T21:43:51.793543Z","iopub.execute_input":"2023-08-10T21:43:51.793802Z","iopub.status.idle":"2023-08-10T21:43:51.889891Z","shell.execute_reply.started":"2023-08-10T21:43:51.793780Z","shell.execute_reply":"2023-08-10T21:43:51.888823Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Load Model","metadata":{}},{"cell_type":"code","source":"from transformers import Wav2Vec2ForCTC\nmodel = Wav2Vec2ForCTC.from_pretrained(model_dir)","metadata":{"execution":{"iopub.status.busy":"2023-08-10T21:43:51.891159Z","iopub.execute_input":"2023-08-10T21:43:51.892228Z","iopub.status.idle":"2023-08-10T21:44:13.045099Z","shell.execute_reply.started":"2023-08-10T21:43:51.892195Z","shell.execute_reply":"2023-08-10T21:44:13.044052Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')","metadata":{"execution":{"iopub.status.busy":"2023-08-10T21:44:13.046445Z","iopub.execute_input":"2023-08-10T21:44:13.046807Z","iopub.status.idle":"2023-08-10T21:44:13.078089Z","shell.execute_reply.started":"2023-08-10T21:44:13.046768Z","shell.execute_reply":"2023-08-10T21:44:13.075240Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model.to(device)","metadata":{"execution":{"iopub.status.busy":"2023-08-10T21:44:13.079408Z","iopub.execute_input":"2023-08-10T21:44:13.080151Z","iopub.status.idle":"2023-08-10T21:44:18.505974Z","shell.execute_reply.started":"2023-08-10T21:44:13.080105Z","shell.execute_reply":"2023-08-10T21:44:18.505017Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from torch.utils.data import DataLoader\ntest_loader = DataLoader(test_dataset,\n                             batch_size=1,\n                             shuffle=False,\n                             num_workers=os.cpu_count())","metadata":{"execution":{"iopub.status.busy":"2023-08-10T21:44:18.510515Z","iopub.execute_input":"2023-08-10T21:44:18.511199Z","iopub.status.idle":"2023-08-10T21:44:18.518467Z","shell.execute_reply.started":"2023-08-10T21:44:18.511161Z","shell.execute_reply":"2023-08-10T21:44:18.516779Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Inference","metadata":{}},{"cell_type":"code","source":"ids = []\nsentences = []","metadata":{"execution":{"iopub.status.busy":"2023-08-10T21:44:18.521718Z","iopub.execute_input":"2023-08-10T21:44:18.522559Z","iopub.status.idle":"2023-08-10T21:44:18.529089Z","shell.execute_reply.started":"2023-08-10T21:44:18.522520Z","shell.execute_reply":"2023-08-10T21:44:18.528068Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"for i, (aud, id_n) in enumerate(tqdm(test_loader)):\n    aud = aud.to(device)\n    logits = model(aud).logits\n    predicted_ids = torch.argmax(logits, dim=-1)\n    transcription = processor.decode(predicted_ids[0])\n    \n    transcription = transcription.replace('[UNK]', '').replace('[PAD]', '')\n    \n    ids.append(id_n[0])\n    sentences.append(transcription)","metadata":{"execution":{"iopub.status.busy":"2023-08-10T21:44:18.530852Z","iopub.execute_input":"2023-08-10T21:44:18.531682Z","iopub.status.idle":"2023-08-10T21:44:24.764747Z","shell.execute_reply.started":"2023-08-10T21:44:18.531649Z","shell.execute_reply":"2023-08-10T21:44:24.763392Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"submission = pd.DataFrame()\nsubmission['id'] = ids\nsubmission['sentence'] = sentences\nsubmission.to_csv(\"submission.csv\", index=False)\nsubmission.head()","metadata":{"execution":{"iopub.status.busy":"2023-08-10T21:44:24.766353Z","iopub.execute_input":"2023-08-10T21:44:24.766757Z","iopub.status.idle":"2023-08-10T21:44:24.796596Z","shell.execute_reply.started":"2023-08-10T21:44:24.766726Z","shell.execute_reply":"2023-08-10T21:44:24.795509Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Done!","metadata":{}}]}