{"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"},"kaggle":{"accelerator":"none","dataSources":[{"sourceId":52324,"databundleVersionId":6229904,"sourceType":"competition"}],"dockerImageVersionId":30528,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"%%capture\n! pip install transformers\n! pip install jiwer\n! pip install --upgrade wandb\n! pip install --upgrade librosa","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2023-08-22T11:29:55.777455Z","iopub.execute_input":"2023-08-22T11:29:55.778009Z","iopub.status.idle":"2023-08-22T11:30:53.755960Z","shell.execute_reply.started":"2023-08-22T11:29:55.777959Z","shell.execute_reply":"2023-08-22T11:30:53.754586Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import os\nimport json\nimport numpy as np\nimport pandas as pd\nfrom tqdm.auto import tqdm\n\nimport torch\nimport torch.nn as nn\nfrom torch.utils.data import Dataset, DataLoader\nfrom torch.optim.lr_scheduler import CosineAnnealingWarmRestarts\n\nimport wandb\nfrom kaggle_secrets import UserSecretsClient\nfrom datetime import datetime\nfrom wandb.keras import WandbCallback\n\nimport librosa\n\nimport warnings\nwarnings.simplefilter('ignore')","metadata":{"execution":{"iopub.status.busy":"2023-08-22T15:19:10.802974Z","iopub.execute_input":"2023-08-22T15:19:10.803703Z","iopub.status.idle":"2023-08-22T15:19:13.586084Z","shell.execute_reply.started":"2023-08-22T15:19:10.803659Z","shell.execute_reply":"2023-08-22T15:19:13.584056Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"<img src=\"https://camo.githubusercontent.com/dd842f7b0be57140e68b2ab9cb007992acd131c48284eaf6b1aca758bfea358b/68747470733a2f2f692e696d6775722e636f6d2f52557469567a482e706e67\">\n\n> I will be integrating W&B for visualizations and logging artifacts!\n> \n> [SEW - Bengali.ai Speech Recognition on W&B Dashboard](https://wandb.ai/usharengaraju/SEW-Bengali.ai)\n> \n> - To get the API key, create an account in the [website](https://wandb.ai/site) .\n> - Use secrets to use API Keys more securely ","metadata":{}},{"cell_type":"code","source":"# Setup user secrets for login\nuser_secrets = UserSecretsClient()\nwandb_api = user_secrets.get_secret(\"api_key\") \n\n# Login\nwandb.login(key = wandb_api)\n\nrun = wandb.init(project = \"SEW-Bengali.ai\",\n                 name = f\"Run_{datetime.now().strftime('%d%m%Y%H%M%S')}\", \n                 notes = \"add some features\",\n                 tags = [],\n                 config = dict(competition = 'Bengali.ai',\n                               _wandb_kernel = 'tensorgirl',\n                               batch_size = 32,\n                               epochs = 30,\n                               learning_rate = 0.005)\n)\n","metadata":{"execution":{"iopub.status.busy":"2023-08-22T15:19:19.411936Z","iopub.execute_input":"2023-08-22T15:19:19.412751Z","iopub.status.idle":"2023-08-22T15:19:58.025667Z","shell.execute_reply.started":"2023-08-22T15:19:19.412714Z","shell.execute_reply":"2023-08-22T15:19:58.024675Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from transformers import (\n    Wav2Vec2ForCTC,\n    Wav2Vec2Processor,\n    Wav2Vec2CTCTokenizer,\n    Wav2Vec2FeatureExtractor\n) ","metadata":{"execution":{"iopub.status.busy":"2023-08-22T11:37:59.311428Z","iopub.execute_input":"2023-08-22T11:37:59.311809Z","iopub.status.idle":"2023-08-22T11:38:09.452637Z","shell.execute_reply.started":"2023-08-22T11:37:59.311780Z","shell.execute_reply":"2023-08-22T11:38:09.451626Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"Config = {\n    'audio_dir': '/kaggle/input/bengaliai-speech/train_mp3s',\n    'model_name': 'facebook/wav2vec2-base',\n    'lr': 3e-4,\n    'wd': 1e-5,\n    'T_0': 10,\n    'T_mult': 2,\n    'eta_min': 1e-6,\n    'nb_epochs': 5,\n    'train_bs': 16,\n    'valid_bs': 16,\n    'sampling_rate': 16000,\n}","metadata":{"execution":{"iopub.status.busy":"2023-08-22T15:20:22.294839Z","iopub.execute_input":"2023-08-22T15:20:22.295214Z","iopub.status.idle":"2023-08-22T15:20:22.300501Z","shell.execute_reply.started":"2023-08-22T15:20:22.295182Z","shell.execute_reply":"2023-08-22T15:20:22.299573Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def read_audio(mp3_path, target_sr=16000):\n    audio, sr = librosa.load(mp3_path, sr=32000)\n    audio_array = librosa.resample(audio, orig_sr=sr, target_sr=target_sr)\n    return audio_array\n\ndef construct_vocab(texts):\n    all_text = \" \".join(texts)\n    vocab = list(set(all_text))\n    return vocab\n\ndef wandb_log(**kwargs):\n    for k, v in kwargs.items():\n        wandb.log({k: v})\n\ndef save_vocab(dataframe):\n    vocab = construct_vocab(dataframe['sentence'].tolist())\n    vocab_dict = {v: k for k, v in enumerate(vocab)}\n    vocab_dict[\"__\"] = vocab_dict[\" \"]\n    _ = vocab_dict.pop(\" \")\n    vocab_dict[\"[UNK]\"] = len(vocab_dict)\n    vocab_dict[\"[PAD]\"] = len(vocab_dict)\n\n    with open('vocab.json', 'w') as fl:\n        json.dump(vocab_dict, fl)\n\n    print(\"Created Vocab file!\")","metadata":{"execution":{"iopub.status.busy":"2023-08-22T15:20:28.871682Z","iopub.execute_input":"2023-08-22T15:20:28.872022Z","iopub.status.idle":"2023-08-22T15:20:28.880573Z","shell.execute_reply.started":"2023-08-22T15:20:28.871996Z","shell.execute_reply":"2023-08-22T15:20:28.879685Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class ASRDataset(Dataset):\n    def __init__(self, df, config, is_test=False):\n        self.df = df\n        self.config = config\n        self.is_test = is_test\n    \n    def __getitem__(self, idx):\n        # First read and pre-process the audio file\n        audio = read_audio(self.df.loc[idx]['path'])\n        audio = processor(\n            audio, \n            sampling_rate=self.config['sampling_rate']\n        ).input_values[0]\n        \n        if self.is_test:\n            return {'audio': audio, 'label': -1}\n        else:\n            # If we are training/validating, also process the labels (actual sentences)\n            with processor.as_target_processor():\n                labels = processor(self.df.loc[idx]['sentence']).input_ids\n            return {'audio': audio, 'label': labels}\n        \n    def __len__(self):\n        return len(self.df)\n    \ndef ctc_data_collator(batch):\n    input_features = [{\"input_values\": sample[\"audio\"]} for sample in batch]\n    label_features = [{\"input_ids\": sample[\"label\"]} for sample in batch]\n    batch = processor.pad(\n        input_features,\n        padding=True,\n        return_tensors=\"pt\",\n    )\n    with processor.as_target_processor():\n        labels_batch = processor.pad(\n            label_features,\n            padding=True,\n            return_tensors=\"pt\",\n        )\n        \n    labels = labels_batch[\"input_ids\"].masked_fill(labels_batch.attention_mask.ne(1), -100)\n    batch[\"labels\"] = labels\n    return batch","metadata":{"execution":{"iopub.status.busy":"2023-08-22T15:20:38.211232Z","iopub.execute_input":"2023-08-22T15:20:38.211604Z","iopub.status.idle":"2023-08-22T15:20:38.221812Z","shell.execute_reply.started":"2023-08-22T15:20:38.211575Z","shell.execute_reply":"2023-08-22T15:20:38.220430Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def train_one_epoch(model, train_loader, optimizer, device='cuda:0'):\n    model.train()\n    pbar = tqdm(train_loader, total=len(train_loader))\n    avg_loss = 0\n    for data in pbar:\n        data = {k: v.to(device) for k, v in data.items()}\n        loss = model(**data).loss\n        loss_itm = loss.item()\n        \n        avg_loss += loss_itm\n        pbar.set_description(f\"loss: {loss_itm:.4f}\")\n        wandb_log(train_step_loss=loss_itm)\n        \n        optimizer.zero_grad(set_to_none=True)\n        loss.backward()\n        optimizer.step()\n        \n    return avg_loss / len(train_loader)\n\n@torch.no_grad()\ndef valid_one_epoch(model, valid_loader, device='cuda:0'):\n    pbar = tqdm(valid_loader, total=len(valid_loader))\n    avg_loss = 0\n    for data in pbar:\n        data = {k: v.to(device) for k, v in data.items()}\n        loss = model(**data).loss\n        loss_itm = loss.item()\n        \n        avg_loss += loss_itm\n        pbar.set_description(f\"val_loss: {loss_itm:.4f}\")\n        wandb_log(valid_step_loss=loss_itm)\n\n    return avg_loss / len(valid_loader)","metadata":{"execution":{"iopub.status.busy":"2023-08-22T15:20:50.224442Z","iopub.execute_input":"2023-08-22T15:20:50.224831Z","iopub.status.idle":"2023-08-22T15:20:50.233873Z","shell.execute_reply.started":"2023-08-22T15:20:50.224798Z","shell.execute_reply":"2023-08-22T15:20:50.232820Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# **<span style=\"color:#F7B2B0;\">SEW (Squeezed and Efficient Wav2vec) </span>**\n\n[Source1](https://arxiv.org/pdf/2109.06870.pdf) \n[Source2](https://huggingface.co/docs/transformers/model_doc/sew)\n\nSEW (Squeezed and Efficient Wav2vec) proposed in [Performance-Efficiency Trade-offs in Unsupervised Pre-training for Speech Recognition](https://arxiv.org/pdf/2109.06870.pdf) by Felix Wu, Kwangyoun Kim, Jing Pan, Kyu Han, Kilian Q. Weinberger, Yoav Artzi introduces a pre-trained model architecture with significant improvements along both performance and efficiency dimensions across a variety of training setups. SEW achieves a 1.9x inference speedup compared to wav2vec 2.0, with a 13.5% relative reduction in word error rate. \n\nW2V2 consists of a waveform feature extractor and a context network. The waveform feature extractor generates a sequence of continuous feature vectors, each encoding a small segment of audio, and a context network that maps these vectors to context-dependent representations. \nThe features extracted by the waveform feature extractor are masked out during pre-trained and discretized as prediction targets and these features are not seen by the context network.\n","metadata":{}},{"cell_type":"markdown","source":"![](https://i.imgur.com/Oet3lDD.png)","metadata":{}},{"cell_type":"code","source":"run = wandb.init(\n    project='bengalispeech',\n    config=Config,\n    name='sew',\n)","metadata":{"_kg_hide-output":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from transformers import SEWForCTC\ndf = pd.read_csv(\"/kaggle/input/bengaliai-speech/train.csv\")\n\n# Get a paths feature for reading in during dataloading\ndf['path'] = df['id'].apply(lambda x: os.path.join(Config['audio_dir'], x+'.mp3'))\ntrain_df = df[df['split'] == 'train'].sample(frac=.005).reset_index(drop=True)\nvalid_df = df[df['split'] == 'valid'].sample(frac=.005).reset_index(drop=True)\nprint(f\"Training on samples: {len(train_df)}, Validation on samples: {len(valid_df)}\")\n\n# Construct and save the vocab file\nsave_vocab(df)\n\n# Init the tokenizer, feature_extractor, processor and model\ntokenizer = Wav2Vec2CTCTokenizer(\n    \"./vocab.json\", \n    unk_token=\"[UNK]\",\n    pad_token=\"[PAD]\",\n    word_delimiter_token=\"__\"\n)\nfeature_extractor = Wav2Vec2FeatureExtractor(\n    feature_size=1, \n    sampling_rate=Config['sampling_rate'], \n    padding_value=0.0, \n    do_normalize=True, \n    return_attention_mask=False\n)\nprocessor = Wav2Vec2Processor(\n    feature_extractor=feature_extractor, \n    tokenizer=tokenizer\n)\n\nmodel = SEWForCTC.from_pretrained(\"asapp/sew-tiny-100k-ft-ls100h\",\n    ctc_loss_reduction=\"mean\", \n    ignore_mismatched_sizes=True,\n    pad_token_id=processor.tokenizer.pad_token_id,\n    vocab_size = len(tokenizer),\n)\nwandb.watch(model)\n\n# Freeze the feature encoder part since we won't be training it\nmodel.to('cuda')\nmodel.freeze_feature_encoder()\noptimizer = torch.optim.AdamW(\n    model.parameters(), \n    lr=Config['lr'], \n    weight_decay=Config['wd']\n)\n\n# Construct training and validation dataloaders\ntrain_ds = ASRDataset(train_df, Config)\nvalid_ds = ASRDataset(valid_df, Config)\n\ntrain_loader = DataLoader(\n    train_ds, \n    batch_size=Config['train_bs'], \n    collate_fn=ctc_data_collator, \n)\nvalid_loader = DataLoader(\n    valid_ds,\n    batch_size=Config['valid_bs'],\n    collate_fn=ctc_data_collator,\n)\n\n# Train the model\nbest_loss = float('inf')\nfor epoch in range(Config['nb_epochs']):\n    print(f\"{'='*40} Epoch: {epoch+1} / {Config['nb_epochs']} {'='*40}\")\n    train_loss = train_one_epoch(model, train_loader, optimizer)\n    valid_loss = valid_one_epoch(model, valid_loader)\n    wandb_log(train_loss=train_loss, val_loss=valid_loss)\n    print(f\"train_loss: {train_loss:.4f}, valid_loss: {valid_loss:.4f}\")\n\n    if valid_loss < best_loss:\n        best_loss = valid_loss\n        torch.save(model.state_dict(), f\"sew_base_bengaliAI.pt\")\n        print(f\"Saved the best model so far with val_loss: {valid_loss:.4f}\")","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Once training is done, \nwandb.finish()","metadata":{"_kg_hide-output":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# **<span style=\"color:#F7B2B0;\">Tune Hyperparameters with Wandb Sweeps </span>**\n\nUse W&B Sweeps to automate hyperparameter search and visualize rich, interactive experiment tracking. Pick from popular search methods such as Bayesian, grid search, and random to search the hyperparameter space. Scale and parallelize sweep across one or more machines.","metadata":{}},{"cell_type":"code","source":"Config['train_bs']=32\nConfig['val_bs']=32","metadata":{"execution":{"iopub.status.busy":"2023-08-22T11:13:34.884897Z","iopub.status.idle":"2023-08-22T11:13:34.885486Z","shell.execute_reply.started":"2023-08-22T11:13:34.885181Z","shell.execute_reply":"2023-08-22T11:13:34.885207Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df = pd.read_csv(\"/kaggle/input/bengaliai-speech/train.csv\")\n\n# Get a paths feature for reading in during dataloading\ndf['path'] = df['id'].apply(lambda x: os.path.join(Config['audio_dir'], x+'.mp3'))\ntrain_df = df[df['split'] == 'train'].sample(frac=.005).reset_index(drop=True)\nvalid_df = df[df['split'] == 'valid'].sample(frac=.005).reset_index(drop=True)\nprint(f\"Training on samples: {len(train_df)}, Validation on samples: {len(valid_df)}\")\n\n# Construct and save the vocab file\nsave_vocab(df)\n\n# Init the tokenizer, feature_extractor, processor and model\ntokenizer = Wav2Vec2CTCTokenizer(\n    \"./vocab.json\", \n    unk_token=\"[UNK]\",\n    pad_token=\"[PAD]\",\n    word_delimiter_token=\"__\"\n)\nfeature_extractor = Wav2Vec2FeatureExtractor(\n    feature_size=1, \n    sampling_rate=Config['sampling_rate'], \n    padding_value=0.0, \n    do_normalize=True, \n    return_attention_mask=False\n)\nprocessor = Wav2Vec2Processor(\n    feature_extractor=feature_extractor, \n    tokenizer=tokenizer\n)\n# Construct training and validation dataloaders\ntrain_ds = ASRDataset(train_df, Config)\nvalid_ds = ASRDataset(valid_df, Config)\n\ntrain_loader = DataLoader(\n    train_ds, \n    batch_size=Config['train_bs'], \n    collate_fn=ctc_data_collator, \n)\nvalid_loader = DataLoader(\n    valid_ds,\n    batch_size=Config['valid_bs'],\n    collate_fn=ctc_data_collator,\n)\n      \n\ndef main():\n    run = wandb.init(project = 'bengalispeech',name='sweep')\n    \n    model = SEWForCTC.from_pretrained(\"asapp/sew-tiny-100k-ft-ls100h\",\n    ctc_loss_reduction=\"mean\", \n    pad_token_id=processor.tokenizer.pad_token_id,\n    vocab_size = len(tokenizer),\n    )\n    wandb.watch(model)\n\n    # Freeze the feature encoder part since we won't be training it\n    model.to('cuda')\n    model.freeze_feature_encoder()\n    optimizer = torch.optim.AdamW(\n        model.parameters(), \n        lr=wandb.config.learning_rate, \n        weight_decay=Config['wd']\n    )\n\n\n    # Train the model\n    best_loss = float('inf')\n    for epoch in range(wandb.config.n_epochs):\n        print(f\"{'='*40} Epoch: {epoch+1} / {wandb.config.n_epochs} {'='*40}\")\n        train_loss = train_one_epoch(model, train_loader, optimizer)\n        valid_loss = valid_one_epoch(model, valid_loader)\n        wandb_log(train_loss=train_loss, val_loss=valid_loss)\n        print(f\"train_loss: {train_loss:.4f}, valid_loss: {valid_loss:.4f}\")\n\n        if valid_loss < best_loss:\n            best_loss = valid_loss\n            torch.save(model.state_dict(), f\"sewtiny.pt\")\n            print(f\"Saved the best model so far with val_loss: {valid_loss:.4f}\")\n        \n    wandb.finish()\n    \nsweep_configuration = {\n    'method': 'bayes',  # random, grid or bayes\n    'name': 'sweep-bayes',\n    'metric': {'goal': 'minimize', 'name': 'loss'},\n    'parameters': \n    {\n        'n_epochs': {'values': [5, 10]},\n        'learning_rate': {'max': 0.001, 'min': 0.00001},\n     }\n}\nsweep_id = wandb.sweep(sweep=sweep_configuration,project='bengalispeech')\nwandb.agent(sweep_id, function=main, count=10)","metadata":{"execution":{"iopub.status.busy":"2023-08-22T11:13:34.887429Z","iopub.status.idle":"2023-08-22T11:13:34.888052Z","shell.execute_reply.started":"2023-08-22T11:13:34.887728Z","shell.execute_reply":"2023-08-22T11:13:34.887776Z"},"_kg_hide-output":true,"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# **<span style=\"color:#F7B2B0;\">WavLM: Large-Scale Self-Supervised Pre-Training for Full Stack Speech Processing </span>**\n\n[Source](https://arxiv.org/pdf/2110.13900.pdf)\n\n\nSelf-supervised learning (SSL) achieves great success in speech recognition, while limited exploration has been attempted for other speech processing tasks. As speech signal contains multi-faceted information including speaker identity,paralinguistics, spoken content, etc., learning universal representations for all speech tasks is challenging. To tackle the problem, we propose a new pre-trained model, WavLM, to solve full-stack downstream speech tasks. WavLM jointly learns\nmasked speech prediction and denoising in pre-training. By this means, WavLM does not only keep the speech content modeling capability by the masked speech prediction, but also improves the\npotential to non-ASR tasks by the speech denoising. In addition,WavLM employs gated relative position bias for the Transformer structure to better capture the sequence ordering of input speech.\nWe also scale up the training dataset from 60k hours to 94k hours. WavLM Large achieves state-of-the-art performance on the SUPERB benchmark, and brings significant improvements for various speech processing tasks on their representative benchmarks.","metadata":{}},{"cell_type":"markdown","source":"![](https://i.imgur.com/RY4Xw3J.png)","metadata":{}},{"cell_type":"code","source":"run = wandb.init(\n    project='bengalispeech',\n    config=Config,\n    name='wavlm',\n)","metadata":{"_kg_hide-output":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"Config['train_bs']=4\nConfig['val_bs']=4","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from transformers import WavLMForCTC\ndf = pd.read_csv(\"/kaggle/input/bengaliai-speech/train.csv\")\n\n# Get a paths feature for reading in during dataloading\ndf['path'] = df['id'].apply(lambda x: os.path.join(Config['audio_dir'], x+'.mp3'))\ntrain_df = df[df['split'] == 'train'].sample(frac=.005).reset_index(drop=True)\nvalid_df = df[df['split'] == 'valid'].sample(frac=.005).reset_index(drop=True)\nprint(f\"Training on samples: {len(train_df)}, Validation on samples: {len(valid_df)}\")\n\n# Construct and save the vocab file\nsave_vocab(df)\n\ntokenizer = Wav2Vec2CTCTokenizer(\n    \"./vocab.json\", \n    unk_token=\"[UNK]\",\n    pad_token=\"[PAD]\",\n    word_delimiter_token=\"__\"\n)\nfeature_extractor = Wav2Vec2FeatureExtractor(\n    feature_size=1, \n    sampling_rate=Config['sampling_rate'], \n    padding_value=0.0, \n    do_normalize=True, \n    return_attention_mask=False\n)\nprocessor = Wav2Vec2Processor(\n    feature_extractor=feature_extractor, \n    tokenizer=tokenizer\n)\n\nmodel = WavLMForCTC.from_pretrained(\"patrickvonplaten/wavlm-libri-clean-100h-base-plus\",\n    ctc_loss_reduction=\"mean\", \n    ignore_mismatched_sizes=True,\n    pad_token_id=processor.tokenizer.pad_token_id,\n    vocab_size = len(tokenizer),\n)\nwandb.watch(model)\n\n# Freeze the feature encoder part since we won't be training it\nmodel.to('cuda')\nmodel.freeze_feature_encoder()\noptimizer = torch.optim.AdamW(\n    model.parameters(), \n    lr=Config['lr'], \n    weight_decay=Config['wd']\n)\n\n# Construct training and validation dataloaders\ntrain_ds = ASRDataset(train_df, Config)\nvalid_ds = ASRDataset(valid_df, Config)\n\ntrain_loader = DataLoader(\n    train_ds, \n    batch_size=Config['train_bs'], \n    collate_fn=ctc_data_collator, \n)\nvalid_loader = DataLoader(\n    valid_ds,\n    batch_size=Config['valid_bs'],\n    collate_fn=ctc_data_collator,\n)\n\n# Train the model\nbest_loss = float('inf')\nfor epoch in range(Config['nb_epochs']):\n    print(f\"{'='*40} Epoch: {epoch+1} / {Config['nb_epochs']} {'='*40}\")\n    train_loss = train_one_epoch(model, train_loader, optimizer)\n    valid_loss = valid_one_epoch(model, valid_loader)\n    wandb_log(train_loss=train_loss, val_loss=valid_loss)\n    print(f\"train_loss: {train_loss:.4f}, valid_loss: {valid_loss:.4f}\")\n\n    if valid_loss < best_loss:\n        best_loss = valid_loss\n        torch.save(model.state_dict(), f\"wavlm_base_bengaliAI.pt\")\n        print(f\"Saved the best model so far with val_loss: {valid_loss:.4f}\")","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"wandb.finish()","metadata":{"_kg_hide-output":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# **<span style=\"color:#F7B2B0;\">References </span>**\n\nAll public notebooks of this competition\n\nhttps://arxiv.org/pdf/2109.06870.pdf\n\nhttps://huggingface.co/docs/transformers/model_doc/sew\n\nhttps://arxiv.org/pdf/2110.13900.pdf\n","metadata":{}}]}