{"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"name":"python","version":"3.11.11","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"none","dataSources":[{"sourceId":91844,"databundleVersionId":11361821,"sourceType":"competition"},{"sourceId":11458343,"sourceType":"datasetVersion","datasetId":6930139},{"sourceId":12045181,"sourceType":"datasetVersion","datasetId":7243711},{"sourceId":12046911,"sourceType":"datasetVersion","datasetId":6930285},{"sourceId":243313432,"sourceType":"kernelVersion"},{"sourceId":243446918,"sourceType":"kernelVersion"}],"dockerImageVersionId":31012,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":false},"papermill":{"default_parameters":{},"duration":64.08478,"end_time":"2025-04-16T20:13:48.159484","environment_variables":{},"exception":null,"input_path":"__notebook__.ipynb","output_path":"__notebook__.ipynb","parameters":{},"start_time":"2025-04-16T20:12:44.074704","version":"2.5.0"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!nvidia-smi","metadata":{"execution":{"iopub.status.busy":"2025-06-03T01:23:42.390745Z","iopub.execute_input":"2025-06-03T01:23:42.391103Z","iopub.status.idle":"2025-06-03T01:23:42.520695Z","shell.execute_reply.started":"2025-06-03T01:23:42.391045Z","shell.execute_reply":"2025-06-03T01:23:42.519437Z"},"papermill":{"duration":0.136763,"end_time":"2025-04-16T20:12:47.634449","exception":false,"start_time":"2025-04-16T20:12:47.497686","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!pip install /kaggle/input/bird-clef-2025-addones/openvino-2025.0.0-17942-cp311-cp311-manylinux2014_x86_64.whl --no-deps\n# !pip install /kaggle/input/bird-clef-2024-addones/onnxruntime-1.17.3-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl --no-deps\n!pip list | grep openvino\n!pip list | grep numpy","metadata":{"execution":{"iopub.status.busy":"2025-06-03T01:23:42.522737Z","iopub.execute_input":"2025-06-03T01:23:42.523122Z","iopub.status.idle":"2025-06-03T01:23:53.228372Z","shell.execute_reply.started":"2025-06-03T01:23:42.523075Z","shell.execute_reply":"2025-06-03T01:23:53.227145Z"},"papermill":{"duration":4.08287,"end_time":"2025-04-16T20:12:51.724666","exception":false,"start_time":"2025-04-16T20:12:47.641796","status":"completed"},"tags":[],"trusted":true,"jupyter":{"source_hidden":true},"_kg_hide-input":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import sys\nsys.path.append('/kaggle/input/bird-clef-2025-code/main_folder/main_folder/')","metadata":{"execution":{"iopub.status.busy":"2025-06-03T01:23:53.229671Z","iopub.execute_input":"2025-06-03T01:23:53.229998Z","iopub.status.idle":"2025-06-03T01:23:53.235675Z","shell.execute_reply.started":"2025-06-03T01:23:53.229952Z","shell.execute_reply":"2025-06-03T01:23:53.234719Z"},"papermill":{"duration":0.013862,"end_time":"2025-04-16T20:12:56.409308","exception":false,"start_time":"2025-04-16T20:12:56.395446","status":"completed"},"tags":[],"trusted":true,"_kg_hide-input":true,"jupyter":{"source_hidden":true}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import pandas as pd\nimport numpy as np\nimport librosa\nimport seaborn as sns\nimport os\nimport json\nimport IPython.display as ipd\nimport soundfile as sf\nimport torch\nimport h5py\n# import onnxruntime as ort\nimport openvino as ov\n\nfrom glob import glob\nfrom tqdm import tqdm\nfrom matplotlib import pyplot as plt\nfrom itertools import chain\nfrom os.path import join as pjoin\nfrom copy import deepcopy\n\n\nfrom code_base.datasets import WaveDataset, WaveAllFileDataset\nfrom code_base.inefernce import BirdsInference\nfrom code_base.utils import load_json, compose_submission_dataframe\n%matplotlib inline","metadata":{"execution":{"iopub.status.busy":"2025-06-03T01:23:53.236568Z","iopub.execute_input":"2025-06-03T01:23:53.236891Z","iopub.status.idle":"2025-06-03T01:24:01.748558Z","shell.execute_reply.started":"2025-06-03T01:23:53.236855Z","shell.execute_reply":"2025-06-03T01:24:01.747551Z"},"papermill":{"duration":7.508944,"end_time":"2025-04-16T20:13:03.924934","exception":false,"start_time":"2025-04-16T20:12:56.41599","status":"completed"},"tags":[],"trusted":true,"jupyter":{"source_hidden":true}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Config","metadata":{"papermill":{"duration":0.006128,"end_time":"2025-04-16T20:13:03.942799","exception":false,"start_time":"2025-04-16T20:13:03.936671","status":"completed"},"tags":[]}},{"cell_type":"code","source":"POSTFIX = \"\"\nDATASET_NAME = \"bird-clef-2025-models-v2\"\nEXP_NAME = \"tf_efficientnetv2_s_in21k_Exp_noamp_64bs_5sec_BasicAug_EqualBalancing_AdamW1e4_CosBatchLR1e6_Epoch50_FocalBCELoss_LSF1005_FromPrebs1_PseudoF2PT05MT01P04I2_AddRareBirdsNoLeak\" + POSTFIX\nTRAIN_PERIOD = 5\nprint(\"Possible checkpoints:\\n\\n{}\".format(\"\\n\".join(set([\n    os.path.basename(el) for el in glob(f\"/kaggle/input/{DATASET_NAME}/{EXP_NAME}/{EXP_NAME}/**\", recursive=True) if \"train\" not in os.path.basename(el)\n]))))","metadata":{"execution":{"iopub.status.busy":"2025-06-03T01:24:01.750944Z","iopub.execute_input":"2025-06-03T01:24:01.751776Z","iopub.status.idle":"2025-06-03T01:24:01.773674Z","shell.execute_reply.started":"2025-06-03T01:24:01.751751Z","shell.execute_reply":"2025-06-03T01:24:01.772319Z"},"papermill":{"duration":0.027164,"end_time":"2025-04-16T20:13:03.976496","exception":false,"start_time":"2025-04-16T20:13:03.949332","status":"completed"},"tags":[],"trusted":true,"_kg_hide-input":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"IS_TEST = len(glob(\"/kaggle/input/birdclef-2025/test_soundscapes/*.ogg\")) > 1\nIS_TEST","metadata":{"execution":{"iopub.status.busy":"2025-06-03T01:24:01.775109Z","iopub.execute_input":"2025-06-03T01:24:01.776143Z","iopub.status.idle":"2025-06-03T01:24:01.790705Z","shell.execute_reply.started":"2025-06-03T01:24:01.776103Z","shell.execute_reply":"2025-06-03T01:24:01.789236Z"},"papermill":{"duration":0.020578,"end_time":"2025-04-16T20:13:04.009088","exception":false,"start_time":"2025-04-16T20:13:03.98851","status":"completed"},"tags":[],"trusted":true,"jupyter":{"source_hidden":true},"_kg_hide-input":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"CONFIG = {\n    # Inference Class\n    \"use_sigmoid\": False,\n    \"use_compiled_fp16\": False,\n    # Data config\n    \"test_data_root\": \"/kaggle/input/birdclef-2025/test_soundscapes/*.ogg\" if IS_TEST else \"/kaggle/input/birdclef-2025/train_soundscapes/*.ogg\",\n    \"label_map_data_path\":'/kaggle/input/bird-clef-2025-models-v2/bird2int_2025.json',\n    \"lookback\":None,\n    \"lookahead\":None,\n    \"segment_len\":5,\n    \"step\": None,\n    \"late_normalize\": True,\n    \"load_normalize\": False,\n    # Model config\n    \"exp_name\":EXP_NAME,\n    # Post Processing\n    \"global_target_smoothing_coef\": None\n}\n\nif CONFIG.get(\"use_sed_mode\", False):\n    assert CONFIG[\"step\"] is not None\nelse:\n    assert CONFIG[\"step\"] is None\n","metadata":{"execution":{"iopub.status.busy":"2025-06-03T01:24:01.792044Z","iopub.execute_input":"2025-06-03T01:24:01.792474Z","iopub.status.idle":"2025-06-03T01:24:01.80015Z","shell.execute_reply.started":"2025-06-03T01:24:01.792436Z","shell.execute_reply":"2025-06-03T01:24:01.798683Z"},"papermill":{"duration":0.015912,"end_time":"2025-04-16T20:13:04.031851","exception":false,"start_time":"2025-04-16T20:13:04.015939","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Data","metadata":{"papermill":{"duration":0.006282,"end_time":"2025-04-16T20:13:04.045411","exception":false,"start_time":"2025-04-16T20:13:04.039129","status":"completed"},"tags":[]}},{"cell_type":"code","source":"bird2id = load_json(CONFIG[\"label_map_data_path\"])","metadata":{"execution":{"iopub.status.busy":"2025-06-03T01:24:01.801626Z","iopub.execute_input":"2025-06-03T01:24:01.802097Z","iopub.status.idle":"2025-06-03T01:24:01.830159Z","shell.execute_reply.started":"2025-06-03T01:24:01.802039Z","shell.execute_reply":"2025-06-03T01:24:01.829161Z"},"papermill":{"duration":0.022074,"end_time":"2025-04-16T20:13:04.075989","exception":false,"start_time":"2025-04-16T20:13:04.053915","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"test_au_pathes = glob(CONFIG[\"test_data_root\"])\nif not IS_TEST:\n    test_au_pathes = test_au_pathes[:10]   ###### 700 for performance test\n\ntest_df = pd.DataFrame({\n    \"filename\": test_au_pathes,\n    #\"duration_s\": [librosa.get_duration(filename=el) for el in test_au_pathes]\n    \"duration_s\": [60.0 for el in test_au_pathes]  # Fernando: prior is slow and all files have 60.0s\n})\nif not IS_TEST:\n    display(test_df)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-06-03T01:24:01.831157Z","iopub.execute_input":"2025-06-03T01:24:01.831418Z","iopub.status.idle":"2025-06-03T01:24:01.994181Z","shell.execute_reply.started":"2025-06-03T01:24:01.831398Z","shell.execute_reply":"2025-06-03T01:24:01.993114Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"ds_config_test = {\n   \"root\": \"\",\n   \"label_str2int_mapping_path\": CONFIG[\"label_map_data_path\"],\n   \"n_cores\": 64,\n   \"use_audio_cache\": True,\n   \"test_mode\": True,\n   \"segment_len\": CONFIG[\"segment_len\"],\n   \"lookback\":CONFIG[\"lookback\"],\n   \"lookahead\":CONFIG[\"lookahead\"],\n    \"sample_id\": None,\n    \"late_normalize\": CONFIG[\"late_normalize\"],\n    \"load_normalize\": CONFIG.get(\"load_normalize\", True),\n    \"step\": CONFIG[\"step\"],\n    \"validate_sr\": 32_000\n}\nloader_config = {\n    \"batch_size\": 24,\n    \"drop_last\": False,\n    \"shuffle\": False,\n    \"num_workers\": 0,\n}","metadata":{"execution":{"iopub.status.busy":"2025-06-03T01:24:01.99528Z","iopub.execute_input":"2025-06-03T01:24:01.995531Z","iopub.status.idle":"2025-06-03T01:24:02.002274Z","shell.execute_reply.started":"2025-06-03T01:24:01.995511Z","shell.execute_reply":"2025-06-03T01:24:02.001181Z"},"papermill":{"duration":0.015432,"end_time":"2025-04-16T20:13:09.164632","exception":false,"start_time":"2025-04-16T20:13:09.1492","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"ds_test = WaveAllFileDataset(df=test_df, **ds_config_test)\nloader_test = torch.utils.data.DataLoader(\n    ds_test,\n    **loader_config,\n)","metadata":{"execution":{"iopub.status.busy":"2025-06-03T01:24:02.003288Z","iopub.execute_input":"2025-06-03T01:24:02.003619Z","iopub.status.idle":"2025-06-03T01:24:03.229192Z","shell.execute_reply.started":"2025-06-03T01:24:02.003591Z","shell.execute_reply":"2025-06-03T01:24:03.22811Z"},"papermill":{"duration":0.02197,"end_time":"2025-04-16T20:13:09.193678","exception":false,"start_time":"2025-04-16T20:13:09.171708","status":"completed"},"tags":[],"trusted":true,"jupyter":{"source_hidden":true}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Model","metadata":{"papermill":{"duration":0.006641,"end_time":"2025-04-16T20:13:09.207448","exception":false,"start_time":"2025-04-16T20:13:09.200807","status":"completed"},"tags":[]}},{"cell_type":"code","source":"def compile_model(exp_name, postfix=POSTFIX):\n    core = ov.Core()\n    model = core.read_model(model=f\"/kaggle/input/{DATASET_NAME}/{exp_name}/{exp_name}/onnx_ensem_5first_folds{postfix}_openvino_fp16/model_simpl.xml\")\n    return core.compile_model(model=model, device_name=\"CPU\")\n\ncompiled_models = [compile_model(e_n) for e_n in [\n    \"tf_efficientnetv2_s_in21k_Exp_noamp_64bs_5sec_BasicAug_EqualBalancing_AdamW1e4_CosBatchLR1e6_Epoch50_FocalBCELoss_LSF1005_FromPrebs1_PseudoF2PT05MT01P04I2_AddRareBirdsNoLeak\",\n    \"eca_nfnet_l0_Exp_noamp_64bs_5sec_BasicAug_SqrtBalancing_Radamlr1e3_CosBatchLR1e6_Epoch50_FocalBCELoss_LSF1005_FromXCV2Best_PseudoF2PT05MT01P04I3_MinorOverSampleV1\"\n]]","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-06-03T01:24:03.23022Z","iopub.execute_input":"2025-06-03T01:24:03.230497Z","iopub.status.idle":"2025-06-03T01:24:14.740698Z","shell.execute_reply.started":"2025-06-03T01:24:03.230474Z","shell.execute_reply":"2025-06-03T01:24:14.739155Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# model = ort.InferenceSession(f\"/kaggle/input/bird-clef-2024-models/{CONFIG['exp_name']}/{CONFIG['exp_name']}/checkpoints/model_simpl.onnx\")","metadata":{"execution":{"iopub.status.busy":"2025-06-03T01:24:14.748696Z","iopub.execute_input":"2025-06-03T01:24:14.749689Z","iopub.status.idle":"2025-06-03T01:24:14.753856Z","shell.execute_reply.started":"2025-06-03T01:24:14.749657Z","shell.execute_reply":"2025-06-03T01:24:14.752973Z"},"papermill":{"duration":0.014662,"end_time":"2025-04-16T20:13:10.926054","exception":false,"start_time":"2025-04-16T20:13:10.911392","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Inference Class\n","metadata":{"papermill":{"duration":0.006732,"end_time":"2025-04-16T20:13:10.9402","exception":false,"start_time":"2025-04-16T20:13:10.933468","status":"completed"},"tags":[]}},{"cell_type":"code","source":"inference_class = BirdsInference(\n    device=\"cpu\",\n    verbose_tqdm=True,\n    use_sigmoid=CONFIG[\"use_sigmoid\"],\n    use_compiled_fp16=CONFIG[\"use_compiled_fp16\"],\n)","metadata":{"execution":{"iopub.status.busy":"2025-06-03T01:24:14.757491Z","iopub.execute_input":"2025-06-03T01:24:14.757731Z","iopub.status.idle":"2025-06-03T01:24:14.775962Z","shell.execute_reply.started":"2025-06-03T01:24:14.757713Z","shell.execute_reply":"2025-06-03T01:24:14.774852Z"},"papermill":{"duration":0.014366,"end_time":"2025-04-16T20:13:10.961674","exception":false,"start_time":"2025-04-16T20:13:10.947308","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Prediction","metadata":{"papermill":{"duration":0.007229,"end_time":"2025-04-16T20:13:10.976049","exception":false,"start_time":"2025-04-16T20:13:10.96882","status":"completed"},"tags":[]}},{"cell_type":"code","source":"test_preds, test_dfidx, test_end = inference_class.predict_test_loader(\n#     nn_models=model,\n    nn_models=compiled_models,\n    data_loader=loader_test,\n#     is_onnx_model=True\n    is_openvino_model=True\n)\ntest_pred_df = compose_submission_dataframe(\n    probs=test_preds,\n    dfidxs=test_dfidx,\n    end_seconds=test_end,\n    filenames=loader_test.dataset.df[loader_test.dataset.name_col].copy(),\n    bird2id=bird2id,\n)\nsample_submission = pd.read_csv(\"/kaggle/input/birdclef-2025/sample_submission.csv\")\nif test_pred_df.shape[1] > sample_submission.shape[1]:\n    print(\"Shrinking columns\")\n    test_pred_df = test_pred_df[sample_submission.columns]","metadata":{"execution":{"iopub.status.busy":"2025-06-03T01:24:14.777011Z","iopub.execute_input":"2025-06-03T01:24:14.777788Z","iopub.status.idle":"2025-06-03T01:30:04.033872Z","shell.execute_reply.started":"2025-06-03T01:24:14.77776Z","shell.execute_reply":"2025-06-03T01:30:04.032435Z"},"papermill":{"duration":33.815759,"end_time":"2025-04-16T20:13:44.798833","exception":false,"start_time":"2025-04-16T20:13:10.983074","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"if not IS_TEST:\n    display(test_pred_df)","metadata":{"execution":{"iopub.status.busy":"2025-06-03T01:30:04.03458Z","iopub.status.idle":"2025-06-03T01:30:04.034857Z","shell.execute_reply.started":"2025-06-03T01:30:04.034732Z","shell.execute_reply":"2025-06-03T01:30:04.034744Z"},"papermill":{"duration":0.036498,"end_time":"2025-04-16T20:13:44.844604","exception":false,"start_time":"2025-04-16T20:13:44.808106","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# -- Fernando start","metadata":{}},{"cell_type":"code","source":"import time\nDEADLINE = time.time() + 90*60 - 2*60\nSKIP = False\n\nimport os\nimport site\nif not os.path.exists(f'{site.getsitepackages()[0]}/onnxruntime'):\n    !pip install -q onnxruntime --no-index --find-links /kaggle/input/b5-onnxruntime-v1-22","metadata":{"trusted":true,"_kg_hide-input":true,"execution":{"iopub.status.busy":"2025-06-03T01:30:04.036147Z","iopub.status.idle":"2025-06-03T01:30:04.036511Z","shell.execute_reply.started":"2025-06-03T01:30:04.03636Z","shell.execute_reply":"2025-06-03T01:30:04.036379Z"},"jupyter":{"source_hidden":true}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"if 'SKIP' not in globals() or not SKIP:   # torch\n    print('libs torch v7.1')\n\n    # 7.1: CE as loss\n    # 7.0: dataset that loads pseudo labels in memory (for colab)\n    # 6.3: time and freq masking\n    # 6.2: multi lr scheduler (in progress)\n    # 6.1: Volodymyr's pseudo label approach in train\n    # 6.0: train with unlabeled using soft pseudo labels\n    # 5.6: add background noise as unlabeled; remove mask\n    # 5.5: support bits in prediction\n    # 5.4: remove unused code\n    # 5.3: different lrs for encoder and head\n    # 5.2: correct support for use_mask=False\n    # 5.1: remove filter_species (for split common/uncommon models)\n    # 5.0: train species separately based on filter + drop_path_rate\n    # 4.2: cleanup of unused and add bits=16\n    # 4.1: pick_audio_segment returns 7s interval instead of offset\n    # 4.0: h5 dataset with custom filter\n    # 3.6: aws changes in train\n    # 3.5: transform_to_spec with parameter on number of bits for audio\n    # 3.4: changed get_checkpoint_path to pick last epoch in case of tie\n    # 3.3: fixed bug in get_checkpoint_path\n    # 3.2: Separate config\n    # 3.1: fix load_state\n    # 3.0: initial\n\n    # ----------------------------------------------------------------------------------------------------\n    # 0.0 imports\n    # ----------------------------------------------------------------------------------------------------\n\n    import warnings\n    warnings.filterwarnings('ignore')\n\n    import pyarrow.parquet as pq\n\n    import copy\n    from contextlib import contextmanager, nullcontext\n    \n    import os\n    import gc\n    import time\n    import pickle\n    import psutil\n    import ast\n    import copy\n    import math\n    import random\n    import pdb\n    import json\n    import glob\n    import shutil\n\n    # from tqdm import tqdm\n    from tqdm.autonotebook import tqdm\n    import wandb\n    from collections import deque, defaultdict\n    import numba as nb\n\n    import numpy as np\n    import numpy.matlib\n    import pandas as pd\n    # import datatable as dt\n    # import polars\n\n    pd.set_option('display.max_columns', 300)\n    pd.set_option('display.max_rows', 100)\n    pd.set_option(\"display.width\", 150)\n    np.set_printoptions(linewidth=140)\n\n    from multiprocessing.pool import ThreadPool\n\n    # torch\n    import torch\n    # from torch import cuda\n    from torch import nn\n    import torch.nn.functional as F\n    from torch.utils.data import Dataset, IterableDataset, DataLoader\n    # from torchinfo import summary\n    import pytorch_lightning as L\n    from pytorch_lightning.utilities.model_summary import ModelSummary\n    from pytorch_lightning.callbacks import (LearningRateMonitor, ModelCheckpoint, TQDMProgressBar, StochasticWeightAveraging)\n    from pytorch_lightning.loggers import WandbLogger\n    # from pytorch_lightning.tuner.tuning import Tuner\n    import torchaudio\n    import torchaudio.transforms as T\n    # from torchvision import transforms\n    import torchvision\n    from torchvision.transforms import v2\n    from torch.distributions import Beta\n    import cv2\n    from PIL import Image, ImageOps\n\n    try:\n        from lion_pytorch import Lion\n    except:\n        pass\n\n    # from kaggle_secrets import UserSecretsClient  # no longer supported in kaggle\n    import os\n    #os.environ[\"TOKENIZERS_PARALLELISM\"] = \"false\"\n    # os.environ['CUDA_VISIBLE_DEVICES'] = '1'\n\n    # visual models\n    import timm\n    from timm.models.layers import drop_path\n\n    import albumentations as A\n\n    # metrics\n    import sklearn.metrics\n\n    # multiprocessing\n    # import joblib\n    from joblib import Parallel, delayed\n    from joblib.externals.loky.backend.context import get_context\n    import concurrent.futures\n\n    # onnx\n    try:\n        import onnx\n        import onnxruntime as ort\n    except:\n        pass\n\n    try:\n        import openvino.runtime as ovrt\n    except:\n        pass\n\n    # ----------------------------------------------------------------------------------------------------\n    # 0.0 constants\n    # ----------------------------------------------------------------------------------------------------\n\n    STRICT = False\n\n    SR = 32_000\n    \n    BIRDS = [\n        '1139490', '1192948', '1194042', '126247', '1346504', '134933', '135045', '1462711', '1462737', '1564122', '21038', '21116',\n        '21211', '22333', '22973', '22976', '24272', '24292', '24322', '41663', '41778', '41970', '42007', '42087', '42113', '46010',\n        '47067', '476537', '476538', '48124', '50186', '517119', '523060', '528041', '52884', '548639', '555086', '555142', '566513',\n        '64862', '65336', '65344', '65349', '65373', '65419', '65448', '65547', '65962', '66016', '66531', '66578', '66893', '67082',\n        '67252', '714022', '715170', '787625', '81930', '868458', '963335', \n        'amakin1', 'amekes', 'ampkin1', 'anhing', 'babwar', 'bafibi1', 'banana', 'baymac', 'bbwduc', 'bicwre1', 'bkcdon', 'bkmtou1', \n        'blbgra1', 'blbwre1', 'blcant4', 'blchaw1', 'blcjay1', 'blctit1', 'blhpar1', 'blkvul', 'bobfly1', 'bobher1', 'brtpar1', 'bubcur1',\n        'bubwre1', 'bucmot3', 'bugtan', 'butsal1', 'cargra1', 'cattyr', 'chbant1', 'chfmac1', 'cinbec1', 'cocher1', 'cocwoo1', 'colara1',\n        'colcha1', 'compau', 'compot1', 'cotfly1', 'crbtan1', 'crcwoo1', 'crebob1', 'cregua1', 'creoro1', 'eardov1', 'fotfly', 'gohman1',\n        'grasal4', 'grbhaw1', 'greani1', 'greegr', 'greibi1', 'grekis', 'grepot1', 'gretin1', 'grnkin', 'grysee1', 'gybmar', 'gycwor1', \n        'labter1', 'laufal1', 'leagre', 'linwoo1', 'littin1', 'mastit1', 'neocor', 'norscr1', 'olipic1', 'orcpar', 'palhor2', 'paltan1',\n        'pavpig2', 'piepuf1', 'pirfly1', 'piwtyr1', 'plbwoo1', 'plctan1', 'plukit1', 'purgal2', 'ragmac1', 'rebbla1', 'recwoo1', 'rinkin1',\n        'roahaw', 'rosspo1', 'royfly1', 'rtlhum', 'rubsee1', 'rufmot1', 'rugdov', 'rumfly1', 'ruther1', 'rutjac1', 'rutpuf1', 'saffin',\n        'sahpar1', 'savhaw1', 'secfly1', 'shghum1', 'shtfly1', 'smbani', 'snoegr', 'sobtyr1', 'socfly1', 'solsan', 'soulap1', 'spbwoo1',\n        'speowl1', 'spepar1', 'srwswa1', 'stbwoo2', 'strcuc1', 'strfly1', 'strher', 'strowl1', 'tbsfin1', 'thbeup1', 'thlsch3', 'trokin',\n        'tropar', 'trsowl', 'turvul', 'verfly', 'watjac1', 'wbwwre1', 'whbant1', 'whbman1', 'whfant1', 'whmtyr1', 'whtdov', 'whttro1',\n        'whwswa1', 'woosto', 'y00678', 'yebela1', 'yebfly1', 'yebsee1', 'yecspi2', 'yectyr1', 'yehbla2', 'yehcar1', 'yelori1', 'yeofly1',\n        'yercac1', 'ywcpar', \n    ]\n\n    N_BIRDS = len(BIRDS)\n    RANGE_BIRDS = range(N_BIRDS)\n    \n    label_to_num = dict(zip(BIRDS, RANGE_BIRDS))\n\n    RARE_BIRDS = [ # 39 birds\n        '1139490', '1192948', '1194042', '126247', '1346504', '134933', '1462711', '1462737', '1564122',\n        '21038', '21116', '24272', '24292', '41778', '42087', '42113', '46010', '47067', '476537',\n        '476538', '523060', '528041', '548639', '555142', '64862', '65336', '65419', '65547', '66016',\n        '66531', '66578', '66893', '67082', '714022', '787625', '81930', '868458', '963335',\n    ]\n    \n    DIFFICULT_BIRDS = [\n        \n    ]\n\n    # ----------------------------------------------------------------------------------------------------\n    # 0.1 utils\n    # ----------------------------------------------------------------------------------------------------\n\n    class dotdict(dict):\n        def __getattr__(self, name):\n            return self.get(name, None)\n\n        def __setattr__(self, name, val):\n            self[name] = val\n\n        def __set__(self, name, val):\n            self.__setattr__(name, val)\n            \n        def to_dict(self):\n            return eval(str(self))\n\n        # the following are required by save_obj\n        def __getstate__(self):\n            return self.__dict__\n\n        def __setstate__(self, d):\n            self.__dict__.update(d)\n\n\n    class Log:\n        def __init__(self, log_path, time_key=True):\n            self.path = log_path\n            if time_key:\n                self.path = self.path.replace('.','{}.'.format(time.strftime('_%Y%m%d%H%M%S', time.localtime(time.time()))))\n            print(time.strftime('%Y-%m-%d %H:%M:%S',time.localtime(time.time())), file=open(self.path,'a+'))\n            print('log path:', self.path)\n            print('------------ begin ------------', file=open(self.path, 'a+'))\n\n        def __call__(self, *content):\n            t1 = time.strftime('%H:%M:%S', time.localtime(time.time()))\n            print(*content)\n            print(t1, *content, file=open(self.path,'a+'))\n\n        def clean(self):\n            print(time.strftime('%Y-%m-%d %H:%M:%S',time.localtime(time.time())), file=open(self.path,'w'))\n            print('------------ begin -------------', file=open(self.path,'a+'))\n\n\n    def is_interactive():\n        return os.environ['KAGGLE_KERNEL_RUN_TYPE'] == 'Interactive'\n\n\n    class Config:\n        def __getattr__(self, name):\n            \"\"\" retun None if attribute doesn't exist \"\"\"\n            return None\n\n        def __repr__(self):\n            out = dict()\n            names = dir(self)\n            for key in names:\n                if not key.startswith('__') and key != 'to_dict':\n                    out[key] = getattr(self, key)\n            return str(out)\n\n        def to_dict(self):\n            config = dotdict(eval(str(self)))\n            if config.model:\n                config.model = dotdict(config.model)\n            if config.aug:\n                config.aug = dotdict(config.aug)\n            return config\n\n\n    def save_config():\n        name = f'config {config.name}.json'\n        folder = f'./{config.name}' if os.path.exists(config.name) else '.'\n            \n        with open(f'{folder}/{name}', 'w') as f:\n            json.dump(config.to_dict(), f)\n\n\n    def load_config(name, folder=None):\n        name = f'{folder}/config {name}.json' if folder is not None else f'config {name}.json'\n        with open(name, 'r') as f:\n            config = dotdict(json.load(f))\n        return config\n\n\n    class Timer:\n        def __init__(self, margin=1.1):\n            self.margin = margin\n\n        def start(self, duration):\n            self.laps = 1\n            self.start_time = time.time()\n            self.duration = duration\n\n        def lap(self):\n            elapsed = time.time() - self.start_time\n            mean_lap = elapsed / self.laps\n            if elapsed + self.margin * mean_lap > self.duration:\n                return False\n\n            self.laps += 1\n            return True\n\n\n    def fillna_df(df, values=0):\n        if not hasattr(values, '__len__'):\n            values = [values] * df.shape[-1]\n        numeric_columns = [c for c in df.columns if df[c].dtype not in ['object', 'bool']]\n        for i, col in enumerate(numeric_columns):\n            a = df[col].values\n            if np.isnan(a.sum()):\n                df[col] = np.where(np.isnan(a), values[i], a)\n\n\n    def wait():\n        for i in range(10000):\n            print('.', end='')\n            time.sleep(60)\n\n\n    def save_obj(obj, name, protocol=4): # pickle.HIGHEST_PROTOCOL):\n        with open('./'+ name + '.pkl', 'wb') as f:\n            pickle.dump(obj, f, protocol)\n\n\n    def load_obj(name, folder=''):\n        name = name.replace('.pkl', '')\n        with open(folder + name + '.pkl', 'rb') as f:\n            return pickle.load(f)\n\n\n    @contextmanager\n    def timer(name=''):\n        s = time.time()\n        mem_before = psutil.virtual_memory().used/1024/1024/1024\n        yield\n        elapsed = time.time() - s\n        mem_after = psutil.virtual_memory().used/1024/1024/1024\n        print(f'[{name}] {elapsed: .3f} s   {mem_before: .2f} GB -> {mem_after: .2f} GB')\n\n\n    def seed_everything(seed_value=42):\n        os.environ['PYTHONHASHSEED'] = str(seed_value)\n        np.random.seed(seed_value)\n        random.seed(seed_value)\n        try:\n            tf.random.set_seed(seed_value)\n        except:\n            pass\n\n        try:\n            torch.manual_seed(seed_value)\n            torch.cuda.manual_seed(seed_value)\n            torch.cuda.manual_seed_all(seed_value)\n            torch.backends.cudnn.deterministic = True\n            torch.backends.cudnn.benchmark = False\n        except:\n            pass\n\n\n    def get_gpu_ram():\n        \"\"\" nvidia GPU memory utilization in MB \"\"\"\n        nvmlInit()\n        handle = nvmlDeviceGetHandleByIndex(0)\n        info = nvmlDeviceGetMemoryInfo(handle)\n        return info.used/1024/1024\n\n\n    def reduce_mem_usage(df, do_categoricals=False, do_float=False):\n        \"\"\" iterate through all the columns of a dataframe and modify the data type\n            to reduce memory usage.\n        \"\"\"\n        start_mem = df.memory_usage().sum() / 1024**2\n        print('Memory usage of dataframe is {:.2f} MB'.format(start_mem))\n\n        for col in df.columns:\n            col_type = df[col].dtype\n\n            if col_type != object:\n                c_min = df[col].min()\n                c_max = df[col].max()\n                if str(col_type)[:3] == 'int':\n                    if c_min >= np.iinfo(np.int8).min and c_max <= np.iinfo(np.int8).max:\n                        df[col] = df[col].astype(np.int8)\n                    elif c_min >= np.iinfo(np.int16).min and c_max <= np.iinfo(np.int16).max:\n                        df[col] = df[col].astype(np.int16)\n                    elif c_min >= np.iinfo(np.int32).min and c_max <= np.iinfo(np.int32).max:\n                        df[col] = df[col].astype(np.int32)\n                    elif c_min >= np.iinfo(np.int64).min and c_max <= np.iinfo(np.int64).max:\n                        df[col] = df[col].astype(np.int64)\n                    #print(col, c_min, c_max, df[col].dtype)\n                elif do_float:\n                    if c_min > np.finfo(np.float16).min and c_max < np.finfo(np.float16).max:\n                        df[col] = df[col].astype(np.float16)\n                    elif c_min > np.finfo(np.float32).min and c_max < np.finfo(np.float32).max:\n                        df[col] = df[col].astype(np.float32)\n                    else:\n                        df[col] = df[col].astype(np.float64)\n            else:\n                if do_categoricals==True:\n                    df[col] = df[col].astype('category')\n\n        end_mem = df.memory_usage().sum() / 1024**2\n        print('Memory usage after optimization is: {:.2f} MB'.format(end_mem))\n        print('Decreased by {:.1f}%'.format(100 * (start_mem - end_mem) / start_mem))\n        return df\n\n\n    def memory_cleanup():\n        \"\"\"\n        Cleans up GPU memory. Call after a fold is trained.\n        https://github.com/huggingface/transformers/issues/1742\n        \"\"\"\n        for obj in gc.get_objects():\n            if torch.is_tensor(obj):\n                del obj\n        gc.collect()\n        torch.cuda.empty_cache()\n\n\n    def clean_config(models_config, models=None):\n        out, out_dict = [], dict()\n        for row in models_config:\n            if len(row) == 5:\n                name, folds, file, folder, config = row\n                folder = f'{ROOT}/{folder}'\n            elif len(row) == 4:\n                name, folds, folder, config = row\n                file = name\n                folder = f'{ROOT}/{folder}'\n            else:\n                name, folds, config = row\n                file = name\n                folder = f'{ROOT}/b5-models'\n                \n            if models is not None and not name in models:\n                continue\n            config = dotdict(config)\n            if config.aug:\n                config.aug = dotdict(config.aug)\n            out.append((name, folds, file, folder, config))\n            out_dict[name] = config\n        return out, out_dict\n        \n    \n    def print_num_parameters(model, note=None):\n        n = sum(p.numel() for p in model.parameters())\n        if note:\n            print(f'{note:<30}: {n:,.0f}')\n        return n\n    \n    \n    def print_model_size(pl_model):\n        \"\"\" print size of lightning model. Adjust example_input_array based on model input \"\"\"\n        size = ModelSummary(pl_model, max_depth=0).trainable_parameters\n        model_name = str(pl_model.model.__class__).split('.')[-1][:-2]\n        print(f'model {model_name}: {size:,.0f}')\n        return size\n\n\n    def print_weights(model):\n        \"\"\" prints model weights and returns them as a dataframe \"\"\"\n        items = []\n        for i, (name, module) in enumerate(model.named_modules()):\n            weight = module.weight.shape if hasattr(module, 'weight') else ''\n            items.append((name, f'{module.__class__.__module__}.{module.__class__.__name__}', weight))\n\n        res = pd.DataFrame(items, columns=['name', 'class', 'weight'])\n        res.to_csv('x3d_l.csv')\n        pd.set_option('display.max_rows', None)\n        display(res)\n        return res\n\n\n    def are_weights_equal(state_dict1, state_dict2):\n        if len(state_dict1) != len(state_dict2):\n            return False\n        for (key1, value1), (key2, value2) in zip(state_dict1.items(), state_dict2.items()):\n            if key1 != key2 or not torch.equal(value1, value2):\n                return False\n        return True\n\n\n    def compare_onnx_models(path1, path2):\n        model1 = onnx.load(path1)\n        model2 = onnx.load(path2)\n\n        # Compare serialized model content\n        if onnx.helper.printable_graph(model1.graph) != onnx.helper.printable_graph(model2.graph):\n            print(\"Graphs are different.\")\n            return False\n\n        # Compare initializer tensors (weights)\n        tensors1 = {init.name: onnx.numpy_helper.to_array(init) for init in model1.graph.initializer}\n        tensors2 = {init.name: onnx.numpy_helper.to_array(init) for init in model2.graph.initializer}\n\n        if tensors1.keys() != tensors2.keys():\n            print(\"Different parameter names.\")\n            return False\n\n        for k in tensors1:\n            if not np.array_equal(tensors1[k], tensors2[k]):\n                print(f\"Mismatch in weights for: {k}\")\n                return False\n\n        print(\"Models are structurally and numerically identical.\")\n        return True\n\n    # ----------------------------------------------------------------------------------------------------\n    # 0.2 Layers\n    # ----------------------------------------------------------------------------------------------------\n\n    def glorot_uniform(parameter):\n        nn.init.xavier_uniform_(parameter.data, gain=1.0)\n            \n\n    class Sinusoidal_positional_embedding(nn.Module):\n        \"\"\" performs fourier positional embedding converting from (B, T) to (B, T, E),\n            batch, time (sequence), embedding.\n            Used by P2 in iceCube to embed continuous variables (categorical used \n            regular embedding. Unlike the positional encoding in the attention is all\n            you need paper, which provides a different embedding for each embedding position \n            and position in the sequence, this function only provides a different embedding\n            for each embedding position (all positions in the sequence are embded in the same way).\n        \"\"\"\n        \n        def __init__(self, e_dim=16, M=10_000):\n            super().__init__()\n            self.e_dim = e_dim  # embedding e_dimension\n            self.M = M\n\n        def forward(self, x):\n            device = x.device\n            half_e_dim = self.e_dim // 2\n            emb = math.log(self.M) / half_e_dim\n            emb = torch.exp(torch.arange(half_e_dim, device=device) * (-emb))\n            emb = x[..., None] * emb[None, ...]\n            emb = torch.cat((emb.sin(), emb.cos()), e_dim=-1)\n            return emb\n\n\n    class DropPath(nn.Module):\n        \"\"\" drop all features for random instances of dim[-2] (regular drop drops random features)\n        \"\"\"\n    \n        def __init__(self, drop_prob=None):\n            super(DropPath, self).__init__()\n            self.drop_prob = drop_prob\n\n        def forward(self, x):\n            return drop_path(x, self.drop_prob, self.training)\n\n        def extra_repr(self) -> str:\n            return \"p={}\".format(self.drop_prob)\n            \n\n    class GaussianNoiseRandom(object):\n        \"\"\" add random noise with random std.\n            Noise is randomized across last two dimensions (h,w)\n        \"\"\"\n        def __init__(self, mean=0., std_range=(0.05, 0.2)):\n            self.mean = mean\n            self.std_range = std_range\n\n        def __call__(self, tensor):\n            std = random.uniform(*self.std_range)  # Randomly choose a std deviation each time\n            # Generate noise only for one channel\n            single_channel_noise = torch.randn(tensor[0, :, :].size()) * std + self.mean\n            # Expand noise to all channels\n            noise = single_channel_noise.repeat(tensor.size(0), 1, 1)\n            return tensor + noise\n\n        def __repr__(self):\n            return f'{self.__class__.__name__}(mean={self.mean}, std_range={self.std_range})'\n\n\n    class GaussianNoiseLearned(nn.Module):\n        \"\"\" add random noise with random std with learnable std. \n            Noise is randomized across all dimensions.\n\n        Args:\n            sigma (float, optional): relative standard deviation used to generate the\n                noise. Relative means that it will be multiplied by the magnitude of\n                the value your are adding the noise to. This means that sigma can be\n                the same regardless of the scale of the vector.\n            is_relative_detach (bool, optional): whether to detach the variable before\n                computing the scale of the noise. If `False` then the scale of the noise\n                won't be seen as a constant but something to optimize: this will bias the\n                network to generate vectors with smaller values.\n        \"\"\"\n\n        def __init__(self, sigma=0.1, is_relative_detach=True):\n            super().__init__()\n            self.sigma = sigma\n            self.is_relative_detach = is_relative_detach\n            self.noise = torch.tensor(0).to(config.device)\n\n        def forward(self, x):\n            if self.training and self.sigma != 0:\n                scale = self.sigma * x.detach() if self.is_relative_detach else self.sigma * x\n                sampled_noise = self.noise.repeat(*x.size()).normal_() * scale\n                x = x + sampled_noise\n            return x\n            \n\n    class Model_head(nn.Module):\n        def __init__(self, input_dim, output_dim=1, dropout_head=None, criterion=None):\n            super(Model_head, self).__init__()\n            self.criterion = criterion\n            self.dropout_head = dropout_head\n            if criterion is not None:\n                if dropout_head is None:\n                    self.dropout1 = nn.Dropout(0.1)\n                    self.dropout2 = nn.Dropout(0.2)\n                    self.dropout3 = nn.Dropout(0.3)\n                    self.dropout4 = nn.Dropout(0.4)\n                    self.dropout5 = nn.Dropout(0.5)\n                else:\n                    self.dropout = nn.Dropout(dropout_head)\n                \n            self.regressor = nn.Linear(input_dim, output_dim)\n            glorot_uniform(self.regressor.weight)\n            #self._init_weights(self.regressor)\n            \n        def _init_weights(self, module):\n            if isinstance(module, nn.Linear):\n                module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)\n                if module.bias is not None:\n                    module.bias.data.zero_()\n            elif isinstance(module, nn.Embedding):\n                module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)\n                if module.padding_idx is not None:\n                    module.weight.data[module.padding_idx].zero_()\n            elif isinstance(module, nn.LayerNorm):\n                module.bias.data.zero_()\n                module.weight.data.fill_(1.0)\n            \n        def forward(self, x, target=None):\n            # x is B x S x C\n            if self.criterion is not None and target is not None:\n                if self.dropout_head is None:\n                    logits1 = self.regressor(self.dropout1(x))\n                    logits2 = self.regressor(self.dropout2(x))\n                    logits3 = self.regressor(self.dropout3(x))\n                    logits4 = self.regressor(self.dropout4(x))\n                    logits5 = self.regressor(self.dropout5(x))\n                    logits = ((logits1 + logits2 + logits3 + logits4 + logits5) / 5)\n\n                    loss1 = self.criterion(logits1, target)\n                    loss2 = self.criterion(logits2, target)\n                    loss3 = self.criterion(logits3, target)\n                    loss4 = self.criterion(logits4, target)\n                    loss5 = self.criterion(logits5, target)\n                    loss = (loss1 + loss2 + loss3  + loss4 + loss5) / 5\n                else:\n                    logits = self.regressor(self.dropout(x))\n                    loss = self.criterion(logits, target)\n\n                #print(loss.shape, attention_mask.shape)\n                loss = loss.mean()\n            else:\n                logits = self.regressor(x)\n                loss = 0\n\n            logits = logits.squeeze()\n            return loss, logits\n\n\n    class GeM(nn.Module):\n        \"\"\" average_pool_2D with power parameter \"\"\"\n        def __init__(self, p=1.5, eps=1e-6):\n            super().__init__()\n            self.p = nn.Parameter(torch.ones(1) * p)\n            self.eps = eps\n\n        def forward(self, x):\n            # return F.avg_pool2d(x.clamp(min=self.eps).pow(self.p), (x.size(-2), x.size(-1))).pow(1.0 / self.p)\n            x = x.clamp(min=self.eps).pow(self.p)\n            return x.mean([2,3]).pow(1.0 / self.p)\n\n        def __repr__(self):\n            return (self.__class__.__name__\n                + f\"(p={self.p.data.tolist()[0]:.4f}, eps={self.eps})\")\n\n    \n    def init_layer(layer):\n        nn.init.xavier_uniform_(layer.weight)\n        if hasattr(layer, \"bias\"):\n            if layer.bias is not None:\n                layer.bias.data.fill_(0.0)\n\n\n    def init_bn(bn):\n        bn.bias.data.fill_(0.0)\n        bn.weight.data.fill_(1.0)\n\n\n    class AttBlockV2(nn.Module):\n        def __init__(self, in_features, out_features=N_BIRDS, activation=\"sigmoid\"):\n            super().__init__()\n            self.activation = activation\n            self.att = nn.Conv1d(\n                in_channels=in_features, out_channels=out_features,\n                kernel_size=1, stride=1, padding=0, bias=True)\n            self.cla = nn.Conv1d(in_channels=in_features, out_channels=out_features,\n                kernel_size=1, stride=1, padding=0, bias=True)\n            self._init_weights()\n\n        def _init_weights(self):\n            init_layer(self.att)\n            init_layer(self.cla)\n\n        def forward(self, x):\n            # x: (n_samples, n_in, n_time)\n            norm_att = torch.softmax(torch.tanh(self.att(x)), dim=-1)\n            cla = self.nonlinear_transform(self.cla(x))\n            x = torch.sum(norm_att * cla, dim=2)\n            return x, norm_att, cla\n\n        def nonlinear_transform(self, x):\n            if self.activation == \"linear\":\n                return x\n            elif self.activation == \"sigmoid\":\n                return torch.sigmoid(x)\n\n\n    class CustomMasking(nn.Module):\n        def __init__(self, mask_max_length: int, mask_max_masks: int, p=1.0, inplace=True):\n            super().__init__()\n            assert isinstance(mask_max_masks, int) and mask_max_masks > 0\n            self.mask_max_masks = mask_max_masks\n            self.mask_max_length = mask_max_length\n            self.mask_module = None\n            self.p = p\n            self.inplace = inplace\n\n        def forward(self, x):\n            if not self.inplace:\n                output = x.clone()\n            for i in range(x.shape[0]):\n                if np.random.binomial(n=1, p=self.p):\n                    n_applies = np.random.randint(\n                        low=1, high=self.mask_max_masks + 1\n                    )\n                    for _ in range(n_applies):\n                        if self.inplace:\n                            x[i : i + 1] = self.mask_module(x[i : i + 1])\n                        else:\n                            output[i : i + 1] = self.mask_module(output[i : i + 1])\n            if self.inplace:\n                return x\n            else:\n                return output\n\n\n    class CustomTimeMasking(CustomMasking):\n        def __init__(self, mask_max_length: int, mask_max_masks: int, p=1.0, inplace=True):\n            super().__init__(mask_max_length=mask_max_length, mask_max_masks=mask_max_masks,\n                             p=p, inplace=inplace)\n            self.mask_module = T.TimeMasking(time_mask_param=mask_max_length)\n\n\n    class CustomFreqMasking(CustomMasking):\n        def __init__(self, mask_max_length: int, mask_max_masks: int, p=1.0, inplace=True):\n            super().__init__(mask_max_length=mask_max_length, mask_max_masks=mask_max_masks,\n                             p=p, inplace=inplace,)\n            self.mask_module = T.FrequencyMasking(freq_mask_param=mask_max_length)\n\n    # ----------------------------------------------------------------------------------------------------\n    # 0.3 Models\n    # ----------------------------------------------------------------------------------------------------\n\n    ENCODER_SIZE = { # dim_channels, channels first\n        'efficientnet_b0': (1280, True),\n        'efficientnet_b1': (1280, True),\n        'efficientnet_b2': (1408, True),\n        'efficientnet_b3': (1536, True),\n        'efficientnet_b4': (1792, True),\n        'efficientnet_b5': (2048, True),\n        'efficientnet_b6': (2304, True),\n        'efficientnet_b7': (2560, True),\n        \n        'swinv2_tiny_window16_256': (768, False),           # (3, 256, 256) -> (8, 8, 768)    h,w = 256; bs=8, lr=1e-4\n        'maxvit_base_tf_512.in21k_ft_in1k': (768, True),    # (3, 512, 512) -> (768, 16, 16)  h,w = 512; bs=1 crashes\n        'tiny_vit_21m_512': (576, True),                    # (3, 512, 512) -> (576, 16, 16)  h,w = multiple of 16; bs=2, lr=1e-3\n        'tiny_vit_21m_224': (576, True),                    # (3, 224, 224) -> (576, 14, 14)  h,w = multiple of 16; bs=8, lr=1e-3\n        'efficientvit_b2': (None, None),\n    }\n\n    class SpecNetImg(nn.Module):\n        \"\"\" linear head \"\"\"\n        def __init__(self, cfg, num_classes=N_BIRDS, channels_first=True, resize=None, add_position=False,\n                     pretrained=True, submit=False, in_chans=3, **kwargs):\n            super().__init__()\n            self.submit = submit\n            self.in_chans = in_chans\n            encoder = cfg.encoder or 'efficientnet_b0'\n            head_dropout = cfg.head_dropout or 0.\n            if add_position:\n                rows, cols = cfg.img_dim\n                row_pos = torch.arange(rows).unsqueeze(1).expand(rows, cols)\n                col_pos = torch.arange(cols).unsqueeze(0).expand(rows, cols)\n                self.pos = torch.stack((row_pos, col_pos), dim=0).unsqueeze(0)\n            else:\n                self.pos = None\n            \n            self.resize = v2.Resize(resize) if resize else None\n            \n            # spec augment\n            spec_augment = []\n            if cfg.freq_mask is not None:\n                spec_augment.append(CustomFreqMasking(**cfg.freq_mask))\n            if cfg.time_mask is not None:\n                spec_augment.append(CustomTimeMasking(**cfg.time_mask))\n            self.spec_augment = nn.Sequential(*spec_augment) if len(spec_augment) > 0 else None\n            \n            drop_path_rate = cfg.drop_path_rate\n            if isinstance(pretrained, str):\n                self.model = timm.create_model(\n                    encoder, pretrained=False, in_chans=in_chans, drop_path_rate=drop_path_rate)\n                load_encoder_state(self.model, pretrained)\n            else:\n                self.model = timm.create_model(\n                    encoder, pretrained=pretrained, in_chans=in_chans, drop_path_rate=drop_path_rate)\n            n_channels = self.model.num_features\n            #if cfg.transformer:\n            #    self.pool = lambda x: x.mean([1])\n            #else:\n            self.pool = nn.AdaptiveAvgPool2d(1) if channels_first else lambda x: x.mean([1,2])\n            self.drop = nn.Dropout(p=head_dropout)\n            self.fc = nn.Linear(n_channels, out_features=num_classes, bias=True)\n\n        def forward(self, x):\n            B = x.size()[0]\n            if self.resize is not None:\n                x = self.resize(x)\n            \n            if self.spec_augment is not None:\n                x = self.spec_augment(x)\n                \n            if self.pos is not None:\n                pos = self.pos.expand(B,-1,-1,-1).to(x.device)\n                x = torch.concat((x, pos), dim=1)\n            elif self.in_chans != x.shape[1]:\n                x = x.expand(-1, 3, -1, -1)     # C3HW\n            x = self.model.forward_features(x)  # BHWC\n            x = self.pool(x)                    # BHWC -> BC or BC11\n            x = x.view(B, -1)  # flatten\n            x = self.drop(x)\n            x = self.fc(x)\n            return x\n           \n\n    class SpecNetImgGem(nn.Module):\n        \"\"\" GeM head \"\"\"\n        def __init__(self, cfg, num_classes=N_BIRDS, channels_first=True, resize=None,\n                     pretrained=True, submit=False, in_chans=3, **kwargs):\n            super().__init__()\n            self.submit = submit\n            self.in_chans = in_chans\n            self.out_indices = cfg.out_indices or 2\n            encoder = cfg.encoder or 'efficientnet_b0'\n            p = cfg.gem_p or 2.\n            head_dropout = cfg.head_dropout or 0.\n            self.resize = v2.Resize(resize) if resize else None\n            \n            drop_path_rate = cfg.drop_path_rate\n            if isinstance(pretrained, str):\n                self.model = timm.create_model(\n                    encoder, pretrained=False, in_chans=in_chans, features_only=True, drop_path_rate=drop_path_rate)\n                load_encoder_state(self.model, pretrained)\n            else:\n                self.model = timm.create_model(\n                    encoder, pretrained=pretrained, in_chans=in_chans, features_only=True, drop_path_rate=drop_path_rate)\n            self.pools = torch.nn.ModuleList([GeM(p=p) for i in range(self.out_indices)])\n            mid_features = self.model.feature_info.channels()[-self.out_indices:]\n            self.n_mid_features = sum(mid_features)\n            self.drop = nn.Dropout(p=head_dropout) if head_dropout > 0 else None\n            self.bn = torch.nn.BatchNorm1d(self.n_mid_features)\n            self.fc = nn.Linear(self.n_mid_features, out_features=num_classes, bias=True)\n\n        def forward(self, x):\n            B = x.size()[0]\n            if self.in_chans != x.shape[1]:\n                x = x.expand(-1, 3, -1, -1)     # C3HW\n            x = self.model(x)[-self.out_indices:]   # n*BHWC\n            if self.out_indices > 1:\n                h = torch.cat([pool(m) for m, pool in zip(x, self.pools)], dim=1)  # n*BHWC -> BC or BC11\n                x = h.view(B, -1)  # flatten\n            else:\n                x = self.pools[0](x[0])\n            if self.drop is not None:\n                x = self.drop(x)\n            x = self.bn(x)\n            x = self.fc(x)\n            return x\n\n\n    class SpecNetImgGemMultiDrop(nn.Module):\n        \"\"\" GeM head with multidrop \"\"\"\n        def __init__(self, backbone=\"resnet34\", num_class=397, in_chans=1):\n            super().__init__()\n            head_dropout = config.head_dropout or 0.\n            drop_path_rate = config.drop_path_rate\n            encoder = timm.create_model(\n                backbone, features_only=False, pretrained=False, in_chans=in_chans, drop_path_rate=drop_path_rate)\n\n            layers = list(encoder.children())[:-2]\n            self.model = nn.Sequential(*layers)\n            if \"efficientnet\" in backbone:\n                dense_input = encoder.num_features\n            elif \"swin\" in backbone:\n                dense_input = encoder.num_features\n            elif hasattr(encoder, \"fc\"):\n                dense_input = encoder.fc.in_features\n            else:\n                dense_input = encoder.feature_info[-1][\"num_chs\"]\n\n            self.global_pool = GeM()\n            self.dropouts = nn.ModuleList([nn.Dropout(p) for p in \\\n                [.2*head_dropout,.4*head_dropout,.6*head_dropout,.8*head_dropout,1*head_dropout]])\n            self.head = nn.Linear(dense_input, num_class)\n\n        def forward(self, batch):\n            x = batch['x']\n            x = self.model(x)\n            x = self.global_pool(x)\n            logit = sum([self.head(dropout(x)) for dropout in self.dropouts]) / 5\n            return logit\n\n\n    class SpecNetImgSed(nn.Module):  # NOT FINISHED\n        \"\"\" attention head. \"\"\"\n        \n        def __init__(self, cfg, num_classes=N_BIRDS, channels_first=True, resize=None, add_position=False,\n                     pretrained=True, submit=False, in_chans=3, **kwargs):\n            super().__init__()\n            self.submit = submit\n            self.in_chans = in_chans\n            self.out_indices = cfg.out_indices or 2\n            encoder = cfg.encoder or 'efficientnet_b0'\n            head_dropout = cfg.head_dropout or 0.\n            drop_path_rate = cfg.drop_path_rate or .2\n            if isinstance(pretrained, str):\n                self.model = timm.create_model(\n                    encoder, pretrained=False, in_chans=in_chans, features_only=True, drop_path_rate=drop_path_rate)\n                load_encoder_state(self.model, pretrained)\n            else:\n                self.model = timm.create_model(\n                    encoder, pretrained=pretrained, in_chans=in_chans, drop_path_rate=drop_path_rate)\n            self.pools = torch.nn.ModuleList([nn.AdaptiveAvgPool2d(1) if channels_first else lambda x: x.mean([1,2])\n                    for i in range(self.out_indices)])\n            mid_features = self.model.feature_info.channels()[-self.out_indices:]\n            self.n_mid_features = sum(mid_features)\n            self.drop = nn.Dropout(p=head_dropout) if head_dropout > 0 else None\n            self.fc1 = nn.Linear(self.n_mid_features, self.n_mid_features, bias=True)\n            self.att_block = AttBlockV2(self.n_mid_features, num_classes, activation='sigmoid')\n            self.bn = torch.nn.BatchNorm1d(self.n_mid_features)\n            self.fc = nn.Linear(self.n_mid_features, out_features=num_classes, bias=True)\n\n        def forward(self, x):\n            B = x.size()[0]\n            x = x.expand(-1, 3, -1, -1)         # C3HW\n            x = self.model(x)[-self.out_indices:]   # n*BHWC\n            h = torch.cat([pool(m) for m, pool in zip(x, self.pools)], dim=1)  # n*BHWC -> BC or BC11\n            x = h.view(B, -1)  # flatten\n            if self.drop is not None:\n                x = self.drop(x)\n            x = self.bn(x)\n            x = self.fc(x)\n            return x\n\n\n    class RawSignalNet(nn.Module):\n        def __init__(self, config, num_classes=N_BIRDS, resize=None, channels_first=True, pretrained=True):\n            super(RawSignalNet, self).__init__()\n            encoder = config.encoder or 'efficientnet_b0'\n            head_dropout = config.head_dropout or 0.\n            self.img_dim = config.img_dim\n            self.resize = v2.Resize(resize) if resize is not None else None\n            self.model = timm.create_model(encoder, pretrained=pretrained, in_chans=3)\n            n_channels = self.model.num_features\n            self.pool = nn.AdaptiveAvgPool2d(1) if channels_first else lambda x: x.mean([1,2])\n            self.drop = nn.Dropout(p=head_dropout)\n            self.fc = nn.Linear(n_channels, out_features=num_classes, bias=True)\n\n        def forward(self, batch):\n            # x = (bs, 1, 32_000*img_duration)\n            x = batch['x']\n            bs = x.size(0)\n            x = x.view(bs, 1, self.img_dim[0], -1).transpose(3,2)\n            x = x.expand(-1, 3, -1, -1)         # C3HW\n            bs = x.size(0)\n            if self.resize:\n                x = self.resize(x)\n            x = self.model.forward_features(x)\n            x = self.pool(x)\n            x = x.view(bs, -1)  # flatten\n            x = self.drop(x)\n            x = self.fc(x)\n            return x\n\n\n    class SpecNetRaw(nn.Module):\n        \"\"\" models that reads a 2D image directly from dataset \"\"\"\n        def __init__(self, config, num_classes=N_BIRDS, channels_first=True, resize=None, add_position=False,\n                     pretrained=True):\n            super().__init__()\n            encoder = config.encoder or 'efficientnet_b0'\n            head_dropout = config.head_dropout or 0.\n            if add_position:\n                rows, cols = config.img_dim\n                row_pos = torch.arange(rows).unsqueeze(1).expand(rows, cols)\n                col_pos = torch.arange(cols).unsqueeze(0).expand(rows, cols)\n                self.pos = torch.stack((row_pos, col_pos), dim=0).unsqueeze(0)\n            else:\n                self.pos = None\n            \n            self.resize = v2.Resize(resize) if resize else None\n            self.model = timm.create_model(encoder, pretrained=pretrained, in_chans=3)\n            n_channels = self.model.num_features\n            self.pool = nn.AdaptiveAvgPool2d(1) if channels_first else lambda x: x.mean([1,2])\n            self.drop = nn.Dropout(p=head_dropout)\n            self.fc = nn.Linear(n_channels, out_features=num_classes, bias=True)\n\n        def forward(self, batch):\n            # do preprocess\n            x = batch['x']\n            B = x.size()[0]\n            # x = torch.unsqueeze(x, dim=1)  # image to dataset adds channel dimension\n            if self.resize:\n                x = self.resize(x)\n            if self.pos is not None:\n                pos = self.pos.expand(B,-1,-1,-1).to(x.device)\n                x = torch.concat((x, pos), dim=1)\n            else:\n                x = x.expand(-1, 3, -1, -1)     # C3HW\n            x = self.model.forward_features(x)  # BHWC\n            x = self.pool(x)                    # BHWC -> BC or BC11\n            x = x.view(B, -1)  # flatten\n            x = self.drop(x)\n            x = self.fc(x)\n            return x\n\n\n    class CWT(nn.Module):\n        # - wavelet_width=7, fs=200, lower_freq=0.5, upper_freq=40, n_scales=40, border_crop=1, stride=16\n\n        def __init__(self, wavelet_width=7, fs=200, lower_freq=.5, upper_freq=40., n_scales=40,\n                     size_factor=1.0, border_crop=1, stride=16):\n            super().__init__()\n\n            self.initial_wavelet_width = wavelet_width\n            self.fs = fs\n            self.lower_freq = lower_freq\n            self.upper_freq = upper_freq\n            self.size_factor = size_factor\n            self.n_scales = n_scales\n            self.wavelet_width = wavelet_width\n            self.border_crop = border_crop\n            self.stride = stride\n            wavelet_bank_real, wavelet_bank_imag = self._build_wavelet_kernel()\n            self.wavelet_bank_real = nn.Parameter(wavelet_bank_real, requires_grad=False)\n            self.wavelet_bank_imag = nn.Parameter(wavelet_bank_imag, requires_grad=False)\n\n            self.kernel_size = self.wavelet_bank_real.size(3)\n\n        def _build_wavelet_kernel(self):\n            s_0 = 1 / self.upper_freq\n            s_n = 1 / self.lower_freq\n\n            base = np.power(s_n / s_0, 1 / (self.n_scales - 1))\n            scales = s_0 * np.power(base, np.arange(self.n_scales))\n\n            frequencies = 1 / scales\n            truncation_size = scales.max() * np.sqrt(4.5 * self.initial_wavelet_width) * self.fs\n            one_side = int(self.size_factor * truncation_size)\n            kernel_size = 2 * one_side + 1\n\n            k_array = np.arange(kernel_size, dtype=np.float32) - one_side\n            t_array = k_array / self.fs\n\n            wavelet_bank_real = []\n            wavelet_bank_imag = []\n\n            for scale in scales:\n                norm_constant = np.sqrt(np.pi * self.wavelet_width) * scale * self.fs / 2.0\n                scaled_t = t_array / scale\n                exp_term = np.exp(-(scaled_t ** 2) / self.wavelet_width)\n                kernel_base = exp_term / norm_constant\n                kernel_real = kernel_base * np.cos(2 * np.pi * scaled_t)\n                kernel_imag = kernel_base * np.sin(2 * np.pi * scaled_t)\n                wavelet_bank_real.append(kernel_real)\n                wavelet_bank_imag.append(kernel_imag)\n\n            wavelet_bank_real = np.stack(wavelet_bank_real, axis=0)\n            wavelet_bank_imag = np.stack(wavelet_bank_imag, axis=0)\n\n            wavelet_bank_real = torch.from_numpy(wavelet_bank_real).unsqueeze(1).unsqueeze(2)\n            wavelet_bank_imag = torch.from_numpy(wavelet_bank_imag).unsqueeze(1).unsqueeze(2)\n            return wavelet_bank_real, wavelet_bank_imag\n\n        def forward(self, x):\n            border_crop = self.border_crop // self.stride\n            start = border_crop\n            end = (-border_crop) if border_crop > 0 else None\n\n            # x [n_batch, n_channels, time_len]\n            out_reals = []\n            out_imags = []\n\n            n_signals = x.size(1)\n            in_width = x.size(2)\n            out_width = int(np.ceil(in_width / self.stride))\n            pad_along_width = np.max((out_width - 1) * self.stride + self.kernel_size - in_width, 0)\n            padding = pad_along_width // 2 + 1\n\n            for i in range(n_signals):\n                # [n_batch, 1, 1, time_len]\n                x_ = x[:, i, :].unsqueeze(1).unsqueeze(2)\n                out_real = F.conv2d(x_, self.wavelet_bank_real, stride=(1, self.stride), padding=(0, padding))\n                out_imag = F.conv2d(x_, self.wavelet_bank_imag, stride=(1, self.stride), padding=(0, padding))\n                out_real = out_real.transpose(2, 1)\n                out_imag = out_imag.transpose(2, 1)\n                out_reals.append(out_real)\n                out_imags.append(out_imag)\n\n            out_real = torch.cat(out_reals, axis=1)\n            out_imag = torch.cat(out_imags, axis=1)\n\n            out_real = out_real[:, :, :, start:end]\n            out_imag = out_imag[:, :, :, start:end]\n\n            scalograms = torch.sqrt(out_real ** 2 + out_imag ** 2)\n            return scalograms\n\n\n    class SpecNetCwt2D(nn.Module):\n        def __init__(self, config, num_classes=N_BIRDS, resize=None, channels_first=True, pretrained=True):\n            super().__init__()\n            encoder = config.encoder or 'maxvit_base_tf_512.in21k_ft_in1k'\n            head_dropout = config.head_dropout or 0.\n            n_scales = config.img_dim[0]\n            stride = math.ceil(32_000 * config.img_duration / config.img_dim[1])\n            self.preprocess = CWT(\n                wavelet_width=7, fs=32_000, lower_freq=50., upper_freq=16_000., n_scales=n_scales,\n                size_factor=1.0, border_crop=1, stride=stride)\n            self.resize = v2.Resize(resize) if resize is not None else None\n            self.model = timm.create_model(encoder, pretrained=pretrained, in_chans=3)\n                                         # drop_rate=FLAGS.drop_rate, drop_path_rate=FLAGS.drop_path_rate\n            n_channels = self.model.num_features\n            self.pool = nn.AdaptiveAvgPool2d(1) if channels_first else lambda x: x.mean([1,2])\n            self.drop = nn.Dropout(p=head_dropout)\n            self.fc = nn.Linear(n_channels, out_features=num_classes, bias=True)\n\n        def forward(self, batch):\n            # do preprocess\n            x = batch['x']\n            x = self.preprocess(x)\n            n, c, f, t = x.size()   # N1HW\n            if self.resize:\n                x = self.resize(x)\n            x = x.expand(-1, 3, -1, -1)         # C3HW\n            x = self.model.forward_features(x)  # BHWC\n            x = self.pool(x)                    # BHWC -> BC or BC11\n            x = x.view(n, -1)  # flatten\n            x = self.drop(x)\n            x = self.fc(x)\n            return x\n\n\n    # ----------------------------------------------------------------------------------------------------\n    # 0.4 augment\n    # ----------------------------------------------------------------------------------------------------\n    \n    def augment_img_train():\n        \"\"\" train augmentation to run on dataset \"\"\"\n        aug = config.aug\n        pipeline = []\n        if aug.Flip > 0:\n            pipeline.append(A.HorizontalFlip(p=aug.Flip))\n        \n        if aug.CoarseDropout is not None:  # (max_height=.375, max_width=.375, max_holes=1, p=.7)\n            max_height = int(config.img_dim[0] * aug.CoarseDropout[0]) # 0.375\n            max_width = int(config.img_dim[1] * aug.CoarseDropout[1]) # 0.375\n            pipeline.append(A.CoarseDropout(max_height=max_height, max_width=max_width,\n                            max_holes=aug.CoarseDropout[2], p=aug.CoarseDropout[3]))\n\n        tranform_fn = A.Compose(pipeline)\n        return tranform_fn\n\n    try:\n        import audiomentations as AA\n\n        def augment_audio_train():\n            \"\"\" train augmentation on audio, which add background noise \"\"\"\n            tranform_fn = AA.Compose([\n                AA.OneOf([\n                    AA.Gain(min_gain_db=-15, max_gain_db=15, p=1.0),\n                    AA.GainTransition(\n                        min_gain_db=-24.0,\n                        max_gain_db=6.0,\n                        min_duration=0.2,\n                        max_duration=6.0,\n                        p=1.0\n                    )\n                ], p=0.5,),\n                AA.OneOf([\n                    AA.AddGaussianNoise(p=1.0),\n                    AA.AddGaussianSNR(p=1.0),\n                ], p=0.3,),\n                AA.OneOf([\n                    AA.AddShortNoises(\n                        sounds_path=f\"{ROOT}/birdclef2021-background-noise/ff1010bird_nocall/nocall\",\n                        min_snr_db=0,\n                        max_snr_db=3,\n                        p=1.0,\n                        lru_cache_size=10,\n                        min_time_between_sounds=4.0,\n                        max_time_between_sounds=16.0,\n                    ),\n                    #AA.AddShortNoises(\n                    #    sounds_path=\"/kaggle/input/birdclef2021-background-noise/esc50/use_label\",\n                    #    min_snr_db=0,\n                    #    max_snr_db=3,\n                    #    p=1.0,\n                    #    lru_cache_size=10,\n                    #    min_time_between_sounds=4.0,\n                    #    max_time_between_sounds=16.0,\n                    #),\n                ], p=0.5,),\n                AA.OneOf([\n                    AA.AddBackgroundNoise(\n                        sounds_path=f\"{ROOT}/birdclef2021-background-noise/train_soundscapes/nocall\",\n                        min_snr_db=0,\n                        max_snr_db=3,\n                        p=1.0,\n                        lru_cache_size=3,),\n                    AA.AddBackgroundNoise(\n                        sounds_path=f\"{ROOT}/birdclef2021-background-noise/aicrowd2020_noise_30sec/noise_30sec\",\n                        min_snr_db=0,\n                        max_snr_db=3,\n                        p=1.0,\n                        lru_cache_size=20,\n                        #lru_cache_size=450,\n                       ),\n                ], p=0.5,),\n                AA.LowPassFilter(p=0.5),\n            ])\n            return tranform_fn\n        \n    except:\n        def augment_audio_train():\n            return lambda x: x\n\n    BETA = Beta(1., 1.)\n    \n    def mixup(batch):\n        \"\"\" mixup augmentation to run during training step \"\"\"\n        x, y = batch['x'], batch['y']\n        perm = torch.randperm(x.size(0))\n\n        # lam = torch.FloatTensor([np.random.beta(alpha, alpha)])\n        lam = BETA.rsample(x.shape[:1]).to(x.device)\n        lam_x = lam.view(-1, 1, 1, 1)\n        lam_y = lam.view(-1, 1)\n        x = x * lam_x + x[perm] * (1 - lam_x)\n        y = y * lam_y + y[perm] * (1 - lam_y)\n\n        batch['x'], batch['y'] = x, y\n\n    # ----------------------------------------------------------------------------------------------------\n    # 0.5 Datasets\n    # ----------------------------------------------------------------------------------------------------\n\n    class SnipetAmplitudeToDB(torch.nn.Module):\n        \"\"\"Turn a tensor from the power/amplitude scale to the decibel scale.\n\n        .. devices:: CPU CUDA\n\n        .. properties:: Autograd TorchScript\n\n        This output operates per image, hence it returns the same values for an audio clip \n        split into snippets vs. a full clip.\n\n        Args:\n            stype (str, optional): scale of input tensor (``\"power\"`` or ``\"magnitude\"``). The\n                power being the elementwise square of the magnitude. (Default: ``\"power\"``)\n            top_db (float or None, optional): minimum negative cut-off in decibels.  A reasonable\n                number is 80. (Default: ``None``)\n\n        Example\n            >>> waveform, sample_rate = torchaudio.load(\"test.wav\", normalize=True)\n            >>> spec = melspec_128_256_5(waveform)\n            >>> transform = StableAmplitudeToDB(stype=\"power\", top_db=80)\n            >>> spec_db = transform(spec)\n        \"\"\"\n        __constants__ = [\"multiplier\", \"amin\", \"ref_value\", \"db_multiplier\"]\n\n        def __init__(self, stype='power', top_db=None):\n            super().__init__()\n            self.stype = stype\n            if top_db is not None and top_db < 0:\n                raise ValueError(\"top_db must be positive value\")\n            self.top_db = top_db\n            self.multiplier = 10.0 if stype == \"power\" else 20.0\n            self.amin = 1e-10\n            self.ref_value = 1.0\n            self.db_multiplier = math.log10(max(self.amin, self.ref_value))\n\n        def forward(self, x):\n            \"\"\"Numerically stable implementation from Librosa.\n\n            https://librosa.org/doc/latest/generated/librosa.amplitude_to_db.html\n\n            Args:\n                x (Tensor): Input tensor before being converted to decibel scale.\n\n            Returns:\n                Tensor: Output tensor in decibel scale.\n            \"\"\"\n            x_db = self.multiplier * torch.log10(torch.clamp(x, min=self.amin))\n            x_db -= self.multiplier * self.db_multiplier\n            if self.top_db is not None:\n                x_db = torch.max(x_db, (x_db.amax(dim=(-2, -1), keepdims=True) - self.top_db))\n            return x_db\n\n\n    hop_length_512_5 = 32_000 * 5 // (512 - 1)\n    hop_length_256_5 = 32_000 * 5 // (256 - 1)\n    hop_length_128_5 = 32_000 * 5 // (128 - 1)\n    hop_length_160_5 = 32_000 * 5 // (160 - 1)\n    hop_length_192_5 = 32_000 * 5 // (192 - 1)\n    hop_length_224_10 = 32_000 * 10 // (224 - 1)\n\n    melspec_128_512_5 = T.MelSpectrogram(\n        n_fft=2048, hop_length=hop_length_512_5, f_min=50, f_max=16_000, sample_rate=32_000,\n        n_mels=128, norm='slaney', mel_scale='slaney', pad_mode='constant')\n    melspec_128_256_5 = T.MelSpectrogram(\n        n_fft=2048, hop_length=hop_length_256_5, f_min=50, f_max=16_000, sample_rate=32_000,\n        n_mels=128, norm='slaney', mel_scale='slaney', pad_mode='constant')\n    melspec_128_128_5 = T.MelSpectrogram(\n        n_fft=2048, hop_length=hop_length_128_5, f_min=50, f_max=16_000, sample_rate=32_000,\n        n_mels=128, norm='slaney', mel_scale='slaney', pad_mode='constant')\n    melspec_160_160_5 = T.MelSpectrogram(\n        n_fft=2048, hop_length=hop_length_160_5, f_min=50, f_max=16_000, sample_rate=32_000,\n        n_mels=160, norm='slaney', mel_scale='slaney', pad_mode='constant')\n    melspec_192_192_5 = T.MelSpectrogram(\n        n_fft=2048, hop_length=hop_length_192_5, f_min=50, f_max=16_000, sample_rate=32_000,\n        n_mels=192, norm='slaney', mel_scale='slaney', pad_mode='constant')\n    melspec_192_256_5 = T.MelSpectrogram(\n        n_fft=2048, hop_length=hop_length_256_5, f_min=50, f_max=16_000, sample_rate=32_000,\n        n_mels=192, norm='slaney', mel_scale='slaney', pad_mode='constant')\n\n    melspec_224_224b_10 = T.MelSpectrogram(\n        n_fft=2048, hop_length=hop_length_224_10, f_min=400, f_max=10_000, sample_rate=32_000,\n        n_mels=224, norm='slaney', mel_scale='slaney', pad_mode='constant')\n    melspec_p4_1 = T.MelSpectrogram(\n        n_fft=2048, hop_length=512, f_min=20, f_max=16_000, sample_rate=32_000,\n        n_mels=128, normalized=True)\n    melspec_p4_1b = T.MelSpectrogram(\n        n_fft=2048, hop_length=hop_length_256_5, f_min=20, f_max=16_000, sample_rate=32_000,\n        n_mels=128, normalized=True)\n    melspec_p4_4 = T.MelSpectrogram(\n        n_fft=1024, hop_length=hop_length_512_5, f_min=50, f_max=14_000, sample_rate=32_000,\n        n_mels=64, normalized=True)\n    melspec_p4_4b = T.MelSpectrogram(\n        n_fft=1024, hop_length=hop_length_512_5, f_min=50, f_max=14_000, sample_rate=32_000,\n        n_mels=64, norm='slaney', mel_scale='slaney', pad_mode='constant')\n    melspec_256_128_5 = T.MelSpectrogram(\n        n_fft=2048, hop_length=hop_length_128_5, f_min=50, f_max=16_000, sample_rate=32_000,\n        n_mels=256, norm='slaney', mel_scale='slaney', pad_mode='constant')\n    melspec_256_256_5 = T.MelSpectrogram(\n        n_fft=2048, hop_length=hop_length_256_5, f_min=50, f_max=16_000, sample_rate=32_000,\n        n_mels=256, norm='slaney', mel_scale='slaney', pad_mode='constant')\n    \n    db_transform = T.AmplitudeToDB(stype='power', top_db=80)\n    db_transform_snipet = SnipetAmplitudeToDB(stype='power', top_db=80)\n        \n    len_audio_5s = 5 * 32_000\n    len_audio_10s = 10 * 32_000\n    len_audio_1m = 1 * 60 * 32_000\n    \n    # Spectrogram types\n    SPEC_128_512_5 = 1\n    SPEC_224_224B_10 = 2\n\n\n    SPEC_DIMS = {\n        'p4_1': (128, 313),\n        'p4_4': (64, 512),\n        'p4_4b': (64, 512),\n        'f180': (180, 180),\n        'f128': (128, 128),\n        'f256': (256, 128),\n        'f256x2': (256, 256),\n        'f192': (192, 256),\n        'f192x2': (192, 192),\n    }\n        \n\n    def transform_to_spec(dim=(128, 256), norm=False, duration=5, snipet=True, spec=None, bits=8):\n        \"\"\" converts an audio array to spectrograms. Audio is a 2D array with a single segment to check \n            against train or multiple segments for submit prediction. Transformation is performed per\n            audio snipet, which means that submitting an audio by itself (1, 32_000*5) or as part of a\n            batch (6, 32_000*5), produces the same result for the common audio snipet.\n        \"\"\"\n        if spec == 'p4_1':\n            melspec_fn = melspec_p4_1\n        elif spec == 'p4_1b':\n            melspec_fn = melspec_p4_1b\n        elif spec == 'p4_4':\n            melspec_fn = melspec_p4_4\n        elif spec == 'p4_4b':\n            melspec_fn = melspec_p4_4b\n        elif spec == 'f256':\n            melspec_fn = melspec_256_128_5\n\n        elif spec == 'f128x2':\n            melspec_fn = melspec_128_128_5\n        elif spec == 'f160x2':\n            melspec_fn = melspec_160_160_5\n        elif spec == 'f192x2':\n            melspec_fn = melspec_192_192_5\n        elif spec == 'f192':\n            melspec_fn = melspec_192_256_5\n\n        elif spec == 'f256x2':\n            melspec_fn = melspec_256_256_5\n        elif dim[-1] == 512:\n            if duration == 5:\n                melspec_fn = melspec_128_512_5\n        elif dim[-1] == 256:\n            if duration == 5:\n                melspec_fn = melspec_128_256_5\n        elif dim[-1] == 224:\n            if duration == 10:\n                melspec_fn = melspec_224_224b_10\n                \n        db_transform_ = db_transform_snipet if snipet else db_transform\n        eps = 1e-6\n        \n        def _process(audio):\n            spec = melspec_fn(audio)\n            spec = db_transform(spec)\n            min_ = torch.amin(spec)\n            max_ = torch.amax(spec)\n            spec = (spec - min_) / (max_ - min_)\n            if bits == 8 or bits is None:\n                spec = (spec * 255).to(torch.uint8) / 255\n            elif bits == 16:\n                spec = (spec * 65535).to(torch.uint32) / 65535\n            return spec\n\n        def _process_norm(audio):\n            spec = melspec_fn(audio)\n            spec = db_transform(spec)\n            mean_, std_ = spec.mean(), spec.std()\n            spec = (spec - mean_) / (std_ + eps)\n            min_ = torch.amin(spec)\n            max_ = torch.amax(spec)\n            spec = (spec - min_) / (max_ - min_)\n            if bits == 8 or bits is None:\n                spec = (spec * 255).to(torch.uint8) / 255\n            elif bits == 16:\n                spec = (spec * 65535).to(torch.uint32) / 65535\n            return spec\n\n        def _process_snipet(audio):\n            spec = melspec_fn(audio)\n            spec = db_transform(spec)\n            min_ = torch.amin(spec, dim=(-2,-1), keepdims=True)\n            max_ = torch.amax(spec, dim=(-2,-1), keepdims=True)\n            spec = (spec - min_) / (max_ - min_)\n            if bits == 8 or bits is None:\n                spec = (spec * 255).to(torch.uint8) / 255\n            elif bits == 16:\n                spec = (spec * 65535).to(torch.uint32) / 65535\n            return spec\n\n        def _process_norm_snipet(audio):\n            spec = melspec_fn(audio)\n            spec = db_transform(spec)\n            mean_, std_ = spec.mean((-2,-1), keepdims=True), spec.std((-2,-1), keepdims=True)\n            spec = (spec - mean_) / (std_ + eps)\n            min_ = torch.amin(spec, dim=(-2,-1), keepdims=True)\n            max_ = torch.amax(spec, dim=(-2,-1), keepdims=True)\n            spec = (spec - min_) / (max_ - min_)\n            if bits == 8 or bits is None:\n                spec = (spec * 255).to(torch.uint8) / 255\n            elif bits == 16:\n                spec = (spec * 65535).to(torch.uint32) / 65535\n            return spec\n\n        if snipet:\n            return _process_norm_snipet if norm else _process_snipet\n        else:\n            return _process_norm if norm else _process\n\n\n    def prepare_data(files, num_workers, verbose=1, norm=False, dim=(128,256), duration=5,\n                     snipet=False, spec=None, bits=8):\n        \"\"\" generate specs for submission in parallel \"\"\"\n        transform_to_spec_ = transform_to_spec(\n            dim=dim, norm=norm, duration=duration, snipet=snipet, spec=spec, bits=bits)\n        len_duration = duration * 32_000\n        \n        def _process_chunk(chunk):\n            out = []\n            for file in tqdm(chunk):\n                audio, sr = torchaudio.load(file)\n                audio = audio[:, :len_audio_1m].view(-1, len_duration)\n                out.append(transform_to_spec_(audio))\n            out = torch.cat(out, axis=0)\n            out = out.unsqueeze(1)\n            return out\n\n        t1 = time.time()\n        chunk_size = math.ceil(len(files) / num_workers)\n        chunks = [files[i : i+chunk_size] for i in range(0, len(files), chunk_size)]\n\n        with concurrent.futures.ThreadPoolExecutor(max_workers=num_workers) as executor:\n            out = list(executor.map(_process_chunk, chunks))\n        out = torch.cat(out, axis=0)\n        duration = time.time() - t1\n        if verbose:\n            print(f'specs: {out.shape}')\n            print(f'time: {duration:,.0f} s for {len(files):,.0f} files  -  all: {duration * 700 / len(files) / 60:,.0f} m')\n        return out\n\n\n    def prepare_data_ragged(files, num_workers, verbose=1, norm=False, dim=(128,256),\n                            duration=5, max_duration=32_000*4*60, snipet=False, spec=None, bits=8):\n        \"\"\" generate specs for prediction in parallel. Audio is limited to 4m and broken into 5s batches.\n            The last batch is padded with 0s.\n            Returns specs and size (number of specs per file, i.e., 5s slots). \n        \"\"\"\n        transform_to_spec_ = transform_to_spec(\n            dim=dim, norm=norm, duration=duration, snipet=snipet, spec=spec, bits=bits)\n        \n        def _process_chunk(chunk):\n            out, size = [], []\n            range_ = tqdm(chunk) if verbose >= 2 else chunk\n            for file in range_:\n                audio, sr = torchaudio.load(file)\n                if max_duration is not None:\n                    audio = audio[:, :max_duration]\n                rest = audio.shape[-1] % len_audio_5s\n                if rest > 0:\n                    audio = F.pad(audio, pad=[0, len_audio_5s - rest], mode='constant', value=0)\n                audio = audio[:, :len_audio_1m].view(-1, len_audio_5s)\n                out.append(transform_to_spec_(audio))\n                size.append(audio.shape[0])\n            out = torch.cat(out, axis=0)\n            out = out.unsqueeze(1)\n            return out, size\n\n        t1 = time.time()\n        chunk_size = math.ceil(len(files) / num_workers)\n        chunks = [files[i : i+chunk_size] for i in range(0, len(files), chunk_size)]\n\n        with concurrent.futures.ThreadPoolExecutor(max_workers=num_workers) as executor:\n            out_size = list(executor.map(_process_chunk, chunks))\n        out, size = zip(*out_size)\n        out = torch.cat(out, axis=0)\n        size = np.concatenate(size, axis=0)\n        duration = time.time() - t1\n        if verbose:\n            print(f'specs: {out.shape}')\n            print(f'time: {duration:,.0f} s for {len(files):,.0f} files  -  all: {duration * 700 / len(files) / 60:,.0f} m')\n        return out, size\n\n\n    class Cache_audio_files:\n        \"\"\" keeps a cache of a random `duration` segment of unlabeled audio files. \n            Refreshes segment every `max_usage` times it is used. \n        \"\"\"\n        \n        def __init__(self, files, duration=5, max_usage=None, size=1000):\n            self.files = files\n            random.shuffle(self.files)\n            self.n_loaded = 0\n            self.duration = duration * 32_000\n            self.max_usage = max_usage or 6\n            self.size = size or 1000\n            self.cache = dict()\n            self.count_used = np.zeros(size)\n            \n        def get_random(self):\n            idx = random.randint(0, self.size - 1)\n            cache = self.cache\n            if not idx in cache or self.count_used[idx] > self.max_usage:\n                file = self.files[self.n_loaded]\n                self.n_loaded += 1\n                if self.n_loaded == len(files):\n                    random.shuffle(self.files)\n                    self.n_loaded = 0\n                audio, sr = torchaudio.load(file)\n\n                self.cache[idx] = audio\n\n            self.count_used[idx] = 1\n            audio = self.cache[idx]\n            \n            # adjust size\n            gap = self.duration - audio.shape[-1]\n            if gap > 0:\n                pad_before = random.randint(0, gap - 1)\n                pad_after = gap - pad_before\n                audio = F.pad(audio, pad=(pad_before, pad_after), mode='constant', value=0)\n            elif gap < 0:\n                start = random.randint(0, -gap - 1)\n                audio = audio[:, start : start + self.duration]\n            return audio\n\n\n    class Ba_dataset_val(Dataset):\n        \"\"\" Uses prebuilt specs and labels\n        \"\"\"\n        def __init__(self, x_file, y_file, idx_birds=None, **kwargs):\n            super().__init__()\n            self.x_file = x_file\n            self.y_file = y_file\n            self.x = load_obj(x_file)\n            self.y = load_obj(y_file)\n            if idx_birds is not None:\n                self.y = self.y[:, idx_birds]\n            #self.ws = torch.ones(x.shape[0], dtype=torch.float32) / x.shape[0]\n            self._tensorize_and_pin()\n\n        def __len__(self):\n            return len(self.y)\n\n        def __getitem__(self, index):\n            x = self.x[index]\n            y = self.y[index]\n            return dict(idx=index, x=x, y=y)\n\n        def _tensorize_and_pin(self):\n            if torch.cuda.is_available():\n                self.x = torch.tensor(self.x).pin_memory()\n                self.y = torch.tensor(self.y).pin_memory()\n                # self.ws = torch.tensor(self.ws).pin_memory()\n            else:\n                self.x = torch.tensor(self.x)\n                self.y = torch.tensor(self.y)\n\n\n    class Ba_data_module_montage_h5(L.LightningDataModule):\n        def __init__(self, files, y, birds, train_idx, val_idx, dataset_val, fold, train_mode='random',\n                     ws=None, files_background_noise=None, custom_filter=None, idx_birds=None, \n                     train_size=None, pseudo_catalog=None, pseudo_audios=None):\n            super().__init__()\n            self.files = files\n            self.y = y\n            self.birds = birds\n            self.ws = ws\n            self.train_idx = train_idx\n            self.val_idx = val_idx\n            self.dataset_val = dataset_val\n            self.fold = fold\n            self.train_mode = train_mode\n            self.files_background_noise = files_background_noise\n            self.custom_filter = custom_filter\n            self.idx_birds = idx_birds\n            self.prefetch_factor = config.prefetch_factor\n            self.train_size = train_size or len(files)\n            self.pseudo_catalog = pseudo_catalog\n            self.pseudo_audios = pseudo_audios\n      \n        def prepare_data(self):\n            \"\"\" processing to be done only in one GPU \"\"\"\n            train_idx = self.train_idx\n            val_idx = self.val_idx\n            if self.pseudo_catalog:\n                self.train_dataset = Ba_dataset_montage_pseudo_h5(\n                    self.files[train_idx], self.y[train_idx], self.birds[train_idx],\n                    self.pseudo_catalog, self.pseudo_audios,\n                    shuffle=True, mode=self.train_mode, cache=False, augment=True,\n                    montage_cache_size=config.montage_cache_size, files_background_noise=self.files_background_noise,\n                    ws=self.ws, custom_filter=self.custom_filter, train_size=self.train_size)\n            else:\n                self.train_dataset = Ba_dataset_montage_h5(\n                    self.files[train_idx], self.y[train_idx], self.birds[train_idx], \n                    shuffle=True, mode=self.train_mode, cache=False, augment=True,\n                    montage_cache_size=config.montage_cache_size, files_background_noise=self.files_background_noise,\n                    ws=self.ws, custom_filter=self.custom_filter, train_size=self.train_size)\n            folder = f'{ROOT}/b5-data-{self.dataset_val}'\n            if self.val_idx is not None:\n                self.val_dataset = Ba_dataset_val(\n                    f'{folder}/x_f{self.fold}', f'{folder}/y_f{self.fold}', idx_birds=self.idx_birds)\n                size = self.val_idx.sum()\n                if size != self.val_dataset.x.shape[0]:\n                    self.val_dataset.x = self.val_dataset.x[:size]\n                    self.val_dataset.y = self.val_dataset.y[:size]\n      \n        def setup(self, stage=None):\n            \"\"\" processing to be done in all GPUs \"\"\"\n            pass\n      \n        def train_dataloader(self):\n            return DataLoader(\n                self.train_dataset, batch_size=config.train_batch_size, shuffle=False, \n                num_workers=config.num_workers, prefetch_factor=self.prefetch_factor,\n                collate_fn=None, drop_last=True)\n\n        def val_dataloader(self):\n            if self.val_idx is None:\n                return None\n            return DataLoader(\n                self.val_dataset, batch_size=config.val_batch_size, shuffle=False, \n                num_workers=config.num_workers, prefetch_factor=self.prefetch_factor,\n                collate_fn=None, drop_last=False)\n      \n        def test_dataloader(self):\n            return None\n\n\n    def pick_audio_segment(custom_filter, mode, width):\n        \"\"\" pick a segment of duration `width` based on the sections configured in `custom_filter`\n            using the following rules based on `mode`:\n            - '2x7s': either first of last segment with a random offset <= 2\n            - '7s': first segment with a random offset <=2\n            - 'random': random segment\n            - '2x7s r': 2x7s, except for segments manually curated, which use 'random'\n        \"\"\"\n        curation = custom_filter.sections[0]\n        sections = custom_filter.sections[1:]\n        if mode.endswith(' r'):\n            mode = 'random' if curation == 'm' else mode.replace(' r', '')\n        \n        if len(sections) == 1:\n            section = sections[0]\n            section_start, section_end = section\n            if mode == '2x7s':\n                if random.random() < .5:  # first 7s\n                    segment_start = section_start\n                    segment_end = min(section_start + width + 2, section_end)\n                else:   # last 7s\n                    segment_start = max(section_start, section_end - width - 2)\n                    segment_end = section_end\n            elif mode == '7s':\n                    segment_start = section_start\n                    segment_end = min(section_start + width + 2, section_end)\n            elif mode == 'random':\n                segment_start = section_start\n                segment_end = section_end\n        else:\n            if mode == '2x7s':\n                if random.random() < .5:  # first 7s\n                    section = sections[0]\n                    section_start, section_end = section\n                    segment_start = section_start\n                    segment_end = min(section_start + width + 2, section_end)\n                else:   # last 7s\n                    section = sections[-1]\n                    section_start, section_end = section\n                    segment_start = max(section_start, section_end - width - 2)\n                    segment_end = section_end\n            elif mode == '7s':\n                section = sections[0]\n                section_start, section_end = section\n                segment_start = section_start\n                segment_end = min(section_start + width + 2, section_end)\n            elif mode == 'random':\n                sections_prob = custom_filter.sections_prob\n                idx_section = random.choices(range(len(sections)), weights=sections_prob, k=1)[0]\n                section = sections[idx_section]\n                section_start, section_end = section\n                segment_start = section_start\n                segment_end = section_end\n                \n\n        return segment_start, segment_end\n\n\n    def test_pick_audio_segment():\n        custom_filter = dotdict(sections=['m', (5, 8), (10, 20), (30, 60)], sections_prob=[3, 10, 30])\n        pick_audio_segment(custom_filter, '2x7s', 5)\n\n\n    class Ba_dataset_montage_h5(Dataset):\n        \"\"\" montage of raw audio files and background noise (if `files_background_noise` is not None). \n            Montage includes additional audios if only one bird in initial audio. \n            If files_background_noise is defined, the montage includes a random unlabeled audio with probability\n            `config.background_noise_prob`. If an unlabeled audio is included, the mask includes only the birds in\n            the labeled audios. Otherwise, it includes all labels.\n            Audio files are h5 and we use a custom filter to pick the audio sections.\n            \n            Parameters:\n                - mode:\n                    - '2x7s' pick a random segment from first or last 7s.\n                    - '7s' pick a random segment from first 7s.\n                    - 'random' pick a segment at random.\n                - false_prob: int:\n                    - >0: include segments with no vocalization with probability `false_prob`\n        \"\"\"\n        def __init__(self, files, y, birds, shuffle=True, ws=None, mode='2x7s', augment=None,\n                     cache=False, montage_cache_size=10, files_background_noise=None, custom_filter=None,\n                     train_size=None, **kwargs):\n            super().__init__()\n            self.files = files\n            self.y = torch.tensor(y, dtype=torch.float32)\n            self.birds = birds\n            self.shuffle = shuffle\n            self.ws = ws\n            self.mode = mode\n            self.custom_filter = custom_filter\n            self.train_size = train_size\n            self.false_prob = config.false_prob or 0\n            \n            self.augment = augment_img_train() if augment else None\n            self.augment_audio = augment_audio_train() if (augment and config.aug is not None and config.aug.audio) else None\n            if config.background_noise_prob and files_background_noise is not None:\n                self.background_noise_cache = Cache_audio_files(\n                    files_background_noise, config.img_duration, config.background_noise_max_usage, config.background_noise_cache_size)\n                self.background_noise_prob = config.background_noise_prob\n            else:\n                self.background_noise_prob = 0\n            self.aug_volume = (math.log(config.aug.volume[0]), math.log(config.aug.volume[1])) \\\n                if config.aug.volume is not None else (math.log(.5), math.log(2.))\n            self.background_noise_reference = config.background_noise_reference\n            self.norm_audio = config.norm_audio or False\n\n            self.label_smoothing = config.label_smoothing\n            self.montage = config.montage if config.montage is not None else (0, 0)\n            self.apply_montage = y.sum(-1) == 1  # exclude from montage saudios with multiple labels\n            self.montage_cache_size = montage_cache_size\n            self.montage_cache = {bird: [] for bird in RANGE_BIRDS}\n            self.snipet = config.snipet if config.snipet is not None else False\n            self.width = config.img_duration\n            self.idxs = np.arange(len(files))\n            bits = config.bits or 8\n            self.transform_to_spec_ = transform_to_spec(\n                dim=config.img_dim, norm=config.spec_norm, duration=config.img_duration,\n                snipet=self.snipet, spec=config.spec, bits=bits)\n            self.on_epoch_end()\n            self.all_true = torch.ones(N_BIRDS, dtype=torch.float32)\n            \n            \n        def __len__(self):\n            return self.train_size\n\n\n        def __getitem__(self, index):\n            # pick primary audio\n            idx = self.idxs[index]  # idxs is already randomized before each epoch\n            x = self._get_audio(idx)\n            if self.y is not None:\n                y = self.y[idx]\n            else:\n                y = ws = None\n\n            x, y = self._get_montage(x, y, idx)\n            if random.random() < self.background_noise_prob:\n                background_noise_audio = self.background_noise_cache.get_random()\n                if self.background_noise_reference:\n                    max_x = x.abs().max()\n                    max_background_noise = background_noise_audio.abs().max()\n                    x = x / max_x * max_background_noise + background_noise_audio\n                else:\n                    x = x + background_noise_audio * math.exp(random.uniform(*self.aug_volume))\n                mask = (y > 0).to(torch.float32)\n            else:\n                mask = self.all_true\n                \n            y = torch.clamp(y, 0., 1.)\n            if self.label_smoothing:\n                tp = y.sum(-1)\n                y = y * (1 - self.label_smoothing) + self.label_smoothing * tp / y.shape[-1]\n\n            if self.augment_audio is not None:\n                dtype = x.dtype\n                x = self.augment_audio(x[0].numpy(), sample_rate=32_000)\n                x = torch.tensor(x, dtype=dtype).unsqueeze(0)\n\n            if self.norm_audio:\n                x /= torch.abs(x).max()\n\n            x = self.transform_to_spec_(x)\n\n            if self.augment:\n                dtype = x.dtype\n                x = self.augment(image=x.permute(1,2,0).numpy())['image'].astype(np.float32)\n                x = torch.tensor(x, dtype=dtype).permute(2,0,1)\n\n            return dict(idx=idx, x=x, y=y, mask=mask)\n\n\n        def on_epoch_end(self):\n            if self.shuffle and self.ws is not None:\n                n = len(self.files)\n                self.idxs = np.random.choice(np.arange(n), size=self.train_size, replace=True, p=self.ws)\n\n\n        def _get_audio(self, idx):\n            file = self.files[idx]\n            width = self.width\n\n            # pick segment start and read it from file\n            segment_start, segment_end = pick_audio_segment(self.custom_filter[idx], self.mode, width)\n            max_offset = max(segment_end - segment_start - width, 0)\n            offset = random.uniform(0, max_offset)\n            # pick montage offset to be within 2s of montage to  minimize size of audio file to load\n            if self.montage[1] > 0:\n                montage_min_offset = max(0, offset - 2)\n                montage_max_offset = min(offset + 2, max_offset)\n                offset_montage = random.uniform(montage_min_offset, montage_max_offset)\n            else:\n                offset_montage = offset\n            \n            min_offset = min(offset, offset_montage)\n            max_offset = max(offset, offset_montage)\n            audio_start = int((segment_start + min_offset) * SR)\n            audio_end = int((segment_start + width + max_offset) * SR)\n            with h5py.File(file, 'r') as f:\n                audio = f['raw'][audio_start : audio_end]\n            audio = torch.as_tensor(audio)\n            \n            # add montage\n            start_montage = int((offset_montage - min_offset) * SR)\n            if audio.shape[-1] < 5 * SR and start_montage > 0:\n                print('######', idx, file, self.custom_filter[idx])\n                print(segment_start, segment_end, audio_start/SR, audio_end/SR, audio.shape, audio.shape[-1]/SR)\n                print(audio_start, audio_end, offset, offset_montage, start_montage/SR)\n                assert 1 == 0\n\n            self._add_montage_cache(audio.unsqueeze(0), start_montage, idx)\n            start = int((offset - min_offset) * SR)\n            width_sr = width * SR\n            x = audio[start : start + width_sr]\n            gap = width_sr - x.shape[-1]\n            if gap > 0:\n                pad_before = random.randint(0, gap - 1)\n                pad_after = gap - pad_before\n                x = F.pad(x, pad=(pad_before, pad_after), mode='constant', value=0)\n            return x.unsqueeze(0)\n\n\n        def _get_montage(self, x, y, idx):\n            \"\"\" build montage. Applies a random volume change to all signals, including the original one \"\"\"\n            aug_volume = self.aug_volume\n            x *= math.exp(random.uniform(*aug_volume))\n            if self.apply_montage[idx]:\n                bird = self.birds[idx]\n                n = random.randint(*self.montage)\n                sample = random.sample(RANGE_BIRDS, n)\n                sample = [x for x in sample if x != bird]\n                # print('n', n, bird, sample)\n                if len(sample) > 0:\n                    x = x.clone()\n                    y = y.clone()\n                    for i in sample:\n                        mx, my = self._get_montage_cache(i)\n                        if mx is not None:\n                            x += mx * math.exp(random.uniform(*aug_volume))\n                            y += my\n            return x, y\n\n\n        def _add_montage_cache(self, audio, offset, idx):\n            \"\"\" add bird instance to cache \"\"\"\n            if not self.apply_montage[idx]:\n                return\n            width_sr = self.width * SR\n            x = audio[..., offset : offset + width_sr]\n            gap = width_sr - x.shape[-1]\n            if gap > 0:\n                pad_before = random.randint(0, gap - 1)\n                pad_after = gap - pad_before\n                x = F.pad(x, pad=(pad_before, pad_after), mode='constant', value=0)\n            \n            bird = self.birds[idx]\n            data = (x, self.y[idx])\n            if len(self.montage_cache[bird]) <= self.montage_cache_size:\n                self.montage_cache[bird].append(data)\n            else:\n                idx_cache = random.randint(0, self.montage_cache_size - 1)\n                self.montage_cache[bird][idx_cache] = data\n\n\n        def _get_montage_cache(self, bird):\n            \"\"\" get random instance of bird from those cached \"\"\"\n            options = self.montage_cache[bird]\n            if len(options) == 0:\n                return None, None\n            else:\n                idx_cache = random.randint(0, len(options) - 1)\n                return self.montage_cache[bird][idx_cache]\n\n\n    class Ba_dataset_montage_pseudo_h5(Dataset):\n        \"\"\" Similar to Ba_dataset_montage_h5 but samples an unlabeled audio with prob `sample_pseudo_prob`.\n            \n            Parameters:\n                - mode:\n                    - '2x7s' pick a random segment from first or last 7s.\n                    - '7s' pick a random segment from first 7s.\n                    - 'random' pick a segment at random.\n                - catalog_pseudo: Dict[bird, list(Audios)], where Audios: list(tuple(path, list(tuple(order, pseudo_label))):\n                \n        \"\"\"\n        def __init__(self, files, y, birds, pseudo_catalog, pseudo_audios,\n                     shuffle=True, ws=None, mode='2x7s', augment=None, cache=False,\n                     montage_cache_size=10, files_background_noise=None, custom_filter=None,\n                     train_size=None, **kwargs):\n            super().__init__()\n            self.files = files\n            self.y = torch.tensor(y, dtype=torch.float32)\n            self.birds = birds\n            self.shuffle = shuffle\n            self.ws = ws\n            self.mode = mode\n            self.custom_filter = custom_filter\n            self.train_size = train_size\n            self.sample_pseudo_prob = config.sample_pseudo_prob or 0.4\n            self.pseudo_catalog = pseudo_catalog\n            self.pseudo_audios = torch.tensor(pseudo_audios).share_memory_()\n            \n            self.augment = augment_img_train() if augment else None\n            self.augment_audio = augment_audio_train() if (augment and config.aug is not None and config.aug.audio) else None\n            if config.background_noise_prob and files_background_noise is not None:\n                self.background_noise_cache = Cache_audio_files(\n                    files_background_noise, config.img_duration, config.background_noise_max_usage, config.background_noise_cache_size)\n                self.background_noise_prob = config.background_noise_prob\n            else:\n                self.background_noise_prob = 0\n            self.aug_volume = (math.log(config.aug.volume[0]), math.log(config.aug.volume[1])) \\\n                if config.aug.volume is not None else (math.log(.5), math.log(2.))\n            self.background_noise_reference = config.background_noise_reference\n            self.norm_audio = config.norm_audio or False\n\n            self.label_smoothing = config.label_smoothing\n            self.montage = config.montage if config.montage is not None else (0, 0)\n            self.apply_montage = y.sum(-1) == 1\n            self.montage_cache_size = montage_cache_size\n            self.montage_cache = {bird: [] for bird in RANGE_BIRDS}\n            self.snipet = config.snipet if config.snipet is not None else False\n            self.width = config.img_duration\n            self.idxs = np.arange(len(files))\n            bits = config.bits or 8\n            self.transform_to_spec_ = transform_to_spec(\n                dim=config.img_dim, norm=config.spec_norm, duration=config.img_duration,\n                snipet=self.snipet, spec=config.spec, bits=bits)\n            self.on_epoch_end()\n            self.all_true = torch.ones(N_BIRDS, dtype=torch.float32)\n            \n            \n        def __len__(self):\n            return self.train_size\n\n\n        def __getitem__(self, index):\n            # pick primary audio\n            idx = self.idxs[index]  # idxs is already randomized before each epoch\n            x, y = self._get_audio(idx)\n            x, y = self._get_montage(x, y, idx)\n            if random.random() < self.background_noise_prob:\n                background_noise_audio = self.background_noise_cache.get_random()\n                if self.background_noise_reference:\n                    max_x = x.abs().max()\n                    max_background_noise = background_noise_audio.abs().max()\n                    x = x / max_x * max_background_noise + background_noise_audio\n                else:\n                    x = x + background_noise_audio * math.exp(random.uniform(*self.aug_volume))\n                mask = (y > 0).to(torch.float32)\n            else:\n                mask = self.all_true\n                \n            y = torch.clamp(y, 0., 1.)\n            if self.label_smoothing:\n                tp = y.sum(-1)\n                y = y * (1 - self.label_smoothing) + self.label_smoothing * tp / y.shape[-1]\n\n            if self.augment_audio is not None:\n                dtype = x.dtype\n                x = self.augment_audio(x[0].numpy(), sample_rate=32_000)\n                x = torch.tensor(x, dtype=dtype).unsqueeze(0)\n\n            if self.norm_audio:\n                x /= torch.abs(x).max()\n\n            x = self.transform_to_spec_(x)\n\n            if self.augment:\n                dtype = x.dtype\n                x = self.augment(image=x.permute(1,2,0).numpy())['image'].astype(np.float32)\n                x = torch.tensor(x, dtype=dtype).permute(2,0,1)\n\n            return dict(idx=idx, x=x, y=y, mask=mask)\n\n\n        def on_epoch_end(self):\n            if self.shuffle and self.ws is not None:\n                n = len(self.files)\n                self.idxs = np.random.choice(np.arange(n), size=self.train_size, replace=True, p=self.ws)\n\n\n        def _get_audio(self, idx):\n            width = self.width\n            bird = self.birds[idx]\n            file = self.files[idx]\n            y = self.y[idx] if self.y is not None else None\n\n            # pick segment start and read it from file\n            segment_start, segment_end = pick_audio_segment(self.custom_filter[idx], self.mode, width)\n            max_offset = max(segment_end - segment_start - width, 0)\n            offset = random.uniform(0, max_offset)\n            # pick montage offset to be within 2s of montage to minimize size of audio file to load\n            if self.montage[1] > 0:\n                montage_min_offset = max(0, offset - 2)\n                montage_max_offset = min(offset + 2, max_offset)\n                offset_montage = random.uniform(montage_min_offset, montage_max_offset)\n            else:\n                offset_montage = offset\n                \n            min_offset = min(offset, offset_montage)\n            max_offset = max(offset, offset_montage)\n            audio_start = int((segment_start + min_offset) * SR)\n            audio_end = int((segment_start + width + max_offset) * SR)\n            start_montage = int((offset_montage - min_offset) * SR)\n            \n            if self.cache_pseudo_catalog is not None and file in self.cache_pseudo_catalog:\n                idx_cache = cache_pseudo_catalog[file][audio_start // (width * SR)]\n                audio = self.cache_pseudo[idx_cache]\n            else:\n                with h5py.File(file, 'r') as f:\n                    audio = f['raw'][audio_start : audio_end]\n                audio = torch.as_tensor(audio)\n            \n            # add montage\n            if audio.shape[-1] < 5 * SR and start_montage > 0:\n                print('######', idx, file, self.custom_filter[idx])\n                print(segment_start, segment_end, audio_start/SR, audio_end/SR, audio.shape, audio.shape[-1]/SR)\n                print(audio_start, audio_end, offset, offset_montage, start_montage/SR)\n                assert 1 == 0\n\n            self._add_montage_cache(audio.unsqueeze(0), start_montage, idx)\n            start = int((offset - min_offset) * SR)\n            width_sr = width * SR\n            x = audio[start : start + width_sr]\n            gap = width_sr - x.shape[-1]\n            if gap > 0:\n                pad_before = random.randint(0, gap - 1)\n                pad_after = gap - pad_before\n                x = F.pad(x, pad=(pad_before, pad_after), mode='constant', value=0)\n            \n            return x.unsqueeze(0), y\n\n\n        def _get_montage(self, x, y, idx):\n            \"\"\" build montage. Applies a random volume change to all signals, including the original one \"\"\"\n            aug_volume = self.aug_volume\n            x *= math.exp(random.uniform(*aug_volume))\n            if self.apply_montage[idx]:\n                bird = self.birds[idx]\n                n = random.randint(*self.montage)\n                sample = random.sample(RANGE_BIRDS, n)\n                sample = [x for x in sample if x != bird]\n                # print('n', n, bird, sample)\n                if len(sample) > 0:\n                    x = x.clone()\n                    y = y.clone()\n                    for i in sample:\n                        mx, my = self._get_montage_cache(i)\n                        if mx is not None:\n                            x += mx * math.exp(random.uniform(*aug_volume))\n                            y += my\n            return x, y\n\n\n        def _add_montage_cache(self, audio, offset, idx):\n            \"\"\" add bird instance to cache \"\"\"\n            if not self.apply_montage[idx]:\n                return\n            width_sr = self.width * SR\n            x = audio[..., offset : offset + width_sr]\n            gap = width_sr - x.shape[-1]\n            if gap > 0:\n                pad_before = random.randint(0, gap - 1)\n                pad_after = gap - pad_before\n                x = F.pad(x, pad=(pad_before, pad_after), mode='constant', value=0)\n            \n            bird = self.birds[idx]\n            data = (x, self.y[idx])\n            if len(self.montage_cache[bird]) <= self.montage_cache_size:\n                self.montage_cache[bird].append(data)\n            else:\n                idx_cache = random.randint(0, self.montage_cache_size - 1)\n                self.montage_cache[bird][idx_cache] = data\n\n\n        def _get_montage_cache(self, bird):\n            \"\"\" get random instance of bird from those cached \"\"\"\n            options = self.montage_cache[bird]\n            if len(options) == 0:\n                return None, None\n            else:\n                idx_cache = random.randint(0, len(options) - 1)\n                return self.montage_cache[bird][idx_cache]\n\n\n    # ----------------------------------------------------------------------------------------------------\n    # 0.6 Loss & metrics\n    # ----------------------------------------------------------------------------------------------------\n\n    def sigmoid(x):\n        return np.where(x >= 0, 1 / (1 + np.exp(-x)), np.exp(x) / (1 + np.exp(x)))\n\n\n    def rank_based_accuracy(preds, labels):\n        \"\"\" score that indicates the average ranking of the correct label:\n            1 indicates the true value is always the top prediction, whereas 0\n            indicates it's the last. labels must be an array of single labels.\n        \"\"\"\n        B = preds.size(0)\n        probabilities, sorted_indices = preds.sort(dim=1, descending=True)\n        correct_indices = labels.view(-1, 1).expand_as(sorted_indices)\n        ranks = (sorted_indices == correct_indices).nonzero(as_tuple=True)[1]\n        scores = 1.0 / (ranks.float() + 1.0)\n        return scores.sum().item() / B\n        \n    RANK = torch.range(1, N_BIRDS).reshape(1, -1).expand(32, -1)\n    \n    def rank_based_accuracy_multi(preds, labels):\n        \"\"\" computes a score that mimics the average ranking of the true labels.\n            Supports multiple true labels. Preds can either be logits or probs.\n        \"\"\"\n        B = preds.size(0)\n        _, sorted_indices = preds.sort(dim=1, descending=True)\n        mask = torch.take_along_dim(labels, sorted_indices, dim=1)\n        rank_ = RANK.to(mask.device)\n        rank = rank_ if B == len(rank_) else rank_[:B]\n        scores = (1 / rank * mask).sum() / mask.sum()\n        return scores.item() \n    \n    \n    RANK_NP = np.arange(N_BIRDS).reshape((1, -1))\n\n    def rank_based_accuracy_multi_np(preds, labels):\n        \"\"\" computes a score that mimics the average ranking of the true labels.\n            Supports multiple true labels. Preds can either be logits or probs.\n        \"\"\"\n        B = preds.shape[0]\n        sorted_indices = np.argsort(preds, axis=1)\n        sorted_indices = sorted_indices[::-1]\n        mask = np.take_along_axis(labels, sorted_indices, axis=1)\n        print(sorted_indices)\n        rank_ = RANK_NP\n        rank = rank_ if B == 32 else rank_[:B]\n        print(rank)\n        scores = (1 / rank * mask).sum() / mask.sum()\n        return scores\n\n    \n    def auc_multi(preds, labels):\n        \"\"\" Version of macro-averaged ROC-AUC score that ignores all classes that have no true positive labels. \"\"\"\n        scored_cols = labels.sum(axis=0) > 0\n        return sklearn.metrics.roc_auc_score(\n            labels[:, scored_cols], preds[:, scored_cols], average='macro')\n\n\n    def get_loss_fn(name):\n        reduction = 'none' if (config.use_mask) else 'mean'\n        if name == 'bce':\n            return nn.BCEWithLogitsLoss(reduction=reduction)\n        if name == 'bces':\n            return BCEWithLabelSmoothingLoss(smoothing=config.label_smoothing or 0, reduction=reduction)\n        if name == 'focal':\n            return BCEFocalLoss(alpha=config.loss_alpha or .25, gamma=config.gamma or 2., reduction=reduction)\n        if name == 'focal_paper':\n            return BCEFocalLossPaper(alpha=config.loss_alpha or .25, gamma=config.gamma or 2., reduction=reduction)\n        if name == 'focal_bce':\n            return FocalLossBCE(reduction=reduction)\n        if name == 'focal_volodymyr':\n            return FocalLossBCEVolodymyr(reduction=reduction)\n        if name == 'cce':\n            return nn.CrossEntropyLoss(label_smoothing=config.label_smoothing or 0, reduction=reduction)\n\n\n    def test_get_loss_fn():\n        t = torch.rand(32, N_BIRDS)\n        p = torch.rand(32, N_BIRDS)\n        print(get_loss_fn('bce')(p, t))\n        print(get_loss_fn('bces')(p, t))\n        print(get_loss_fn('focal')(p, t))\n        print(get_loss_fn('cce')(p, t))\n\n\n    class BCEWithLabelSmoothingLoss(torch.nn.Module):\n        def __init__(self, smoothing=0.05, reduction='mean'):\n            super(BCEWithLabelSmoothingLoss, self).__init__()\n            self.smoothing = smoothing\n            self.reduction = reduction\n\n        def forward(self, pred, target):\n            target_smoothed = target * (1.0 - self.smoothing) + self.smoothing / target.shape[-1]\n            loss = F.binary_cross_entropy_with_logits(pred, target_smoothed, reduction=self.reduction)\n            return loss\n\n\n    class BCEFocalLoss(nn.Module):\n        \"\"\" BCE focalized on positives (alpha applied only to positive) \"\"\"\n        def __init__(self, alpha=0.25, gamma=2.0, reduction='mean'):\n            super().__init__()\n            self.alpha = alpha\n            self.gamma = gamma\n            self.loss_fn = nn.BCEWithLogitsLoss(reduction='none')\n            self.reduction = reduction\n\n        def forward(self, preds, targets):\n            bce_loss = self.loss_fn(preds, targets)\n            proba = torch.sigmoid(preds)\n            loss = (targets * self.alpha * (1.0 - proba) ** self.gamma * bce_loss +\n                    (1.0 - targets) * proba**self.gamma * bce_loss)\n            if self.reduction != 'none':\n                loss = loss.mean()\n            return loss\n\n\n    class BCEFocalLossPaper(nn.Module):\n        \"\"\" BCE focused on positives (as in paper) \"\"\"\n        def __init__(self, alpha=0.25, gamma=2.0, reduction='mean'):\n            super().__init__()\n            self.alpha = alpha\n            self.gamma = gamma\n            self.loss_fn = nn.BCEWithLogitsLoss(reduction='none')\n            self.reduction = reduction\n\n        def forward(self, preds, targets):\n            bce_loss = self.loss_fn(preds, targets)\n            proba = torch.sigmoid(preds)\n            #loss = (targets * self.alpha * (1.0 - proba) ** self.gamma * bce_loss +\n            #        (1.0 - targets) * (1 - self.alpha) * proba**self.gamma * bce_loss)\n            loss = (targets * self.alpha * (1.0 - proba) ** self.gamma +\n                    (1.0 - targets) * (1 - self.alpha) * proba ** self.gamma)\n            loss *= bce_loss\n            if self.reduction != 'none':\n                loss = loss.mean()\n            return loss\n\n\n    class FocalLossBCE(torch.nn.Module):\n        \"\"\" focal + BCE. other 0.66 2024 using default parameters.\n            lr_max = 1e-5\n            lr_min = 1e-7\n            weight_decay = 1e-6\n        \"\"\"\n        def __init__(self, alpha=0.25, gamma=2, reduction=\"mean\", bce_weight=1.0, focal_weight=1.0):\n            super().__init__()\n            self.alpha = alpha\n            self.gamma = gamma\n            self.reduction = reduction\n            self.bce = torch.nn.BCEWithLogitsLoss(reduction=reduction)\n            self.bce_weight = bce_weight\n            self.focal_weight = focal_weight\n\n        def forward(self, logits, targets):\n            focall_loss = torchvision.ops.focal_loss.sigmoid_focal_loss(\n                inputs=logits,\n                targets=targets,\n                alpha=self.alpha,\n                gamma=self.gamma,\n                reduction=self.reduction,\n            )\n            bce_loss = self.bce(logits, targets)\n            return self.bce_weight * bce_loss + self.focal_weight * focall_loss\n        \n\n    class FocalLossBCEVolodymyr(torch.nn.Module):\n        def __init__(\n            self,\n            alpha: float = 0.25,\n            gamma: float = 2,\n            reduction: str = 'mean',\n            bce_weight: float = 1.0,\n            focal_weight: float = 1.0,\n        ):\n            super().__init__()\n            self.alpha = alpha\n            self.gamma = gamma\n            self.reduction = reduction\n            self.bce = torch.nn.BCEWithLogitsLoss(reduction=reduction)\n            self.bce_weight = bce_weight\n            self.focal_weight = focal_weight\n\n        def forward(self, inputs, targets):\n            focall_loss = torchvision.ops.focal_loss.sigmoid_focal_loss(\n                inputs=inputs,\n                targets=targets,\n                alpha=self.alpha,\n                gamma=self.gamma,\n                reduction=self.reduction,\n            )\n            bce_loss = self.bce(inputs, targets)\n            return self.bce_weight * bce_loss + self.focal_weight * focall_loss\n            \n\n    class BCEKDLoss(nn.Module):  ## UNFINISHED\n        \"\"\" weighted evarage of cce and kl_div with teacher labels.\n            p1 2023 with default parameters. \"\"\"\n        def __init__(self, weights=[0.1, 0.9], class_weights=None):\n            super().__init__()\n\n            self.weights = weights\n            self.T = 20\n\n        def forward(self, output):\n            input_ = output[\"logit\"]\n            target = output[\"target\"].float()\n            rating = output[\"rating\"]\n            teacher_preds = output[\"teacher_preds\"]\n\n            rating = rating.unsqueeze(1).repeat(1, CFG.num_classes)\n            loss = nn.BCEWithLogitsLoss(\n                weight=rating,\n                reduction='mean',\n            )(input_, target)\n\n            KD_loss = nn.KLDivLoss()(\n                F.log_softmax(input_ / self.T, dim=1),\n                F.softmax(teacher_preds / self.T, dim=1)\n                ) * (self.weights[1] * self.T * self.T)\n\n            return self.weights[0] * loss + KD_loss\n        \n\n    class Multi_label_loss:\n        \"\"\" Unlike cross entropy loss, which only takes into account one class, this loss also considers the neighboring\n            classes by applying a weight that corresponds to normal distribution centered on the true value, which assigns \n            a larger weight to the true value, but also includes the neighboring values with a weight that decreases as they\n            get farther away from the truth.\n            For azimuth a 1Dcnn with circular padding is used. For zenith that wouldn't work because the vaues are the two \n            edges are not continous. Hence, we extend the edges with fictitious classes with a value that mirrors the real\n            neighboring classes:\n            - zenith:  [.2, .1, | 0, .1, .2, ..., 2.9, 3.0, 3.1, | 3.0, 2.9]\n                                 <-------- real classes -------->\n        \"\"\"\n        def __init__(self, num_bins, kernel_length, device):\n            self.num_bins = num_bins\n            # azimuth weights            \n            local_weights = np.float32(gaussian_kernel(l=kernel_length, sig=kernel_length/5))\n\n            kernel = torch.nn.Parameter(torch.tensor(local_weights)[None, None, :].to(device), requires_grad=False)\n            self.az_conv = torch.nn.Conv1d(1, 1, kernel_size=kernel_length, bias=False,\n                                            padding_mode='circular', padding='same', device=device)\n            self.az_conv.weight = kernel\n            \n            # zenith weights\n            self.ze_conv = torch.nn.Conv1d(1, 1, kernel_size=kernel_length,  bias=False,\n                                           padding_mode='zeros', padding='same', device=device)\n            self.ze_conv.weight = kernel\n\n        def nll(self, pred, target):\n            pred = pred.log_softmax(dim=1)\n            loss = torch.mean(torch.sum(-target * pred, dim=1))\n            return loss\n\n        def __call__(self, pred, target):\n            paz, pze = pred\n\n            with torch.no_grad():\n                az_class = target[:, 0].type(torch.int64)\n                ze_class = target[:, 1].type(torch.int64)\n                az_onehot = F.one_hot(az_class, self.num_bins).unsqueeze(1).type(torch.float32)\n                ze_onehot = F.one_hot(ze_class, self.num_bins).unsqueeze(1).type(torch.float32)\n                az_smoothed = self.az_conv(az_onehot).squeeze(1)\n                ze_smoothed = self.ze_conv(ze_onehot).squeeze(1)\n\n            loss = self.nll(paz, az_smoothed) + self.nll(pze, ze_smoothed)\n            return loss\n\n\n    def criterion_bce(logits, labels):\n        loss_fct = nn.BCEWithLogitsLoss(reduction='none')\n        truth = labels.view(-1, 1)\n        loss = loss_fct(logits.view(-1, 1), truth)\n        loss = torch.masked_select(loss, truth != -1).mean()\n        return loss\n\n\n    def criterion_mse(logits, labels):\n        loss_fct = nn.MSELoss(reduction='none')\n        truth = labels.view(-1, 1)\n        loss = loss_fct(logits.view(-1, 1), truth)\n        loss = torch.masked_select(loss, truth != -1).mean()\n        return loss\n\n    # ----------------------------------------------------------------------------------------------------\n    # 0.7 Optimizers & schedulers\n    # ----------------------------------------------------------------------------------------------------\n\n    class AdamW(torch.optim.Optimizer):\n        \"\"\"\n        Implements Adam algorithm with weight decay fix as introduced in\n        `Decoupled Weight Decay Regularization <https://arxiv.org/abs/1711.05101>`__.\n\n        Parameters:\n            params (:obj:`Iterable[torch.nn.parameter.Parameter]`):\n                Iterable of parameters to optimize or dictionaries defining parameter groups.\n            lr (:obj:`float`, `optional`, defaults to 1e-3):\n                The learning rate to use.\n            betas (:obj:`Tuple[float,float]`, `optional`, defaults to (0.9, 0.999)):\n                Adam's betas parameters (b1, b2).\n            eps (:obj:`float`, `optional`, defaults to 1e-6):\n                Adam's epsilon for numerical stability.\n            weight_decay (:obj:`float`, `optional`, defaults to 0):\n                Decoupled weight decay to apply.\n            correct_bias (:obj:`bool`, `optional`, defaults to `True`):\n                Whether ot not to correct bias in Adam (for instance, in Bert TF repository they use :obj:`False`).\n        \"\"\"\n\n        def __init__(\n            self,\n            params,\n            lr = 1e-3,\n            betas = (0.9, 0.999),\n            eps= 1e-6,\n            weight_decay = 0.0,\n            correct_bias = True,\n        ):\n            if lr < 0.0:\n                raise ValueError(\"Invalid learning rate: {} - should be >= 0.0\".format(lr))\n            if not 0.0 <= betas[0] < 1.0:\n                raise ValueError(\"Invalid beta parameter: {} - should be in [0.0, 1.0[\".format(betas[0]))\n            if not 0.0 <= betas[1] < 1.0:\n                raise ValueError(\"Invalid beta parameter: {} - should be in [0.0, 1.0[\".format(betas[1]))\n            if not 0.0 <= eps:\n                raise ValueError(\"Invalid epsilon value: {} - should be >= 0.0\".format(eps))\n            defaults = dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay, correct_bias=correct_bias)\n            super().__init__(params, defaults)\n\n        def step(self, closure = None):\n            \"\"\"\n            Performs a single optimization step.\n\n            Arguments:\n                closure (:obj:`Callable`, `optional`): A closure that reevaluates the model and returns the loss.\n            \"\"\"\n            loss = None\n            if closure is not None:\n                loss = closure()\n\n            for group in self.param_groups:\n                for p in group[\"params\"]:\n                    if p.grad is None:\n                        continue\n                    grad = p.grad.data\n                    if grad.is_sparse:\n                        raise RuntimeError(\"Adam does not support sparse gradients, please consider SparseAdam instead\")\n\n                    state = self.state[p]\n\n                    # State initialization\n                    if len(state) == 0:\n                        state[\"step\"] = 0\n                        state[\"exp_avg\"] = torch.zeros_like(p.data) # Exponential moving average of gradient values\n                        state[\"exp_avg_sq\"] = torch.zeros_like(p.data) # Exponential moving average of squared gradient values\n\n                    exp_avg, exp_avg_sq = state[\"exp_avg\"], state[\"exp_avg_sq\"]\n                    beta1, beta2 = group[\"betas\"]\n\n                    state[\"step\"] += 1\n\n                    # Decay the first and second moment running average coefficient\n                    # In-place operations to update the averages at the same time\n                    exp_avg.mul_(beta1).add_(grad, alpha=1.0 - beta1)\n                    exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1.0 - beta2)\n                    denom = exp_avg_sq.sqrt().add_(group[\"eps\"])\n\n                    step_size = group[\"lr\"]\n                    if group[\"correct_bias\"]:  # No bias correction for Bert\n                        bias_correction1 = 1.0 - beta1 ** state[\"step\"]\n                        bias_correction2 = 1.0 - beta2 ** state[\"step\"]\n                        step_size = step_size * math.sqrt(bias_correction2) / bias_correction1\n\n                    p.data.addcdiv_(exp_avg, denom, value=-step_size)\n\n                    # Just adding the square of the weights to the loss function is *not*\n                    # the correct way of using L2 regularization/weight decay with Adam,\n                    # since that will interact with the m and v parameters in strange ways.\n                    #\n                    # Instead we want to decay the weights in a manner that doesn't interact\n                    # with the m/v parameters. This is equivalent to adding the square\n                    # of the weights to the loss with plain (non-momentum) SGD.\n                    # Add weight decay at the end (fixed version)\n                    if group[\"weight_decay\"] > 0.0:\n                        p.data.add_(p.data, alpha=-group[\"lr\"] * group[\"weight_decay\"])\n\n            return loss\n\n\n    def fetch_optimizer(model, cfg):\n        no_decay = [\"bias\", \"LayerNorm.bias\", \"LayerNorm.weight\"]\n        param_optimizer = list(model.named_parameters())\n        if isinstance(cfg.lr, (int, float)):\n            optimizer_parameters = [\n                {'params': [p for n, p in param_optimizer if not any(nd in n for nd in no_decay)],\n                 'weight_decay': cfg.weight_decay, 'lr': cfg.lr},\n                {'params': [p for n, p in param_optimizer if any(nd in n for nd in no_decay)],\n                 'weight_decay': 0.0, 'lr': cfg.lr},\n            ]\n        else:\n            param_optimizer_encoder = [x for x in param_optimizer if x[0].startswith('model.')]\n            param_optimizer_head = [x for x in param_optimizer if not x[0].startswith('model.')]\n            optimizer_parameters = [\n                {'params': [p for n, p in param_optimizer_encoder if not any(nd in n for nd in no_decay)],\n                 'weight_decay': cfg.weight_decay, 'lr': cfg.lr[0][0]},\n                {'params': [p for n, p in param_optimizer_encoder if any(nd in n for nd in no_decay)],\n                 'weight_decay': 0.0, 'lr': cfg.lr[0]},\n                {'params': [p for n, p in param_optimizer_head if not any(nd in n for nd in no_decay)],\n                 'weight_decay': cfg.weight_decay, 'lr': cfg.lr[0][1]},\n                {'params': [p for n, p in param_optimizer_head if any(nd in n for nd in no_decay)],\n                 'weight_decay': 0.0, 'lr': cfg.lr[1]},\n            ]\n        if cfg.adam8:\n            optimizer = bnb.optim.Adam8bit(optimizer_parameters)\n        elif cfg.optimizer == 'AdamW':\n            optimizer = AdamW(optimizer_parameters, eps=cfg.eps, betas=cfg.betas)\n        elif cfg.optimizer == 'AdamW0':\n            optimizer = torch.optim.AdamW(optimizer_parameters, eps=cfg.eps, betas=cfg.betas)\n        elif cfg.optimizer == 'Adamw':\n            optimizer = torch.optim.Adam(optimizer_parameters, lr=cfg.lr)\n        elif cfg.optimizer == 'Adam':\n            optimizer = torch.optim.Adam(model.parameters(), lr=cfg.lr)\n        elif cfg.optimizer == 'RAdam':\n            if isinstance(cfg.lr, (int, float)):\n                optimizer = torch.optim.RAdam(model.parameters(), lr=cfg.lr, weight_decay=cfg.weight_decay or 0)\n            else:\n                optimizer = torch.optim.RAdam(optimizer_parameters)\n        elif cfg.optimizer == 'Lion':\n            optimizer = Lion(model.parameters(), lr=cfg.lr, weight_decay=cfg.weight_decay or 0)\n\n        # restore state if applicable\n        optimizer_loaded = False\n        if cfg.saved_optimizer:\n            filename = cfg.saved_optimizer\n            if os.path.exists(filename):\n                print(f\"Loading optimizer '{filename}'\")\n                try:\n                    checkpoint = torch.load(filename)\n                except RuntimeError:\n                    checkpoint = torch.load(filename, map_location=torch.device('cpu'))\n                optimizer.load_state_dict(checkpoint['optimizer_states'][-1])\n                optimizer_loaded = True\n            else:\n                # set initial learning rate if resuming training without a saved optimizer state.\n                print('optimizer not found.')\n        if cfg.previous_step is not None and cfg.previous_step != -1 and not optimizer_loaded:\n            for group in optimizer.param_groups:\n                group.setdefault('initial_lr', group['lr'])\n        return optimizer\n\n\n    def fetch_scheduler(cfg, optimizer):\n        previous_step = cfg.previous_step or -1\n        if cfg.schedule_type in ['linear', 'cos']:\n            pct_start = cfg.num_warmup_steps / cfg.num_train_steps\n            if pct_start > 0 and pct_start * cfg.num_train_steps <= 1:\n                pct_start = 1 / cfg.num_train_steps + .001   # to avoid exception with small samples\n\n            if isinstance(cfg.lr, (int, float)):\n                start_lr = cfg.start_lr or (cfg.lr / 25 if cfg.warmup_share else cfg.lr)\n                end_lr = cfg.end_lr or cfg.lr / 25\n                div_factor = cfg.lr / start_lr\n                final_div_factor = start_lr / end_lr\n                max_lr = cfg.lr\n            else:\n                div_factor = 25 if cfg.warmup_share else 1\n                final_div_factor = 25\n                max_lr = (cfg.lr[0], cfg.lr[0], cfg.lr[1], cfg.lr[1])\n                \n            scheduler = torch.optim.lr_scheduler.OneCycleLR(\n                optimizer, anneal_strategy=cfg.schedule_type, last_epoch=previous_step,\n                total_steps=cfg.num_train_steps, pct_start=pct_start,\n                max_lr=max_lr, div_factor=div_factor, final_div_factor=final_div_factor)\n        \n        elif cfg.schedule_type == 'multi_cos':\n            gamma = get_gamma(\n                cfg.lr, cfg.end_lr, cfg.num_train_steps - cfg.num_warmup_steps)\n            num_warmup_steps = cfg.num_warmup_steps\n            start_factor = 1e-5 / cfg.lr if cfg.start_lr is None else cfg.start_lr / cfg.lr\n            scheduler1 = torch.optim.lr_scheduler.LinearLR(\n                optimizer, start_factor=start_factor, end_factor=1, total_iters=num_warmup_steps)\n            cycle_steps = int(cfg.num_train_steps / cfg.num_cycles)\n            scheduler2 = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(\n                optimizer, cycle_steps, T_mult=1, eta_min=cfg.end_lr)\n            scheduler = torch.optim.lr_scheduler.SequentialLR(\n                optimizer, schedulers=[scheduler1, scheduler2],\n                milestones=[cfg.num_warmup_steps], last_epoch=-1)\n        \n        if cfg.schedule_type == 'multi_lr':  # lr = [(1e-5, 1e-3), (3, 2.5e-4)]\n            pct_start = cfg.num_warmup_steps / cfg.num_train_steps\n            # scheduler 1: split encoder / head\n            lr1 = cfg.lr[0]\n            div_factor = 25 if cfg.warmup_share else 1\n            final_div_factor = 25\n            max_lr = (lr1[0], lr1[0], lr1[1], lr1[1])\n            scheduler_1 = torch.optim.lr_scheduler.OneCycleLR(\n                optimizer, anneal_strategy='cos', last_epoch=previous_step,\n                total_steps=cfg.num_train_steps, pct_start=pct_start,\n                max_lr=max_lr, div_factor=div_factor, final_div_factor=final_div_factor)\n            \n            # scheduler 2: unique lr\n            lr2_epoch, lr2 = cfg.lr[1]\n            lr2_step = int(lr2_epoch * config.train_size / (config.train_batch_size * config.accumulation_steps))\n            start_lr = cfg.start_lr or (lr2 / 25 if cfg.warmup_share else lr2)\n            end_lr = cfg.end_lr or lr2 / 25\n            div_factor = 25 if cfg.warmup_share else 1\n            final_div_factor = lr2 / end_lr\n            max_lr = lr2\n            scheduler_2 = OneCycleLRCustom(\n                optimizer, anneal_strategy='cos', last_epoch=lr2_step - 1,\n                total_steps=cfg.num_train_steps - lr2_step,\n                max_lr=max_lr, div_factor=div_factor, final_div_factor=final_div_factor)\n            \n            # combined scheduler\n            scheduler = torch.optim.lr_scheduler.SequentialLR(\n                optimizer, schedulers=[scheduler_1, scheduler_2],\n                milestones=[lr2_step], last_epoch=-1)\n        \n        elif cfg.schedule_type == 'StepLR':\n            gamma = cfg.gamma or 0.999\n            scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma=gamma, last_epoch=previous_step)\n        \n        elif cfg.schedule_type == 'exponential':\n            gamma = cfg.gamma or get_gamma(cfg.lr, cfg.end_lr, cfg.num_train_steps)\n            print(gamma)\n            scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=gamma, last_epoch=previous_step)\n        \n        elif cfg.schedule_type == 'exponential_warmup':\n            gamma = cfg.gamma or get_gamma(\n                cfg.lr, cfg.end_lr, cfg.num_train_steps - cfg.num_warmup_steps)\n            num_warmup_steps = cfg.num_warmup_steps\n            start_factor = 1e-5 / cfg.lr if cfg.start_lr is None else cfg.start_lr / cfg.lr\n            scheduler1 = torch.optim.lr_scheduler.LinearLR(\n                optimizer, start_factor=start_factor, end_factor=1, total_iters=num_warmup_steps)\n            scheduler2 = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=gamma, last_epoch=previous_step)\n            scheduler = torch.optim.lr_scheduler.SequentialLR(\n                optimizer, schedulers=[scheduler1, scheduler2],\n                milestones=[cfg.num_warmup_steps], last_epoch=previous_step)\n            #scheduler = torch.optim.lr_scheduler.ChainedScheduler([scheduler1, scheduler2])\n        \n        elif cfg.schedule_type == 'rop':\n            scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, **cfg.rop)\n\n        # restore state if applicable\n        if cfg.saved_scheduler:\n            filename = cfg.saved_scheduler\n            if os.path.exists(filename):\n                print(f\"Loading scheduler '{filename}'\")\n                try:\n                    checkpoint = torch.load(filename)\n                except RuntimeError:\n                    checkpoint = torch.load(filename, map_location=torch.device('cpu'))\n                scheduler.load_state_dict(checkpoint['lr_schedulers'][-1])\n                # if hasattr(scheduler, 'last_epoch'):\n                #     cfg.previous_step = scheduler.last_epoch\n                # else:\n                #     cfg.previous_step = scheduler._schedulers[0].last_epoch\n            else:\n                print('scheduler not found.')\n        return scheduler\n\n\n    def get_gamma(start_lr, end_lr, steps):\n        \"\"\" return 1 - gamma such that gamma gets us from start_lr to end_lr after 'steps' \"\"\"\n        diff = math.log(end_lr) - math.log(start_lr)\n        rate = diff / steps\n        return math.exp(rate)\n\n\n    class OneCycleLRCustom(torch.optim.lr_scheduler._LRScheduler):\n        \"\"\" similar to default OneCycleLR but uses it's own lr rather than the optimizers.\n            To use as second scheduler in a SequentialLR.\n        \"\"\"\n\n        def __init__(self, optimizer, max_lr, total_steps=None, epochs=None, steps_per_epoch=None,\n                     pct_start = 0.3, anneal_strategy=\"cos\", cycle_momentum=True, base_momentum=0.85,\n                     max_momentum=0.95, div_factor=25.0, final_div_factor=1e4, three_phase=False,\n                     last_epoch=-1, verbose=\"deprecated\"):\n            # Validate optimizer\n            # if not isinstance(optimizer, Optimizer):\n            #     raise TypeError(f\"{type(optimizer).__name__} is not an Optimizer\")\n            self.optimizer = optimizer\n\n            # Validate total_steps\n            if total_steps is not None:\n                if total_steps <= 0 or not isinstance(total_steps, int):\n                    raise ValueError(\n                        f\"Expected positive integer total_steps, but got {total_steps}\"\n                    )\n                self.total_steps = total_steps\n            elif epochs is not None and steps_per_epoch is not None:\n                if not isinstance(epochs, int) or epochs <= 0:\n                    raise ValueError(f\"Expected positive integer epochs, but got {epochs}\")\n                if not isinstance(steps_per_epoch, int) or steps_per_epoch <= 0:\n                    raise ValueError(\n                        f\"Expected positive integer steps_per_epoch, but got {steps_per_epoch}\"\n                    )\n                self.total_steps = epochs * steps_per_epoch\n            else:\n                raise ValueError(\n                    \"You must define either total_steps OR (epochs AND steps_per_epoch)\"\n                )\n\n            self._schedule_phases: List[_SchedulePhase]\n            self._schedule_phases = [{\n                \"end_step\": self.total_steps - 1,\n                \"start_lr\": \"max_lr\",\n                \"end_lr\": \"min_lr\",\n                \"start_momentum\": \"base_momentum\",\n                \"end_momentum\": \"max_momentum\"}]\n\n            # Validate anneal_strategy\n            if anneal_strategy not in [\"cos\", \"linear\"]:\n                raise ValueError(f\"anneal_strategy must be one of 'cos' or 'linear', instead got {anneal_strategy}\")\n            else:\n                self._anneal_func_type = anneal_strategy\n\n            # Initialize learning rate variables\n            max_lrs = max_lr\n            if last_epoch == -1:\n                for idx, group in enumerate(self.optimizer.param_groups):\n                    group[\"initial_lr\"] = max_lrs[idx] / div_factor\n                    group[\"max_lr\"] = max_lrs[idx]\n                    group[\"min_lr\"] = group[\"initial_lr\"] / final_div_factor\n\n            # Initialize momentum variables\n            self.cycle_momentum = cycle_momentum\n            if self.cycle_momentum:\n                if (\"momentum\" not in self.optimizer.defaults and \"betas\" not in self.optimizer.defaults):\n                    raise ValueError(\"optimizer must support momentum or beta1 with `cycle_momentum` option enabled\")\n                self.use_beta1 = \"betas\" in self.optimizer.defaults\n                max_momentums = torch.optim.lr_scheduler._format_param(\"max_momentum\", optimizer, max_momentum)\n                base_momentums = torch.optim.lr_scheduler._format_param(\"base_momentum\", optimizer, base_momentum)\n                if last_epoch == -1:\n                    for m_momentum, b_momentum, group in zip(max_momentums, base_momentums, optimizer.param_groups):\n                        if self.use_beta1:\n                            group[\"betas\"] = (m_momentum, *group[\"betas\"][1:])\n                        else:\n                            group[\"momentum\"] = m_momentum\n                        group[\"max_momentum\"] = m_momentum\n                        group[\"base_momentum\"] = b_momentum\n\n            self.start_lr = max_lr\n            self.end_lr = max_lr / final_div_factor\n            super().__init__(optimizer, last_epoch, verbose)\n\n        def _anneal_func(self, *args, **kwargs):\n            if hasattr(self, \"_anneal_func_type\"):\n                if self._anneal_func_type == \"cos\":\n                    return self._annealing_cos(*args, **kwargs)\n                elif self._anneal_func_type == \"linear\":\n                    return self._annealing_linear(*args, **kwargs)\n                else:\n                    raise ValueError(f\"Unknown _anneal_func_type: {self._anneal_func_type}\")\n            else:\n                # For BC\n                return self.anneal_func(*args, **kwargs)  # type: ignore[attr-defined]\n\n        @staticmethod\n        def _annealing_cos(start, end, pct):\n            \"\"\"Cosine anneal from `start` to `end` as pct goes from 0.0 to 1.0.\"\"\"\n            cos_out = math.cos(math.pi * pct) + 1\n            return end + (start - end) / 2.0 * cos_out\n\n        @staticmethod\n        def _annealing_linear(start, end, pct):\n            \"\"\"Linearly anneal from `start` to `end` as pct goes from 0.0 to 1.0.\"\"\"\n            return (end - start) * pct + start\n\n        def get_lr(self):\n            \"\"\"Compute the learning rate of each parameter group.\"\"\"\n            torch.optim.lr_scheduler._warn_get_lr_called_within_step(self)\n\n            lrs = []\n            step_num = self.last_epoch\n\n            if step_num > self.total_steps:\n                raise ValueError(\n                    f\"Tried to step {step_num} times. The specified number of total steps is {self.total_steps}\"  # noqa: UP032\n                )\n\n            for group in self.optimizer.param_groups:\n                start_step = 0.0\n                for i, phase in enumerate(self._schedule_phases):\n                    end_step = phase[\"end_step\"]\n                    if step_num <= end_step or i == len(self._schedule_phases) - 1:\n                        pct = (step_num - start_step) / (end_step - start_step)\n                        computed_lr = self._anneal_func(self.start_lr, self.end_lr, pct)\n                        if self.cycle_momentum:\n                            computed_momentum = self._anneal_func(\n                                group[phase[\"start_momentum\"]],\n                                group[phase[\"end_momentum\"]],\n                                pct,\n                            )\n                        break\n                    start_step = phase[\"end_step\"]\n\n                lrs.append(computed_lr)  # type: ignore[possibly-undefined]\n                if self.cycle_momentum:\n                    if self.use_beta1:\n                        group[\"betas\"] = (computed_momentum, *group[\"betas\"][1:])  # type: ignore[possibly-undefined]\n                    else:\n                        group[\n                            \"momentum\"\n                        ] = computed_momentum  # type: ignore[possibly-undefined]\n\n            return lrs\n\n\n    class CosineAnnealingWarmupRestarts(torch.optim.lr_scheduler._LRScheduler):\n        \"\"\"\n            optimizer (Optimizer): Wrapped optimizer.\n            first_cycle_steps (int): First cycle step size.\n            cycle_mult(float): Cycle steps magnification. Default: -1.\n            max_lr(float): First cycle's max learning rate. Default: 0.1.\n            min_lr(float): Min learning rate. Default: 0.001.\n            warmup_steps(int): Linear warmup step size. Default: 0.\n            gamma(float): Decrease rate of max learning rate by cycle. Default: 1.\n            last_epoch (int): The index of last epoch. Default: -1.\n        \"\"\"\n\n        def __init__(self,\n                     optimizer: torch.optim.Optimizer,\n                     first_cycle_steps: int,\n                     cycle_mult: float = 1.,\n                     max_lr: float = 0.1,\n                     min_lr: float = 0.001,\n                     warmup_steps: int = 0,\n                     gamma: float = 1.,\n                     last_epoch: int = -1\n                     ):\n            assert warmup_steps < first_cycle_steps\n\n            self.first_cycle_steps = first_cycle_steps  # first cycle step size\n            self.cycle_mult = cycle_mult  # cycle steps magnification\n            self.base_max_lr = max_lr  # first max learning rate\n            self.max_lr = max_lr  # max learning rate in the current cycle\n            self.min_lr = min_lr  # min learning rate\n            self.warmup_steps = warmup_steps  # warmup step size\n            self.gamma = gamma  # decrease rate of max learning rate by cycle\n\n            self.cur_cycle_steps = first_cycle_steps  # first cycle step size\n            self.cycle = 0  # cycle count\n            self.step_in_cycle = last_epoch  # step size of the current cycle\n\n            super(CosineAnnealingWarmupRestarts, self).__init__(optimizer, last_epoch)\n\n            # set learning rate min_lr\n            self.init_lr()\n\n        def init_lr(self):\n            self.base_lrs = []\n            for param_group in self.optimizer.param_groups:\n                param_group['lr'] = self.min_lr\n                self.base_lrs.append(self.min_lr)\n\n        def get_lr(self):\n            if self.step_in_cycle == -1:\n                return self.base_lrs\n            elif self.step_in_cycle < self.warmup_steps:\n                return [(self.max_lr - base_lr) * self.step_in_cycle / self.warmup_steps + base_lr for base_lr in self.base_lrs]\n            else:\n                return [base_lr + (self.max_lr - base_lr)\n                        * (1 + math.cos(math.pi * (self.step_in_cycle - self.warmup_steps)\n                                        / (self.cur_cycle_steps - self.warmup_steps))) / 2\n                        for base_lr in self.base_lrs]\n\n        def step(self, epoch=None):\n            if epoch is None:\n                epoch = self.last_epoch + 1\n                self.step_in_cycle = self.step_in_cycle + 1\n                if self.step_in_cycle >= self.cur_cycle_steps:\n                    self.cycle += 1\n                    self.step_in_cycle = self.step_in_cycle - self.cur_cycle_steps\n                    self.cur_cycle_steps = int((self.cur_cycle_steps - self.warmup_steps) * self.cycle_mult) + self.warmup_steps\n            else:\n                if epoch >= self.first_cycle_steps:\n                    if self.cycle_mult == 1.:\n                        self.step_in_cycle = epoch % self.first_cycle_steps\n                        self.cycle = epoch // self.first_cycle_steps\n                    else:\n                        n = int(math.log((epoch / self.first_cycle_steps * (self.cycle_mult - 1) + 1), self.cycle_mult))\n                        self.cycle = n\n                        self.step_in_cycle = epoch - int(self.first_cycle_steps * (self.cycle_mult ** n - 1) / (self.cycle_mult - 1))\n                        self.cur_cycle_steps = self.first_cycle_steps * self.cycle_mult ** (n)\n                else:\n                    self.cur_cycle_steps = self.first_cycle_steps\n                    self.step_in_cycle = epoch\n\n            self.max_lr = self.base_max_lr * (self.gamma**self.cycle)\n            self.last_epoch = math.floor(epoch)\n            for param_group, lr in zip(self.optimizer.param_groups, self.get_lr()):\n                param_group['lr'] = lr\n\n    # ----------------------------------------------------------------------------------------------------\n    # 0.8 Lightning training\n    # ----------------------------------------------------------------------------------------------------\n\n    def get_amp_context(device_type, datatype):\n        \"\"\" get context for \"\"\"\n        ptdtype = {'float32': torch.float32, 'bfloat16': torch.bfloat16, 'float16': torch.float16}[datatype]\n        device_type = 'cuda' if 'cuda' in str(device_type) else str(device_type)\n        #context = nullcontext() if device_type == 'cpu' else torch.amp.autocast(device_type=device_type, dtype=ptdtype)\n        context = nullcontext() if device_type == 'cpu' else torch.cuda.amp.autocast(dtype=ptdtype)  # for torch v1.11\n        return context\n\n\n    def pb_format_num(n):\n        \"\"\" format numbers in progress bar \"\"\"\n        f = '{0:.4g}'.format(n).replace('+0', '+').replace('-0', '-')\n        n = str(n)\n        return f if len(f) < len(n) else n\n\n\n    class My_progress_bar(TQDMProgressBar):\n        \"\"\" customize TQDMProgressBar to add lr \"\"\"\n        def get_metrics(self, trainer, pl_module):\n            items = super().get_metrics(trainer, pl_module)\n            items.pop('v_num', None)  # added in base class; contains version number\n            scheduler = trainer.model.lr_schedulers()\n            lr = scheduler.get_last_lr()[0] if hasattr(scheduler, 'get_last_lr') else trainer.model.optimizers().param_groups[0]['lr']\n            items['lr'] = lr[0] if isinstance(lr, list) else lr\n            items = {k: pb_format_num(v) for k, v in items.items()}\n            return items\n\n        def init_validation_tqdm(self):\n            \"\"\" disable validation progress bar \"\"\"\n            bar = tqdm(disable=True)\n            return bar\n            \n\n    def fg_format_seconds(seconds):\n        seconds = int(seconds)\n        minutes = seconds // 60\n        remaining_seconds = seconds % 60\n        return f\"{minutes:02}:{remaining_seconds:02}\"\n\n\n    class Exponential_average(object):\n        def __init__(self, decay=.01):\n            self.decay = decay\n            self.val = None\n\n        def reset(self):\n            pass\n\n        def update(self, pred=None):\n            val = pr_auc_score(true, pred)\n            if self.keep_pred:\n                self.preds.append(pred)\n            self.val = (self.val*(1-self.decay) + val*self.decay) if self.val is not None else val\n\n        @property\n        def avg(self):\n            return self.val\n\n        def get_pred(self):\n            res = np.concatenate(self.preds, axis=0)\n            self.preds = []\n            return res\n\n\n    class AverageMeter:\n        def __init__(self, labels=None):\n            self.reset()\n            self.auc = None\n            self.labels = labels if labels is not None else None\n            \n        def reset(self):\n            self.loss_sum = 0\n            self.acc_sum = 0\n            self.count = 0\n            self.preds = []\n        \n        def update(self, loss, acc):\n            self.loss_sum += loss\n            self.acc_sum += acc\n            self.count += 1\n            return loss\n        \n        def update_preds(self, preds):\n            self.preds.append(preds.detach().cpu())\n        \n        @property\n        def loss(self):\n            return self.loss_sum / self.count if self.count > 0 else 0\n\n        @property\n        def acc(self):\n            return self.acc_sum / self.count if self.count > 0 else 0\n        \n        def calculate_auc(self):\n            self.preds = np.concatenate(self.preds, axis=0)\n            len_preds = self.preds.shape[0]  # to support sanity check\n            self.auc = auc_multi(self.preds, self.labels[:len_preds])\n            \n\n    class Pl_model(L.LightningModule):\n        def __init__(self, model, cfg, val_labels=None, silent=None):\n            super().__init__()\n            self.model = model\n            self.cfg = cfg\n            self.silent = silent if silent is not None else False\n            self.save_hyperparameters()\n\n            self.use_mask = cfg.use_mask\n            self.automatic_optimization = False\n            self.current_amp_dtype = None\n            self.current_loss_name = None\n            self.saved_inputs = [] if cfg.awp or cfg.fgm else None\n            self._setup_amp()\n            self._setup_loss()\n            self.train_meter = AverageMeter()\n            self.val_meter = AverageMeter(val_labels) if val_labels is not None else None\n            self.t0 = time.time()\n\n        def _setup_amp(self):\n            \"\"\" adjust automatic mode precision dtype based on epoch \"\"\"\n            dtype = self.cfg.precision_schedule.get(self.current_epoch, None)\n            if dtype is not None and self.current_amp_dtype != dtype:\n                self.amp_context = get_amp_context(self.device, dtype)\n                if not self.silent:\n                    print(f\"Switched amp context to {dtype}\")\n                self.scaler = torch.cuda.amp.GradScaler(enabled=(dtype == 'float16'),\n                                                        growth_interval=600)\n            assert hasattr(self, 'amp_context'), \"Amp context not set, check precision schedule\"\n            self.current_amp_dtype = dtype\n            os.environ['EPOCH'] = str(self.current_epoch)\n\n        def _setup_loss(self):\n            \"\"\" adjust loss based on epoch \"\"\"\n            loss_name = self.cfg.loss_schedule.get(self.current_epoch, None)\n            if loss_name is not None and self.current_loss_name != loss_name:\n                self.loss_fn = get_loss_fn(loss_name)\n                self.current_loss_name = loss_name\n\n        def _compute_loss_metrics(self, pred, data, meter):\n            y = data['y']\n            mask = data.get('mask', None)\n            loss = self.loss_fn(pred, y)\n            if self.use_mask and mask is not None:\n                loss = (loss * mask).sum() / mask.sum()\n            else:\n                loss = loss.mean()\n                \n            # assert loss >= 0\n            acc = rank_based_accuracy_multi(pred, y)\n            meter.update(loss.item(), acc)\n            return loss\n\n        def on_train_start(self):\n            pass\n            # self._setup_amp()\n\n        def on_train_epoch_start(self):\n            self._setup_amp()\n            self._setup_loss()\n            self.train_meter.reset()\n            if self.val_meter:\n                self.val_meter.reset()\n            self.t0 = time.time()                     \n        \n        def training_step(self, batch, batch_idx):\n            cfg = self.cfg\n            meter = self.train_meter\n            if cfg.aug.mixup:\n                mixup(batch)\n            if self.saved_inputs is not None:\n                self.saved_inputs.append(batch)\n            with self.amp_context:\n                preds = self(batch['x'])\n                loss = self._compute_loss_metrics(preds, batch, meter)\n            out = dict(loss=meter.loss, acc=meter.acc)\n            if torch.isnan(loss).any():\n                print('ERROR: loss is nan. Interrupting training.')\n                self.trainer.should_stop = True\n            self.log_dict(out, prog_bar=True)  # prog_bar shows most recent value, not mean, which is fine \n                                               # because out already contains the mean.\n            # print(loss.item(), meter.acc, trainer.callback_metrics)\n            self.manual_backward(self.scaler.scale(loss))\n\n            if (batch_idx + 1) % cfg.accumulation_steps == 0:  # replace by GradientAccumulationScheduler callback?\n                score = self.val_meter.auc if self.val_meter else self.train_meter.loss\n                if cfg.awp and score is not None and score >= cfg.adv_th:\n                    self.awp.attack_backward(self.saved_inputs)\n                if cfg.fgm and score is not None and score >= cfg.adv_th:\n                    self.fgm.attack_backward(self.saved_inputs)\n\n                optimizer = self.optimizers()\n                self.scaler.unscale_(optimizer)\n                if cfg.max_grad_value is not None:\n                    torch.nn.utils.clip_grad_value_(self.parameters(), cfg.max_grad_value)\n                elif cfg.max_grad_norm is not None:\n                    norm_grad = torch.nn.utils.clip_grad_norm_(parameters=self.parameters(),\n                                                               max_norm=cfg.max_grad_norm)\n                self.scaler.step(optimizer)\n                self.scaler.update()\n                optimizer.zero_grad(set_to_none=True)\n                if cfg.step_scheduler_after == 'step':\n                    self.lr_schedulers().step()\n                if self.saved_inputs is not None:\n                    self.saved_inputs.clear()\n            # return loss_metrics  # don't return because we do the backpass here\n\n        def on_train_epoch_end(self):\n            \"\"\" invoked after validation to process outputs of training_step. \"\"\"\n            ds = self.trainer.train_dataloader.dataset\n            if hasattr(ds, 'on_epoch_end'):\n                ds.on_epoch_end()\n            if self.cfg.step_scheduler_after == 'epoch':  # needed due to manual optimization  \n                self.lr_schedulers().step()\n            if self.val_meter is None:\n                loss_metrics_str = f'loss: {self.train_meter.loss:.4f} - acc: {self.train_meter.acc:.4f}'\n                scheduler = self.lr_schedulers()\n                lr = scheduler.get_last_lr()[0] if hasattr(scheduler, 'get_last_lr') else self.optimizers().param_groups[0]['lr']\n                lr = lr[0] if isinstance(lr, list) else lr\n                lr = f'{lr:.3e}'.replace('+0', '+').replace('-0', '-')\n                if not self.silent:                   \n                    print(f'epoch:{self.current_epoch:>3} - {loss_metrics_str} - lr: {lr}')\n            return super().on_train_epoch_end()\n\n        def on_validation_start(self):\n            self._setup_amp()\n\n        def validation_step(self, batch, batch_idx):\n            \"\"\" no logging per step \"\"\"\n            meter = self.val_meter\n            with self.amp_context:\n                preds = self(batch['x'])\n                self._compute_loss_metrics(preds, batch, meter)\n                meter.update_preds(preds)\n            # out = dict(val_loss=meter.loss, val_kl=meter.kl)\n            # self.log_dict(out, on_step=True, prog_bar=True)\n\n        def on_validation_epoch_end(self):\n            \"\"\" process outputs of validation_step \"\"\"\n            self.val_meter.calculate_auc()\n            out = dict(val_loss=self.val_meter.loss, val_acc=self.val_meter.acc, val_auc=self.val_meter.auc)\n            self.log_dict(out, prog_bar=False)  # disable prog_bar to simplify progress bar; we write it to the screen anyway\n            loss_metrics_str = f'loss: {self.train_meter.loss:.4f} - acc: {self.train_meter.acc:.4f} - ' \\\n                               f'val_loss: {self.val_meter.loss:.4f} - val_acc: {self.val_meter.acc:.4f} - ' \\\n                               f'val_auc: {self.val_meter.auc:.4f}'\n            scheduler = self.lr_schedulers()\n            lr = scheduler.get_last_lr()[0] if hasattr(scheduler, 'get_last_lr') else self.optimizers().param_groups[0]['lr']\n            lr = lr[0] if isinstance(lr, list) else lr\n            lr = f'{lr:.3e}'.replace('+0', '+').replace('-0', '-')\n            duration = time.time() - self.t0\n            if not self.silent:                                \n                print(f'epoch:{self.current_epoch:>3} - {loss_metrics_str} - lr: {lr} - {duration:,.0f}s', ' '*40)\n            \n            # print() # Keep the previous progbar\n            # self.trainer.val_dataloaders[0].dataset.on_epoch_end()\n            return super().on_validation_epoch_end()\n\n        def forward(self, batch):\n            return self.model(batch)\n\n        def configure_optimizers(self):\n            cfg = self.cfg\n            optimizer = fetch_optimizer(self.model, cfg)\n            if self.cfg.awp:\n                self.awp = AWP(self, optimizer, adv_lr=cfg.adv_lr, adv_eps=cfg.adv_eps, cfg=cfg)\n            if self.cfg.fgm:\n                self.fgm = FGM(self, optimizer, self.scaler, epsilon=cfg.fgm_epsilon, emb_name='word_embeddings')\n            scheduler = fetch_scheduler(self.cfg, optimizer)\n            lr_scheduler_config = {\"scheduler\": scheduler, \"interval\": self.cfg.step_scheduler_after}\n            if cfg.monitor:\n                lr_scheduler_config['monitor'] = cfg.monitor\n            return {\"optimizer\": optimizer, \"lr_scheduler\": lr_scheduler_config}\n\n    \n    def train(meta, files, no_try=False, debug=None):\n        global config, RANK\n        RANK = torch.range(1, N_BIRDS).reshape(1, -1).expand(config.train_batch_size, -1)\n        config.img_dim = SPEC_DIMS.get(config.spec, (128, 256))\n        \n        my_log = Log(f'{config.name}.log', time_key=False)\n        logger = get_logger(config)\n        my_log(f'config: {config}\\n')\n        folds = config.folds\n        idx_birds = None\n        \n        for fold in folds:\n            print(f'{\"-\"*40} {config.name}_f{fold} {\"-\"*40}')\n            if config.monitor == 'val_auc':\n                train_idx = (meta.fold.values != fold)\n                val_idx = (meta.fold.values == fold)\n                if config.filter[-1] > 0:\n                    train_idx &= meta.rating >= config.filter[-1]\n                if config.dataset_val:\n                    val_labels = load_obj(f'{ROOT}/b5-data-{config.dataset_val}/y_f{fold}')\n                else:\n                    val_labels = labels[val_idx]\n            else:\n                train_idx, val_idx = np.full(meta.shape[0], True), None\n                val_labels = None\n            if debug:\n                print(f'\\n*** DEBUG: {debug} ***\\n')\n                n = meta.shape[0]\n                train_idx = np.array([True] * debug + [False] * (n - debug))\n                np.random.shuffle(train_idx)\n                val_idx = np.array([True] * debug + [False] * (n - debug))\n                np.random.shuffle(val_idx)\n            meta_train = meta[train_idx]\n            train_size = config.train_size or int(meta_train.ws.sum())\n            if debug:\n                train_size = debug\n            if val_idx is not None:\n                print(f'train: {train_size:,} | {len(meta_train):,} -  val: {val_idx.sum():,}')\n            else:\n                print(f'train: {train_size:,} | {len(meta_train):,} - val: n/a')\n            \n            # set ws for train sample\n            ws_power = config.ws_power or 1/2\n            bird_sum = meta_train.groupby('primary_label')['ws'].transform('sum').values\n            sampling_th = 6\n            ws = meta_train.ws.values / np.where(\n                bird_sum <= sampling_th, bird_sum, sampling_th + np.power(bird_sum - sampling_th, ws_power))\n            times_picked = ws * bird_sum\n            print(f'raw weights: {times_picked.min():.2f} - {times_picked.max():.2f}')\n            ws /= ws.sum()\n            \n            epochs = config.epochs\n            config.num_train_steps = int(train_size / (config.train_batch_size * config.accumulation_steps) * epochs)\n            config.num_warmup_steps = int(config.num_train_steps * config.warmup_share)\n            print(f'train steps: {config.num_train_steps:,} - warmup steps: {config.num_warmup_steps:,}')\n            custom_filter = build_custom_filter(meta[train_idx])\n            data_module = Ba_data_module_montage_h5(\n                files, labels, birds, train_idx, val_idx, config.dataset_val, fold, train_mode=config.train_mode,\n                ws=ws, files_background_noise=files_background_noise if config.background_noise_prob else None, idx_birds=idx_birds,\n                custom_filter=custom_filter, train_size=train_size)\n\n            if 'resume' in mode:\n                print('Resume training')\n                trainer = L.Trainer()\n                pl_model = Pl_model.load_from_checkpoint(config.checkpoint)\n                trainer.fit(pl_model, data_module, ckpt_path=config.checkpoint)\n            else:\n                model_class = globals()[config.model_name]\n                pretrained = config.pretrained_encoder if config.pretrained_encoder is not None else True\n                in_chans = config.in_chans or 3\n                model = model_class(\n                    config, resize=config.resize, add_position=config.add_position, num_classes=N_BIRDS,\n                    out_indices=config.out_indices, pretrained=pretrained, in_chans=in_chans)\n                inp_dim = (1, config.img_dim[0], config.img_dim[1] * config.img_duration // 5)\n                example_input_array = [torch.zeros(8, *inp_dim)]\n                out_ = model(example_input_array[0])  # sanity check\n                print(f'model out shape: {out_.shape}')\n\n                if config.checkpoint:\n                    print('load existing model')\n                    try:\n                        chk = torch.load(config.checkpoint)\n                    except:\n                        chk = torch.load(config.checkpoint, map_location=torch.device('cpu'))\n                    \n                    # model_state = {k.replace('model.', ''): v for k, v in chk['state_dict'].items()}\n                    model_state = {k[6:]: v for k, v in chk['state_dict'].items()}\n                    model.load_state_dict(model_state)\n\n                pl_model = Pl_model(model, config, val_labels=val_labels)\n                pl_model.example_input_array = example_input_array\n                print_model_size(pl_model)\n                # display(summary(model, col_names=[\"num_params\",\"trainable\"]))  # show layer sizes\n                # model  # show dimensions\n\n                if val_idx is not None:\n                    loss_file = '{epoch}-{val_loss:.4f}'\n                    loss_monitor = 'val_loss'\n                    metric_file = '{epoch}-{val_auc:.4f}'\n                    metric_monitor = 'val_auc'\n                else:\n                    loss_file = '{epoch}-{loss:.4f}'\n                    loss_monitor = 'loss'\n                    metric_file = '{epoch}-{acc:.4f}'\n                    metric_monitor = 'acc'\n                chk_callback_metric = ModelCheckpoint(\n                    dirpath='./', filename=f'{config.name}_f{fold}-{metric_file}',\n                    save_top_k=3, monitor=metric_monitor, mode=config.mode, \n                    save_last=True, save_on_train_epoch_end=True, save_weights_only=True)\n                chk_callback_loss = ModelCheckpoint(\n                    dirpath='./', filename=f'{config.name}_f{fold}-{loss_file}',\n                    save_top_k=3, monitor=loss_monitor, mode='min', \n                    save_last=False, save_on_train_epoch_end=True, save_weights_only=True)\n                callbacks = [chk_callback_metric, chk_callback_loss, \n                             LearningRateMonitor(logging_interval='step'),\n                             My_progress_bar(refresh_rate=config.log_every_n_steps),\n                            ]\n                if config.patience and val_idx is not None:\n                    early_stop_callback = EarlyStopping(\n                        monitor=config.monitor, min_delta=1e-5, patience=config.patience,\n                        verbose=False, mode=config.mode)\n                    callbacks.append(early_stop_callback)\n                if config.swa is not None:\n                    swa_lrs, swa_epoch_start, annealing_epochs = config.swa\n                    swa_callback = StochasticWeightAveraging(\n                        swa_lrs=swa_lrs, swa_epoch_start=swa_epoch_start, annealing_epochs=annealing_epochs,\n                        annealing_strategy='cos')\n                    callbacks.append(swa_callback)\n                    print('***** swa *****')\n                    \n                max_epochs = config.max_epochs if config.max_epochs is not None else epochs\n                num_sanity_val_steps = 1 if is_interactive() and config.monitor != 'acc' else 0\n                # check_val_every_n_epoch = config.check_val_every_n_epoch if config.monitor != 'acc' else None\n                limit_val_batches = None if config.monitor != 'acc' else 0\n                trainer = L.Trainer(max_epochs=max_epochs, max_time=config.duration, logger=logger,\n                                    log_every_n_steps=config.log_every_n_steps,\n                                    check_val_every_n_epoch=config.check_val_every_n_epoch,\n                                    accelerator='auto', devices='auto',\n                                    num_sanity_val_steps=num_sanity_val_steps,\n                                    limit_train_batches=None, limit_val_batches=limit_val_batches,\n                                    enable_checkpointing=True, callbacks=callbacks,\n                                    # gradient_clip_val=1.0,\n                                   )\n                print('start training...')\n                if no_try:\n                    trainer.fit(pl_model, data_module)\n                else:\n                    try:\n                        trainer.fit(pl_model, data_module)\n                    except Exception as e:\n                        print(f'\\ntraining interrupted: {e}')\n\n                trainer.save_checkpoint(f'last_train_{config.name}_f{fold}.ckpt', weights_only=True)  # needed if swa\n                if pl_model.val_meter:\n                    save_obj(pl_model.val_meter.preds, f'preds_{config.name}_f{fold}')\n\n        print('Finished training...')\n        del data_module\n        gc.collect()\n        if config.wandb_project:\n            wandb.finish()\n    \n    # ----------------------------------------------------------------------------------------------------\n    # 0.9 adversarial training classes\n    # ----------------------------------------------------------------------------------------------------\n\n    class FGM():\n        \"\"\" FGM attack. Rather than updating the model using the gradient, it perturbs the weights of the embedding layer\n            using the computed gradients, calculates a perturbed loss and gradients, restores the weights and uses the\n            perturbed gradients to update the weights.\n\n            The emb_name parameter should be replaced with the parameter name of the embedding in your model\n        \"\"\"\n        def __init__(self, model, optimizer, scaler=None,\n                     epsilon=1., emb_name='word_embeddings', do_add_gradients=True):\n            self.model = model\n            self.optimizer = optimizer\n            self.scaler = scaler\n            self.emb_name = emb_name\n            self.epsilon = epsilon\n            self.emb_name = emb_name\n            self.do_add_gradients = do_add_gradients\n            self.backup = {}\n\n        def attack_backward(self, inputs):\n            accumulation_steps = config.accumulation_steps\n            self._attack_step()\n            if not self.do_add_gradients:\n                self.optimizer.zero_grad()  # usually we add the gradients\n            for input in inputs:\n                with torch.cuda.amp.autocast():\n                    tr_logits = self.model(input)\n                    adv_loss = Von_mises_fisher_3D_loss(tr_logits, input)\n                if accumulation_steps > 1:\n                    adv_loss /= accumulation_steps\n                if self.scaler is not None:\n                    self.scaler.scale(adv_loss).backward()\n                else:\n                    adv_loss.backward()\n                torch.nn.utils.clip_grad_norm_(parameters=model.parameters(),\n                                               max_norm=config.max_grad_norm)\n            self._restore()\n            #print(f'fgm {epoch} {adv_loss.item()}')\n\n\n        def _attack_step(self):\n            emb_name = self.emb_name\n            epsilon = self.epsilon\n            for name, param in self.model.named_parameters():\n                if param.requires_grad and emb_name in name:\n                    self.backup[name] = param.data.clone()\n                    norm = torch.norm(param.grad)\n                    if norm != 0 and not torch.isnan(norm):\n                        r_at = epsilon * param.grad / norm\n                        param.data.add_(r_at)\n\n        def _restore(self):\n            # The emb_name parameter should be replaced with the parameter name of the embedding in your model\n            emb_name = self.emb_name\n            for name, param in self.model.named_parameters():\n                if param.requires_grad and emb_name in name:\n                    assert name in self.backup\n                    param.data = self.backup[name]\n            self.backup = {}\n\n        # example of usage\n        # fgm = FGM(model)\n        # for batch_input, batch_label in data:\n        #       loss = model(batch_input, batch_label)\n        #       loss.backward()\n        #\n        #       # adversarial training\n        #       fgm.attack()\n        #       loss_adv = model(batch_input, batch_label)\n        #       loss_adv.backward()\n        #       fgm.restore()\n        #\n        #       optimizer.step()\n        #       model.zero_grad()\n\n    # --------------------\n    # 0.10 AWP\n    # --------------------\n\n    class AWP:\n        \"\"\" adversarial weight perturbation. Rather than updating the model using the gradient, it perturbs the weights\n            using the computed gradients, calculates a perturbed loss and gradients, restores the weights and uses the\n            perturbed gradients to update the weights.\n        \"\"\"\n        def __init__(self, pl_model, optimizer, adv_param=\"weight\", adv_lr=1, adv_eps=0.01, adv_step=1, cfg=None):\n            self.pl_model = pl_model\n            self.optimizer = optimizer\n            self.adv_param = adv_param\n            self.adv_lr = adv_lr\n            self.adv_eps = adv_eps\n            self.adv_step = adv_step\n            self.cfg = cfg\n            self.backup = {}\n            self.backup_eps = {}\n\n        def attack_backward(self, inputs):\n            if (self.adv_lr == 0):\n                return None\n\n            cfg = self.cfg\n            pl_model = self.pl_model\n            model = pl_model.model\n            amp_context = pl_model.amp_context\n            scaler = pl_model.scaler\n            loss_fn = pl_model.loss_fn if pl_model.loss_fn is not None else model.compute_loss\n            accumulation_steps = cfg.accumulation_steps\n\n            self._save()\n            for i in range(self.adv_step):\n                self._attack_step()\n                self.optimizer.zero_grad()\n                for batch in inputs:   # to support accumulation_steps > 1\n                    with amp_context:\n                        preds = pl_model(batch['x'])\n                        y = batch['y']\n                        mask = batch.get('mask', None)\n                        adv_loss = loss_fn(preds, y)\n                        if mask is not None:\n                            adv_loss = (adv_loss * mask).sum() / mask.sum()\n                        else:\n                            adv_loss = adv_loss.mean()\n                        if isinstance(adv_loss, dict):\n                            adv_loss = adv_loss['loss']\n                    if accumulation_steps > 1:\n                        adv_loss /= accumulation_steps\n                    if scaler is not None:\n                        scaler.scale(adv_loss).backward()\n                    else:\n                        adv_loss.backward()\n                    if cfg.max_grad_value is not None:\n                        torch.nn.utils.clip_grad_value_(model.parameters(), cfg.max_grad_value)\n                    elif cfg.max_grad_norm is not None:\n                        norm_grad = torch.nn.utils.clip_grad_norm_(parameters=model.parameters(),\n                                                                   max_norm=cfg.max_grad_norm)\n            self._restore()\n            #print(f'awp {epoch} {adv_loss.item()}')\n\n        def _attack_step(self):\n            e = 1e-6\n            for name, param in self.pl_model.model.named_parameters():\n                if param.requires_grad and param.grad is not None and self.adv_param in name:\n                    norm1 = torch.norm(param.grad)\n                    norm2 = torch.norm(param.data.detach())\n                    if norm1 != 0 and not torch.isnan(norm1):\n                        r_at = self.adv_lr * param.grad / (norm1 + e) * (norm2 + e)\n                        param.data.add_(r_at)\n                        param.data = torch.min(\n                            torch.max(param.data, self.backup_eps[name][0]),\n                            self.backup_eps[name][1]\n                        )\n                    # param.data.clamp_(*self.backup_eps[name])\n\n        def _save(self):\n            for name, param in self.pl_model.model.named_parameters():\n                if param.requires_grad and param.grad is not None and self.adv_param in name:\n                    if name not in self.backup:\n                        self.backup[name] = param.data.clone()\n                        grad_eps = self.adv_eps * param.abs().detach()\n                        self.backup_eps[name] = (\n                            self.backup[name] - grad_eps,\n                            self.backup[name] + grad_eps,\n                        )\n\n        def _restore(self,):\n            for name, param in self.pl_model.model.named_parameters():\n                if name in self.backup:\n                    param.data = self.backup[name]\n            self.backup = {}\n            self.backup_eps = {}\n\n    # ----------------------------------------------------------------------------------------------------\n    # 0.11 dummy lightning logger\n    # ----------------------------------------------------------------------------------------------------\n\n    from pytorch_lightning.loggers.logger import Logger, rank_zero_experiment\n    from pytorch_lightning.utilities import rank_zero_only\n\n    class Dummy_logger(Logger):\n        \"\"\" Needed to bypass bug in default logger class.\n            for pytorch_lightning==1.9.4. \n        \"\"\"\n        \n        @property\n        def name(self):\n            return \"MyLogger\"\n\n        @property\n        def version(self):\n            return \"0.1\"\n\n        @rank_zero_only\n        def log_hyperparams(self, params):\n            # params is an argparse.Namespace\n            pass\n\n        @rank_zero_only\n        def log_metrics(self, metrics, step):\n            # metrics is a dictionary of metric names and values\n            pass\n\n        @rank_zero_only\n        def save(self):\n            # Optional. Any code necessary to save logger data\n            pass\n\n        @rank_zero_only\n        def finalize(self, status):\n            # Optional. Any code that needs to be run after training finishes.\n            pass\n\n    def get_logger(config, group='att', job_type='train'):\n        if config.wandb_project:\n            os.environ['WANDB_SILENT'] = 'true'\n            key = UserSecretsClient().get_secret('wandb')\n            wandb.login(key=key)\n            logger = WandbLogger(name=config.name, project=config.wandb_project, config=config, group=group, job_type=job_type)\n        else:\n            logger = Dummy_logger()\n        return logger\n\n    # ----------------------------------------------------------------------------------------------------\n    # 0.12 submit\n    # ----------------------------------------------------------------------------------------------------\n\n    def merge_stat_dict(m1, m2, layers=['_conv_layers', '_post_processing']):\n        \"\"\" takes the state_dict of two models and replaces 'm1' parameters\n            whose key matches 'layers' with the values from 'm2'\n        \"\"\"\n        for k, v in m1.items():\n            for match in layers:\n                if match in k and k in m2:\n                    m1[k] = m2[k]\n                    continue\n        return m1\n\n\n    def merge_stat_dict_files(m1_path, m2_path, layers=['_conv_layers', '_post_processing']):\n        \"\"\" reads the state_dict of two models given their paths and replaces 'm1' parameters\n            whose key matches 'layers' with the values from 'm2'\n        \"\"\"\n        m1 = torch.load(m1_path)\n        m2 = torch.load(m2_path)\n        return merge_stat_dict(m1, m2, layers)\n\n\n    def merge_models(model, m2_path, layers):\n        \"\"\" takes a model and merges it with the model whose stat_dict is saved in m2_path\n            for the layers whose key matches 'layers'.\n        \"\"\"\n        m1_path = 'temp.pth'\n        model.save_state_dict(m1_path)\n        m1 = merge_stat_dict_files(m1_path, m2_path, layers)\n        model.load_state_dict(m1)\n\n\n    def get_best_checkpoints(model_name):\n        file = f'/kaggle/working/{model_name}.ckpt'\n        if os.path.exists(file):\n            return [file]\n        \n        files = glob.glob(f'/kaggle/working/{model_name}_epoch*.ckpt')\n        sorted_files = []\n        for file in files:\n            metrics = file.split('_')[1][:-5]\n            epoch, acci = metrics.split('-')\n            epoch = int(epoch.split('=')[1])\n            acci = float(acci.split('=')[1])\n            sorted_files.append((acci, epoch, file))\n        sorted_files = sorted(sorted_files)\n        return [x[-1] for x in sorted_files]\n\n\n    def get_checkpoint_path(name, filename, fold, folder, th_id=34):\n        \"\"\" gets best checkpoint path but ignores differences <= `min_delta` and prefers\n            models with more epochs if their model id >= `th_id` (to avoid changing\n            past models).\n        \"\"\"\n        path = f'{folder}/{filename}.ckpt'\n        if not os.path.exists(path):\n            files = glob.glob(f'{folder}/{filename}_f{fold}*.ckpt')\n            scores = [float(x.split('=')[-1][:-5].split('-')[0]) if 'val_auc' in x or 'acc' in x else\n                      -np.inf for x in files]\n            id_ = int(name.split('.')[-1])\n            if id_ >= th_id:\n                epochs = [float(x.split('=')[1].split('-')[0]) for x in files]\n                scores = [score + epoch * 1e-7 for score, epoch in zip(scores, epochs)]\n            if len(scores) > 0:\n                idx = np.argsort(scores)[-1]\n                path = files[idx]\n            else:\n                print(f\"No file found for '{folder}/{filename}_f{fold}' -> {files}\")\n                path = None\n        return path\n        \n\n    def load_model_state(model, chkp_path):\n        \"\"\" load checkpoint into model \"\"\"\n        chkp = torch.load(chkp_path, map_location=torch.device('cpu'))\n        model_state = {k[6:]: v for k, v in chkp['state_dict'].items()}  # to remove initial 'model.'\n        model.load_state_dict(model_state)\n\n\n    def average_model_state(chkp_paths, save=None):\n        \"\"\" average various checkpoints \"\"\"\n        chkp = torch.load(chkp_paths[0], map_location=torch.device('cpu'))\n        for chkp_path in chkp_paths[1:]:\n            temp_chk = torch.load(chkp_path, map_location=torch.device('cpu'))\n            for k in chkp['state_dict']:\n                if isinstance(chkp['state_dict'][k], dict):\n                    for k2 in chkp['state_dict'][k]:\n                        chkp[k][k2] += temp_chk[k][k2]\n                else:\n                    chkp[k] += temp_chk[k]\n        for k in chkp:\n            chkp[k] /= len(chkp_paths)\n        if save is not None:\n            torch.save(chkp)\n        return chkp\n                \n\n    def load_model_state_average(model, chkp_paths):\n        \"\"\" load checkpoints, average their weights and load into model \"\"\"\n        chkp = average_model_state(chkp_paths)\n        model_state = {k[6:]: v for k, v in chkp['state_dict'].items()}  # to remove initial 'model.'\n        model.load_state_dict(model_state)\n        \n\n    def load_encoder_state(encoder, chkp_path):\n        # chkp_path = f'{ROOT}/b5-train-cola-1-0-0/last_train_f0.ckpt'\n        try:\n            chkp = torch.load(chkp_path)\n        except:\n            chkp = torch.load(chkp_path, map_location=torch.device('cpu'))\n        \n        if 'state_dict' in chkp:\n            encoder_state = {k[6:]: v for k, v in chkp['state_dict'].items()}\n            encoder_state = {k[6:]: v for k, v in encoder_state.items() if k[:6] == 'model.'}\n        else:\n            encoder_state = chkp.copy()\n            for k, v in encoder_state.items():\n                if '.running_mean' in k and torch.isnan(v).any():\n                    encoder_state[k] = torch.ones_like(v, dtype=torch.float32) * 0\n                if '.running_var' in k and torch.isnan(v).any():\n                    encoder_state[k] = torch.ones_like(v, dtype=torch.float32)\n\n        encoder_state = {\n             k.replace('backbone.0', 'stem').replace('backbone.1', 'stages')\n              .replace('backbone.2', 'final_conv').replace('stem_', 'stem.')\n              .replace('stages_', 'stages.').replace('global_pool.', ''): v \n                 for k, v in encoder_state.items()}\n        encoder.load_state_dict(encoder_state, strict=STRICT)  # for Vladimir pretrained encoders\n        print(f\"loaded encoder '{chkp_path}'\")\n        return encoder_state\n        \n\n    def create_models(models_config, half=True, verbose=False):\n        models = dict()\n        for name, folds, filename, folder, config in models_config:\n            if not os.path.exists(folder):\n                print(f\"folder '{folder}' not found\")\n                continue\n            for fold in folds:\n                chkp_path = get_checkpoint_path(name, filename, fold, folder)\n                if chkp_path is None:\n                    continue\n                if verbose:\n                    print(f\"{name} f{fold} -> '{chkp_path}'\")\n                cls = globals()[config.model_name] if config.model_name else SpecNetImg\n                in_chans = config.in_chans or 3\n                out_indices = config.out_indices or None\n                model = cls(config, resize=config.resize, add_position=config.add_position,\n                            out_indices=out_indices, pretrained=False, in_chans=in_chans)\n                load_model_state(model, chkp_path)\n                if half:\n                    model = model.half().float()\n                model.eval()\n\n                models[f'{name}_f{fold}'] = model\n        print(f'models: {list(models.keys())}')\n        return models\n\n\n    def save_models_jit(models_config):\n        models = create_models(models_config)\n        for name, model in models.items():\n            model_jit = torch.jit.script(model)\n            torch.jit.save(model_jit, f'{name}.pt')\n            \n\n    def load_models_jit(models_config):\n        models = dict()\n        for name, folds, filename, folder, config in models_config:\n            for fold in folds:\n                name_fold = f'{name}_f{fold}'\n                model = torch.jit.load(f'./{name_fold}.pt')\n                model.eval()\n                # model = torch.jit.optimize_for_inference(torch.jit.script(model.eval()))\n                models[name_fold] = model\n        return models\n\n\n    def load_models_onnx(models_config):\n        models = dict()\n        for name, folds, filename, folder, config in models_config:\n            for fold in folds:\n                filename = f'{name}_f{fold}'\n                onnx_model = onnx.load(f'{folder}/{name}/{filename}.onnx')\n                # onnx_model_graph = onnx_model.graph\n                onnx_session = ort.InferenceSession(onnx_model.SerializeToString())\n                models[filename] = onnx_session\n        print(f'models: {list(models.keys())}')\n        return models\n\n\n    def load_model_onnx(name, folder, fold):\n        filename = f'{name}_f{fold}'\n        onnx_model = onnx.load(f'{folder}/{name}/{filename}.onnx')\n        # onnx_model_graph = onnx_model.graph\n        onnx_session = ort.InferenceSession(onnx_model.SerializeToString())\n        return onnx_session\n\n\n    def load_models_vino(models_config, verbose=1):\n        models = dict()\n        for name, folds, filename, folder, config in models_config:\n            for fold in folds:\n                filename = f'{name}_f{fold}'\n                ir_path = f'{ROOT}/b5-models-1-2/{name}/{filename}.xml'\n                core = ovrt.Core()\n                model = core.read_model(model=ir_path)\n                compiled_model = core.compile_model(\n                    model=model, device_name='CPU',\n                    config={ovrt.properties.inference_num_threads(): 4,\n                            ovrt.properties.hint.enable_hyper_threading(): True,\n                            ovrt.properties.hint.enable_cpu_pinning(): False})\n                # compiled_model = compiled_model.create_infer_request()\n                models[filename] = compiled_model\n        if verbose:\n            print(f'models: {list(models.keys())}')\n        return models\n\n\n    def reshape_spec(spec):\n        \"\"\" converts batch, chunk, heigh, width to BHW \"\"\"\n        B, C, H, W = spec.size()\n        return spec.reshape((B * C, H, W))        \n\n\n    def predict_fn_single(ds_test, models):\n        \"\"\" run prediction for `models` using dataset `ds_test`. returns a list[models, batch]\n        \"\"\"\n        dl_test = DataLoader(ds_test, batch_size=ARGS.loader_batch_size, num_workers=ARGS.num_workers,\n                             # prefetch_factor=ARGS.prefetch_factor, \n                             drop_last=False,\n                             pin_memory=True,\n                             # multiprocessing_context=get_context('loky')\n                            )\n        autocast_ = ARGS.autocast\n        range_ = tqdm(dl_test) if ARGS.verbose else dl_test\n        batch_size = ARGS.batch_size\n        preds = []\n        for spec in range_:\n            #if SPEC_128_512_5 in spec_types:\n            B, C, H, W = spec.size()\n            n_batch = B * C / batch_size\n            spec = spec.reshape((int(n_batch), batch_size, 1, H, W))\n            # assert n_batch == int(n_batch)\n\n            # spec = specs[SPEC_224_224B_10]\n            # B, C, H, W = spec.size()\n            # n_batch = B * C / batch_size\n            # spec_224_224b_10 = spec.reshape((int(n_batch), batch_size, 1, H, W))\n            \n            models_preds = []\n            for name, model in models.items():\n                batch_out = []\n                for batch in spec:\n                    with torch.no_grad():\n                        out = model(batch)\n                    batch_out.append(out)\n                batch_out = torch.cat(batch_out, 0)\n                #if n_repeat:\n                #    batch_out = batch_out.repeat_interleave(n_repeat, dim=1)\n                models_preds.append(batch_out)\n            preds.append(torch.stack(models_preds))\n        return preds\n\n\n    def predict_fn_regular(ds_test, models):\n        \"\"\" run prediction for `models` using dataset `ds_test`. returns a list[models, batch]\n        \"\"\"\n        dl_test = DataLoader(ds_test, batch_size=ARGS.loader_batch_size, num_workers=ARGS.num_workers,\n                             # prefetch_factor=ARGS.prefetch_factor, \n                             drop_last=False,\n                             pin_memory=True,\n                             # multiprocessing_context=get_context('loky')\n                            )\n        autocast_ = ARGS.autocast\n        range_ = tqdm(dl_test) if ARGS.verbose else dl_test\n        batch_size = ARGS.batch_size\n        preds = []\n        for spec in range_:\n            #if SPEC_128_512_5 in spec_types:\n            B, C, H, W = spec.size()\n            n_batch = B * C / batch_size\n            spec = spec.reshape((int(n_batch), batch_size, 1, H, W))\n            # assert n_batch == int(n_batch)\n\n            # spec = specs[SPEC_224_224B_10]\n            # B, C, H, W = spec.size()\n            # n_batch = B * C / batch_size\n            # spec_224_224b_10 = spec.reshape((int(n_batch), batch_size, 1, H, W))\n            \n            models_preds = []\n            for name, model in models.items():\n                batch_out = []\n                for batch in spec:\n                    with torch.no_grad():\n                        out = model(batch)\n                    batch_out.append(out)\n                batch_out = torch.cat(batch_out, 0)\n                #if n_repeat:\n                #    batch_out = batch_out.repeat_interleave(n_repeat, dim=1)\n                models_preds.append(batch_out)\n            preds.append(torch.stack(models_preds))\n        return preds\n\n\n    def predict_fn_opt(specs):\n        \"\"\" run prediction for specs using opt models. \"\"\"\n        out = []\n        specs = specs.numpy()\n        n = math.ceil(specs.shape[0] / bs)\n        for model in models.values():\n            out_model = []\n            for i in tqdm(range(n)):\n                y = model.run(['y'], {'x': specs[i*bs : (i+1)*bs]})[0]\n                out_model.append(y)\n            out_model = np.concatenate(out_model, axis=0)\n            out.append(out_model)\n        out = np.stack(out, axis=0)\n        return out\n\n\n    def predict_fn_onnx(specs, verbose=False):\n        \"\"\" run prediction for specs using onnx models. \"\"\"\n        batch_size = ARGS.batch_size\n        out = []\n        specs = specs.numpy()\n        n = math.ceil(specs.shape[0] / batch_size)\n        for model in models:\n            out_model = []\n            range_ = tqdm(range(n)) if verbose else range(n)\n            for i in range_:\n                y = model.run(['y'], {'x': specs[i*batch_size : (i+1)*batch_size]})[0]\n                out_model.append(y)\n            out_model = np.concatenate(out_model, axis=0)\n            out.append(out_model)\n        out = np.stack(out, axis=0)\n        return out\n\n\n    def predict_fn_onnx_timeout(specs):\n        \"\"\" run prediction for specs using onnx models. \"\"\"\n        deadline = DEADLINE\n        stop = False\n        batch_size = ARGS.batch_size\n        out = []\n        specs = specs.numpy()\n        n = math.ceil(specs.shape[0] / batch_size)\n        \n        for model in models:\n            out_model = []\n            for i in tqdm(range(n)):\n                if time.time() >= deadline:\n                    stop = True\n                    print(f'STOP: deadline breached {time.time() - deadline:,.0f}')\n                    break\n                y = model.run(['y'], {'x': specs[i*batch_size : (i+1)*batch_size]})[0]\n                out_model.append(y)\n            out_model = np.concatenate(out_model, axis=0)\n            out.append(out_model)\n            if stop:\n                gap = len(specs) - len(out[-1])\n                if gap > 0:\n                    out[-1] = np.pad(out[-1], pad_width=((0, gap), (0, 0)), mode='constant', constant_values=np.nan)\n                break\n        out = np.stack(out, axis=0)\n        return out\n\n\n    def predict_fn_onnx_timeout_bits(specs):\n        \"\"\" run prediction for specs using onnx models. \"\"\"\n        deadline = DEADLINE\n        stop = False\n        batch_size = ARGS.batch_size\n        out = []\n        specs = specs.numpy()\n        n = math.ceil(specs.shape[0] / batch_size)\n        \n        for model, bits in zip(models, all_bits):\n            out_model = []\n            for i in tqdm(range(n)):\n                if time.time() >= deadline:\n                    stop = True\n                    print(f'STOP: deadline breached {time.time() - deadline:,.0f}')\n                    break\n                x = specs[i*batch_size : (i+1)*batch_size]\n                if bits is None or bits == 8:\n                    x = ((x * 255).astype(np.uint8) / 255).astype(np.float32)\n                elif bits == 16:\n                    x = ((x * 65535).astype(np.uint32) / 65535).astype(np.float32)\n                y = model.run(['y'], {'x': x})[0]\n                out_model.append(y)\n            out_model = np.concatenate(out_model, axis=0)\n            out.append(out_model)\n            if stop:\n                gap = len(specs) - len(out[-1])\n                if gap > 0:\n                    out[-1] = np.pad(out[-1], pad_width=((0, gap), (0, 0)), mode='constant', constant_values=np.nan)\n                break\n        out = np.stack(out, axis=0)\n        return out\n        \n\n    def predict_fn_vino(specs):\n        \"\"\" run prediction for specs using vino models. \"\"\"\n        models = load_models_vino(models_config, verbose=0)\n        batch_size = ARGS.batch_size\n        out = []\n        specs = specs.numpy()\n        n = math.ceil(specs.shape[0] / batch_size)\n        for model in models.values():\n            out_model = []\n            for i in tqdm(range(n)):\n                y = model([specs[i*batch_size : (i+1)*batch_size]])['y']\n                # y = model.infer(inputs=[specs[i*batch_size : (i+1)*batch_size]])['y']\n                out_model.append(y)\n            out_model = np.concatenate(out_model, axis=0)\n            out.append(out_model)\n        out = np.stack(out, axis=0)\n        return out\n\n\n    def prepare_predict():\n        \"\"\" create models \"\"\"\n        if ARGS.engine == 'opt':\n            models = load_models_jit(models_config)\n            predict_fn = predict_fn_regular\n        elif ARGS.engine == 'onnx':\n            models = load_models_onnx(models_config)\n            predict_fn = predict_fn_onnx\n        elif ARGS.engine == 'vino':\n            models = None # load_models_vino(models_config)\n            predict_fn = predict_fn_vino\n        return models, predict_fn\n\n\n    def split_spec_configs(ws, models_config_dict):\n        \"\"\" split models based on the melspec parameters \"\"\"\n        spec_configs = dict()\n        for name in ws.keys():\n            cfg = models_config_dict[name.split('_')[0]]\n            spec = (cfg.spec, cfg.img_dim, cfg.img_duration, cfg.spec_norm, cfg.snipet)\n            if spec not in spec_configs:\n                spec_configs[spec] = [name]\n            else:\n                spec_configs[spec].append(name)\n        return spec_configs\n\n    # ----------------------------------------------------------------------------------------------------\n    # 0.13 end\n    # ----------------------------------------------------------------------------------------------------\n\nprint('done')","metadata":{"trusted":true,"jupyter":{"source_hidden":true},"_kg_hide-input":true,"execution":{"iopub.status.busy":"2025-06-03T01:30:04.038397Z","iopub.status.idle":"2025-06-03T01:30:04.038694Z","shell.execute_reply.started":"2025-06-03T01:30:04.038565Z","shell.execute_reply":"2025-06-03T01:30:04.038578Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"if 'SKIP' not in globals() or not SKIP:     # features\n    print('features v5.0')\n    \n    # 5.1: build unlabeled dataset to be loaded in memory\n    # 5.0: split pseudo in folds\n    # 4.4: ROOT path\n    # 4.3: google drive functions\n    # 4.2: ensemble of pseudo labels\n    # 4.1: fix pseudo label using Volodymyr's sampling\n    # 4.0: adjust read to support pseudo train audios with soft pseudo labels\n    # 3.5: unlabeled with soft pseudo labels\n    # 3.4: read function\n    # 3.3: flag for manual curation\n    # 3.2: remove unused code\n    # 3.1: additional files\n    # 3.0: h5 data + curated v2\n    # 2.5: enhanced filter speech\n    # 2.4: filter speech\n    # 2.3: speech filter\n    # 2.2: update ensemble\n    # 2.1: update methods\n    # 2.0: initial\n\n    from typing import Any, Dict, List, Optional, Sequence, Tuple, Union, NamedTuple\n    import copy\n    from contextlib import contextmanager\n    import warnings\n    warnings.filterwarnings('ignore')\n    \n    import os\n    import sys\n    import gc\n    import time\n    import torch\n    import joblib\n    import pickle\n    import psutil\n    import ast\n    import copy\n    import glob\n    import numpy as np\n    import pandas as pd\n    from tqdm.autonotebook import tqdm\n    import multiprocessing as mp\n    import concurrent.futures\n    from itertools import groupby as grp_by\n    from datetime import datetime, timedelta\n    import h5py\n\n    from itertools import permutations\n    import math\n    import heapq\n    import re\n\n    # import sklearn\n    from sklearn.model_selection import KFold, GroupKFold, GroupShuffleSplit, StratifiedKFold, StratifiedGroupKFold\n    from sklearn.model_selection._split import _BaseKFold\n    from sklearn.utils.validation import _num_samples, check_array\n    from sklearn.utils.multiclass import type_of_target\n    from sklearn.utils import check_random_state\n    # from sklearn.model_selection._split import _BaseKFold, indexable\n    \n    try:\n        import polars as pl\n    except:\n        pass\n        # !pip install polars\n        # import polars\n\n    try:\n        import pyarrow.parquet as pq\n    except:\n        pass\n        # !pip install pyarrow\n        # import pyarrow.parquet as pq\n\n    try:\n        import optuna\n    except:\n        pass\n    \n    from numba import njit\n\n\n        # -------------------\n        # plot\n        # -------------------\n\n    import matplotlib\n    import matplotlib.pyplot as plt\n    from IPython.display import Audio\n    try:\n        import seaborn as sns\n        import plotly\n        from plotly.offline import iplot\n        from plotly.subplots import make_subplots\n        import plotly.express as px\n        import plotly.graph_objects as go\n\n        import ipywidgets as widgets\n        from ipywidgets import Layout\n\n        # the following 2 lines resolve issue with charts not showing\n        import plotly.io as pio\n        pio.renderers.default = \"iframe\"\n    except:\n        pass\n\n    try:\n        # pydot-ng is a fork of pydot that is better maintained.\n        import pydot_ng as pydot\n    except ImportError:\n        # pydotplus is an improved version of pydot\n        try:\n            import pydotplus as pydot\n        except ImportError:\n            # Fall back on pydot if necessary.\n            try:\n                import pydot\n            except ImportError:\n                pydot = None\n\n    try:\n        import cupy, cudf, cuml\n        from cuml.linear_model import LinearRegression\n        from cuml.neighbors import NearestNeighbors, KNeighborsRegressor, KNeighborsClassifier\n        CUDF = True\n    except:\n        try:\n            from sklearn.linear_model import LinearRegression\n            from sklearn.neighbors import NearestNeighbors, KNeighborsRegressor, KNeighborsClassifier\n        except:\n            pass\n        CUDF = False\n    from sklearn.cluster import KMeans\n\n    # scipy\n    from scipy.stats import linregress\n    from scipy.interpolate import interp1d\n    from scipy import signal\n    from scipy.signal import argrelmax\n    from scipy.signal import spectrogram\n    from scipy.signal import find_peaks\n    from scipy.signal import butter, lfilter\n    # import scipy.fft as fft\n    from scipy import optimize\n    \n    # sklearn\n    try:\n        from sklearn import linear_model\n        from sklearn.preprocessing import normalize, QuantileTransformer, MinMaxScaler, RobustScaler, normalize, minmax_scale\n        from sklearn.model_selection import train_test_split\n        from sklearn.metrics import make_scorer, \\\n                                    mean_squared_error, mean_absolute_error, r2_score, \\\n                                    accuracy_score, classification_report, precision_score, recall_score, f1_score, \\\n                                    roc_auc_score, precision_recall_curve, roc_curve, auc, average_precision_score, confusion_matrix\n        from sklearn.decomposition import PCA, TruncatedSVD\n    except:\n        pass\n\n    # skspatial\n    # from skspatial.objects import Line, Points\n\n    import pickle\n    import random\n    import pdb\n    import warnings\n    import math\n    warnings.filterwarnings('ignore')\n\n    # competition\n    try:\n        import pywt\n    except:\n        pass\n    import librosa\n    \n    COLAB = \"google.colab\" in sys.modules\n\n    # ----------------------------------------------------------------------------------------------------\n    # 0.1 utils\n    # ----------------------------------------------------------------------------------------------------\n\n    import pickle\n\n    def save_obj(obj, name, protocol=4): # pickle.HIGHEST_PROTOCOL):\n        name = name.replace('.pkl', '')\n        with open('./'+ name + '.pkl', 'wb') as f:\n            pickle.dump(obj, f, protocol)\n\n    def load_obj(name, folder=''):\n        name = name.replace('.pkl', '')\n        with open(folder + name + '.pkl', 'rb') as f:\n            return pickle.load(f)\n\n    def save_obj2(obj, name):\n        joblib.dump(obj, './'+ name + '.jl')\n\n    def load_obj2(name, folder=''):\n        name = name.replace('.jl', '')\n        return joblib.load(folder + name + '.jl')\n\n\n    @contextmanager\n    def timer(name=''):\n        s = time.time()\n        mem_before = psutil.virtual_memory().used/1024/1024/1024\n        yield\n        elapsed = time.time() - s\n        mem_after = psutil.virtual_memory().used/1024/1024/1024\n        print(f'[{name}] {elapsed: .3f} s   {mem_before: .2f} GB -> {mem_after: .2f} GB')\n\n\n    def split_letters_numbers(text):\n        if len(text) == 0:\n            return text\n\n        if text[0].isalpha():\n            # e.g. 'eb13'\n            for index, char in enumerate(text):\n                if char.isdigit():\n                    return text[:index], text[index:]\n            return text, ''\n        else:\n            # e.g. '19a'\n            for index, char in enumerate(text):\n                if char.isalpha():\n                    return text[:index], text[index:]\n            return text, ''\n        \n    \n    def sort_names(names):\n        \"\"\" sort model names \"\"\"\n        sortable_names = []\n        for name in names:\n            type, id_ = name.split('.')\n            # split letters from numbers in type\n            letters, numbers = split_letters_numbers(type)\n            type = f'{letters:<6}{numbers:>6}'\n            numbers, letters = split_letters_numbers(id_)\n            id_ = f'{numbers:>6}{letters:<6}'\n            sortable_names.append(f'{type}.{id_}')\n            \n        sorted_pairs = sorted(zip(sortable_names, names))\n        return [x[1] for x in sorted_pairs]\n\n\n    def argsort_names(names):\n        \"\"\" sort model names \"\"\"\n        sortable_names = []\n        for name in names:\n            type, id_ = name.split('.')\n            # split letters from numbers in type\n            letters, numbers = split_letters_numbers(type)\n            type = f'{letters:<6}{numbers:>6}'\n            numbers, letters = split_letters_numbers(id_)\n            id_ = f'{numbers:>6}{letters:<6}'\n            sortable_names.append(f'{type}.{id_}')\n            \n        idx = np.argsort(sortable_names)\n        return idx\n\n\n    def list_match(str_, list_):\n        for x in list_:\n            if str_ in x:\n                return True\n        return False\n\n\n    def wait():\n        for i in range(1000):\n            print('.', end='')\n            time.sleep(60)\n\n\n    def is_interactive():\n        if COLAB:\n            return hasattr(sys, 'ps1')\n        else:\n            return os.environ['KAGGLE_KERNEL_RUN_TYPE'] == 'Interactive'\n        \n\n    def sort_external(data, order, reverse=False, return_order=False):\n        \"\"\" sort list `data` based on the content of list `order` \"\"\"\n        merged = list(zip(data, order))\n        sorted_merged = sorted(merged, key=lambda x: x[1], reverse=reverse)\n        if return_order:\n            return [x[0] for x in sorted_merged], [x[1] for x in sorted_merged]\n        else:\n            return [x[0] for x in sorted_merged]\n        \n\n    def filter_list(l, filter):\n        \"\"\" filter list l using filter (either a list of bool or indices) \"\"\"\n        if isinstance(filter, bool):\n            return [x for x, idx in zip(l, filter) if idx]\n        return [x for i, x in enumerate(l) if i in filter]\n\n\n    def apply_temperature(probs, temp):\n        \"\"\" apply temperature to probs:\n            temp = 2: smothes probs and tends to also pick less likely options (more exploratory)\n            temp = 1: same\n            temp = .5: sharpens probs and tends to pick most probable (more deterministic)\n        \"\"\"\n        scaled_probs = np.exp(np.log(probs) / temp)\n        return scaled_probs / np.sum(scaled_probs)\n\n\n    def apply_temperature_log(log_probs, temp):\n        \"\"\" apply temperature to log of probs:\n            temp = 2: smothes probs and tends to also pick less likely options (more exploratory)\n            temp = 1: same\n            temp = .5: sharpens probs and tends to pick most probable (more deterministic)\n        \"\"\"\n        if not isinstance(log_probs, np.ndarray):\n            log_probs = np.array(log_probs)\n        scaled_probs = np.exp(log_probs / temp)\n        return scaled_probs / np.sum(scaled_probs)    \n\n\n    def seed_python(seed=1234):\n        random.seed(seed)\n        os.environ['PYTHONHASHSEED'] = str(seed)\n        np.random.seed(seed)\n\n\n    def mySystem(cmd):\n        \"\"\" run shell command from python \"\"\"\n        print(cmd)\n        process = Popen(cmd, stdout=PIPE, stderr=STDOUT, shell=True)\n        for line in iter(process.stdout.readline, b''):\n            print(line.decode(\"utf-8\"), end='')\n        assert(process.wait() == 0)\n\n\n    def fg_format_seconds(seconds):\n        seconds = int(seconds)\n        minutes = seconds // 60\n        remaining_seconds = seconds % 60\n        return f\"{minutes:02}:{remaining_seconds:02}\"\n\n\n    def random_choice(range_, k):\n        \"\"\" faster version of random choice without repetition \"\"\"\n        while True:\n            choice = random.choices(range_, k=k)\n            if len(set(choice)) == k:\n                return choice\n\n\n    def gaussian_kernel(l=5, sig=1.):\n        \"\"\" create 1D gaussian kernel with side length `l` and a sigma of `sig` \"\"\"\n        ax = np.linspace(-(l - 1) / 2., (l - 1) / 2., l)\n        gauss = np.exp(-0.5 * np.square(ax) / np.square(sig))\n        return gauss / np.sum(gauss)\n\n\n    def pad_tuples(tuples, pad_value=-1):\n        \"\"\"\n        Pads a list of tuples to the size of the longest tuple and returns the padded list,\n        with tuples converted to lists, along with the original sizes.\n\n        Args:\n            tuples (list of tuples): The list of tuples to pad.\n            pad_value (any): The value used for padding.\n\n        Returns:\n            tuple: A tuple containing:\n                - padded_list (list of lists): The padded tuples as lists.\n                - original_sizes (list of int): The original sizes of the tuples.\n        \"\"\"\n        # Get the size of the longest tuple\n        max_size = max(len(t) for t in tuples)\n        \n        # Create the padded list and record original sizes\n        padded_list = [list(t) + [pad_value] * (max_size - len(t)) for t in tuples]\n        original_sizes = [len(t) for t in tuples]\n        return padded_list, original_sizes\n\n\n    def first_smaller_index(lst, th):\n        return next((i for i, x in enumerate(lst) if x < th), None)\n\n\n    def get_scaler(type='quantile'):\n         if type == 'std':\n            return StandardScaler()\n         if type == 'minmax':\n            return MinMaxScaler()\n         if type == 'max':\n            return MaxAbsScaler() \n         if type == 'robust':\n            return RobustScaler()  # iqr\n         if type == 'norm':\n            return Normalizer()\n         if type == 'quantile':\n            return QuantileTransformer(n_quantiles=100, random_state=0, output_distribution='normal')\n         if type == 'power':\n            return PowerTransformer()\n\n\n    def logit(p):\n        return np.log(p) - np.log(1 - p)\n\n\n    def select_uncorrelated_indices(corr_matrix, th_corr=.95):\n        \"\"\"\n        Selects indices starting from 0 and adds indices one by one \n        if their correlation with all previously selected indices is below th_corr.\n\n        :param corr_matrix: 2D numpy array (correlation matrix)\n        :param th_corr: float (correlation threshold)\n        :return: List of selected indices\n        \"\"\"\n        n = corr_matrix.shape[0]\n        selected = [0]  # Start with the first index\n\n        for i in range(1, n):\n            if all(abs(corr_matrix[i, j]) < th_corr for j in selected):\n                selected.append(i)\n\n        return selected\n\n\n    def split_list(lst, n):\n        k, m = divmod(len(lst), n)\n        return [lst[i*k + min(i, m):(i+1)*k + min(i+1, m)] for i in range(n)]\n\n\n    def split_balanced_by_species(files, n):\n        # Group by species\n        species_groups = defaultdict(list)\n        for file in files:\n            species = file.strip('/').split('/')[-2]\n            species_groups[species].append(file)\n        \n        # sort species groups by descending size (for better balancing)\n        species_group_list = sorted(species_groups.values(), key=len, reverse=True)\n        \n        # Distribute into n lists using a min-heap based on list size\n        result = [[] for _ in range(n)]\n        heap = [(0, i) for i in range(n)]  # (size, index)\n        heapq.heapify(heap)\n\n        for group in species_group_list:\n            size, idx = heapq.heappop(heap)\n            result[idx].extend(group)\n            heapq.heappush(heap, (size + len(group), idx))\n\n        return result\n           \n\n    def get_subfolders_size(root='/content/drive/MyDrive/b5'):\n        size = dict()\n        for subfolder in os.listdir(root):\n            subfolder_path = os.path.join(root, subfolder)\n            if os.path.isdir(subfolder_path):\n                total_size = 0\n                for dirpath, _, filenames in os.walk(subfolder_path):\n                    for f in filenames:\n                        fp = os.path.join(dirpath, f)\n                        if os.path.isfile(fp):\n                            total_size += os.path.getsize(fp)\n                print(f'{subfolder}: {total_size/1e9:,.2f} GB')\n                size[subfolder] = total_size\n        size = pd.Series(size)\n        size.to_csv('folder sizes.csv')\n        print(f'All: {size.sum()/1e9:,.2f} GB')\n           \n    # torch\n\n    def print_num_parameters(model, note=None):\n        n = sum(p.numel() for p in model.parameters())\n        if note:\n            print(f'{note:<30}: {n:,.0f}')\n        return n\n\n    \n    def print_weights(model):\n        \"\"\" prints model weights and returns them as a dataframe \"\"\"\n        items = []\n        for i, (name, module) in enumerate(model.named_modules()):\n            weight = module.weight.shape if hasattr(module, 'weight') else ''\n            items.append((name, f'{module.__class__.__module__}.{module.__class__.__name__}', weight))\n\n        res = pd.DataFrame(items, columns=['name', 'class', 'weight'])\n        res.to_csv('x3d_l.csv')\n        pd.set_option('display.max_rows', None)\n        display(res)\n        return res\n\n    # ----------------------------------------------------------------------------------------------------\n    # 0.2 target\n    # ----------------------------------------------------------------------------------------------------\n\n    def get_pseudo_labels(coef, folder=f'/kaggle/input/ba-models-2/pseudo'):\n        \"\"\" load and average pseuso labels \"\"\"\n        if isinstance(coef, list):\n            coef = {name: 1 for name in coef}\n        names = list(coef.keys())\n        oof = []\n        for name in names:\n            file = f'{folder}/{name}/oof {name}.pkl'\n            if not os.path.exists(file):\n                file = f'{folder}/oof {name}.pkl'\n            oof.append(load_obj(file))\n        oof = np.stack(oof, axis=0)\n        pseudo_labels = ensemble_oof(coef, oof, names)\n        return pseudo_labels\n\n\n    def expand_labels(labels, sizes):\n        \"\"\" expand labels per file such that each label is repeated `sizes` \"\"\"\n        expanded = []\n        for l, n in zip(labels, sizes):\n            expanded += [l] * n\n        return np.stack(expanded, axis=0)\n\n    # ----------------------------------------------------------------------------------------------------\n    # 0.3 meta\n    # ----------------------------------------------------------------------------------------------------\n\n    def get_audio_dims(files):\n        def _process(file):\n            audio, sr = torchaudio.load(file)\n            len_ = audio.shape[1] / SR\n            return len_\n        \n        out = Parallel(n_jobs=-1)(delayed(_process)(file) for file in tqdm(files))\n        return out\n\n\n    def print_dup():\n        metak = meta.set_index('filename').to_dict(orient='index')\n        dup = []\n        for a, b in DUPLICATED_FILES:\n            l1, l2 = a.split('/')[0], b.split('/')[0]\n            is_same = l1 == l2\n            dup.append((*metak[a].values(), *metak[b].values(), is_same))\n        dup = pd.DataFrame(dup, columns=['lat_a', 'long_a', 'label_a', 'lat_b', 'long_b', 'label_b', 'same'])\n        dup.to_csv('duplicates')\n\n\n    def read_meta(split_author=True, pivot_folds=False, author_to_num=True):\n        file = f'{ROOT}/birdclef-2025/train.csv'\n        meta = pd.read_csv(file).sort_values(['filename'])\n        meta['i'] = np.arange(meta.shape[0])\n        authors = meta.author.unique()\n        if author_to_num:\n            author_to_num = dict(zip(authors, range(len(authors))))\n            meta['author'] = meta.author.map(author_to_num)\n        \n        # split author in case there is only one per primary_label\n        if split_author:\n            meta['count_unique_author'] = meta.groupby('primary_label')['author'].transform('nunique')\n            meta['order'] = meta.groupby('primary_label').cumcount()\n            meta['author'] = np.where(meta.count_unique_author == 1, meta.author + meta.order / 1000, meta.author)\n\n        meta = meta[['i', 'primary_label', 'secondary_labels', 'author', 'latitude', 'longitude', 'rating', 'filename']]\n        \n        # merge labels\n        labels = []\n        for row in meta.to_dict(orient='records'):\n            row = dotdict(row)\n            secondary = row.secondary_labels[1:-1].replace(\"'\", '')\n            secondary = secondary.split(',') if len(secondary) > 0 else []\n            secondary = [x.replace(' ', '') for x in secondary]\n            labels.append([row.primary_label] + secondary)\n        meta['labels'] = labels\n        meta['n_labels'] = [len(x) for x in labels]\n        meta = meta.drop(['secondary_labels'], axis='columns')\n        \n        # save pivot of folds\n        if pivot_folds:\n            meta.pivot_table(index='primary_label', columns='fold', values='filename',\n                             aggfunc='count').to_csv('folds.csv')\n        return meta\n        \n\n    def build_previus_meta():\n        \"\"\" build meta for previous files. \"\"\"\n        previous_lens = pd.read_csv(f'{ROOT}/b5-data-2021-to-2025/len_previous_audios')\n        previous = pd.read_csv(f'{ROOT}/b5-data-2021-to-2025/previous.csv')\n        cols = ['primary_label', 'secondary_labels', 'record_name', 'author', 'latitude', 'longitude']\n        meta = previous[cols]\n        filenames = [f'{primary_label}/{record_name}' for primary_label, record_name in \\\n                         zip(meta.primary_label, meta.record_name)]\n        meta['filename'] = filenames\n        files_exist = [os.path.exists(f'{ROOT}/b5-data-2021-to-2025/{x}') for x in tqdm(filenames)]\n        meta = meta[files_exist]\n        \n        # merge labels\n        labels = []\n        for row in meta.to_dict(orient='records'):\n            row = dotdict(row)\n            secondary = row.secondary_labels[1:-1].replace(\"'\", '')\n            secondary = secondary.split(',') if len(secondary) > 0 else []\n            secondary = [x.replace(' ', '') for x in secondary]\n            labels.append([row.primary_label] + secondary)\n        meta['labels'] = labels\n        meta['n_labels'] = [len(x) for x in labels]\n        \n        if 'rating' not in meta.columns:\n            meta['rating'] = 5\n        meta['i'] = np.arange(50_000, 50_000 + len(meta))\n        meta['fold'] = -2\n        meta = meta.merge(previous_lens, on='filename', how='left')\n        meta['filename'] = meta.filename.astype(str) + '.ogg'\n        \n        meta_cols = ['i', 'primary_label', 'author', 'latitude', 'longitude', 'rating',\n               'filename', 'labels', 'n_labels', 'dim', 'fold']\n        meta = meta[meta_cols]\n        return meta\n        \n\n    def build_meta_add(folder_base, previous_filenames):\n        meta_add = pd.read_csv(f'{folder_base}/train_and_prev_comps_extendedv1_pruneSL_XConly2025_snipet28032025_hdf5.csv')\n        filter_new = (meta_add.dataset != 'comp_2025') & ~meta_add.filename.isin(previous_filenames)\n        meta_add = meta_add[filter_new]\n        meta_add['fold'] = -2\n        meta_add['dim'] = meta_add.duration_s\n        \n        # merge labels and build list of files\n        labels, files = [], []\n        for row in meta_add.to_dict(orient='records'):\n            row = dotdict(row)\n            secondary = row.secondary_labels[1:-1].replace(\"'\", '')\n            secondary = secondary.split(',') if len(secondary) > 0 else []\n            secondary = [x.replace(' ', '') for x in secondary]\n            secondary = [x for x in secondary if x in BIRDS]\n            labels.append([row.primary_label] + secondary)\n            files.append(f'{folder_base}/{row.data_root_id}/{row.data_root_id}/{row.filename}')\n        \n        meta_add['filename'] = [x.replace('.mp3', '.ogg') for x in meta_add.filename.values]\n        meta_add['labels'] = labels\n        meta_add['n_labels'] = [len(x) for x in labels]\n        cols = ['primary_label', 'author', 'latitude', 'longitude', 'rating', 'filename', 'labels', 'n_labels', 'dim', 'fold']\n        meta_add = meta_add[cols]\n        return meta_add, files\n\n\n    def get_series_mode(x):\n        counts = x.value_counts()\n        return counts.index[0] if len(counts) > 0 else None\n\n\n    def add_geo_label(meta, n=60, verbose=False):\n        \"\"\" generate `n` clusters of geo-coordinates and add them as a geo_label.\n            Fill records with missing geo-coordinates using existing records in the following order:\n            - same author and primary\n            - same author\n            - same primary\n        \"\"\"\n        n = 60\n        geo = meta.loc[~meta.latitude.isnull(), ['i', 'latitude', 'longitude']]\n        model = KMeans(n_clusters=n, max_iter=300, random_state=42)\n        model.fit(geo[['latitude', 'longitude']])\n        centers = model.cluster_centers_\n        i_to_geo_label = dict(zip(geo.i.values, model.labels_))\n        meta['geo_label'] = meta.i.map(i_to_geo_label)\n        meta['mode_primary'] = meta.groupby('primary_label')['geo_label'].transform(get_series_mode)\n        meta['mode_author'] = meta.groupby('author')['geo_label'].transform(get_series_mode)\n        meta['mode_author_primary'] = meta.groupby(['author', 'primary_label'])['geo_label'].transform(get_series_mode)\n        f = meta.geo_label.isnull()\n        if verbose: print('null coordinates', f.sum())\n        meta.loc[f, 'geo_label'] = meta.loc[f, 'mode_author_primary']\n        f = meta.geo_label.isnull()\n        if verbose: print('null coordinates', f.sum())\n        meta.loc[f, 'geo_label'] = meta.loc[f, 'mode_author']\n        f = meta.geo_label.isnull()\n        if verbose: print('null coordinates', f.sum())\n        meta.loc[f, 'geo_label'] = meta.loc[f, 'mode_primary']\n        f = meta.geo_label.isnull()\n        if verbose: print('null coordinates', f.sum())\n        meta['geo_label'] = meta.geo_label.astype('int16')\n        meta = meta.drop(labels=['mode_author_primary', 'mode_author', 'mode_primary'], axis=1)\n        return meta\n\n\n    def get_meta(geo_label=0, split_author=False, author_to_num=True):\n        meta = read_meta(split_author=split_author, author_to_num=author_to_num)\n        audios = pd.read_csv(f'{ROOT}/b5-cache/len_train_audios.csv')\n        audios.columns = ['filename', 'dim']\n        meta = meta.merge(audios)\n        if geo_label:\n            meta = add_geo_label(meta, n=geo_label)\n        return meta\n\n\n    def filter_meta(meta, filters):\n        filter_secondary, min_rating = filters\n        if filter_secondary:\n            meta = meta[meta.n_labels == 1]\n        #if min_rating > 0:\n        #    meta = meta[meta.rating >= min_rating]\n        return meta\n\n    \n    def single_edge_files(files):\n        \"\"\" takes a list of edge files and leaves only the end (if there is one) or teh start otherwise \"\"\"\n        if files[0][-1] not in ['s', 'e']:\n            return files\n        \n        single = []\n        for f in files:\n            if f[-1] != 's' or f[:-1]+'e' in files:\n                single.append(f)\n        return np.array(single)\n                \n\n    def get_meta_previous(folder='raw-30s'):\n        previous_files = glob.glob(f'{ROOT}/b5-data-previous-{folder}/*/*.ogg*')\n        previous_files = single_edge_files(previous_files)\n        primary_label = [x.split('/')[-2] for x in previous_files]\n        meta_previous = pd.DataFrame(dict(primary_label=primary_label, filename=previous_files))\n        meta_previous['labels'] = [[x] for x in primary_label]\n        meta_previous['rating'] = 5\n        meta_previous['n_labels'] = 1\n        meta_previous['fold'] = -1 # only use for training\n        return meta_previous\n\n\n    def get_meta_previous_from_source():\n        previous_files = glob.glob(f'{ROOT}/audio-bc-212223/AUDIO/*/*.ogg*')\n        previous_label = [x.split('/')[-2] for x in previous_files]\n        meta_previous = pd.DataFrame(dict(\n            primary_label=previous_label, \n            filename=previous_files,\n            dim=len_previous))\n        meta_previous['labels'] = [[x] for x in previous_label]\n        meta_previous['rating'] = 5\n        meta_previous['n_labels'] = 1\n        meta_previous['fold'] = -1 # only use for training\n        return meta_previous\n\n\n    def get_merged_meta(additional_data_path=None):\n        \"\"\" get merged train meta + meta for additional data \"\"\"\n        m1 = get_meta(geo_label=0, split_author=False, author_to_num=False)\n        if additional_data_path is not None:\n            m2 = load_obj(additional_data_path)\n            meta = pd.concat([m1, m2])\n        meta = meta.drop_duplicates(subset=['filename'], keep='first')\n        meta['source'] = np.where(meta.fold.isnull(), 'b25', 'previous')\n        meta.loc[meta.i.isnull(), 'source'] = 'other'\n        meta['counter'] = meta.groupby('primary_label')['filename'].transform('count')\n        meta['group'] = meta.author\n        meta['i'] = np.arange(len(meta))\n        meta = meta.drop(['fold'], axis='columns')\n        return meta\n\n\n    def get_meta_pseudo():\n        \"\"\" build meta and pseudo labels for pseudo labeeld data \"\"\"\n        \n        # build meta\n        if config.pseudo_preds.endswith('.csv'):  # Volodymyr's preds\n            oof = pd.read_csv(f'{ROOT}/{config.pseudo_preds}')\n            cols = list(oof.columns[1:])\n            idx_cols = [cols.index(x) for x in BIRDS]\n            pseudo_labels = oof.iloc[:, 1:].values[:, idx_cols]\n            ids = oof.row_id.values\n            meta = []\n            for id_ in ids:\n                filename = '_'.join(id_.split('_')[:-1]) + '.ogg'\n                order = int(id_.split('_')[-1]) // 5 - 1\n                meta.append(dict(filename=filename, order=order))\n            meta = pd.DataFrame(meta)        \n        \n        else: # my preds\n            # get pseudo oof\n            files = glob.glob(f'{ROOT}/{config.pseudo_preds}/oof*')\n            names, oof = [], []\n            for file in files:\n                # {ROOT}/b5-predict-unlabeled-0-3/oof_unlabeled_ebs.123_f0.pkl\n                name = file.split('/')[-1].split('_')[2]\n                if config.pseudo_models is not None and name not in config.pseudo_models:\n                    continue\n                names.append(name)\n                oof.append(load_obj(file))\n            oof = np.stack(oof)\n            oof = sigmoid(oof)\n            print(f'Pseudo labels ensembling: {sort_names(names)}')\n            \n            # ensemble oof and build pseudo labels\n            if config.pseudo_ensemble_mode == 'max':\n                pseudo_labels = oof.max(axis=0)\n            elif config.pseudo_ensemble_mode == 'mean':\n                pseudo_labels = oof.mean(axis=0)\n        \n            res_file = glob.glob(f'{ROOT}/{config.pseudo_preds}/res*')[0]\n            meta = load_obj(res_file)\n            meta = meta[['filename', 'order']]\n            meta['filename'] = [f'{x}.ogg' for x in meta.filename]\n\n        # filter and adjust pseudo labels\n        f_primary = np.max(pseudo_labels, axis=1) > config.pseudo_label_primary_th\n        pseudo_labels = pseudo_labels[f_primary]\n        pseudo_labels[pseudo_labels < config.pseudo_label_zero_th] = 0\n        \n        meta = meta[f_primary]\n        meta['primary_label'] = [BIRDS[x] for x in np.argmax(pseudo_labels, axis=1)]\n        meta['labels'] = [[x] for x in meta.primary_label]\n        meta['n_labels'] = 1\n        meta['i'] = np.arange(len(meta))\n        meta['group'] = np.arange(len(meta))\n        if config.pseudo_folds:\n            folds = meta[['i', 'primary_label', 'group']]\n            folds = add_mlstratified_folds(folds, n_folds=config.n_folds, random_state=config.seed, group='group', verbose=0)\n            meta = pd.merge(meta, folds[['i', 'fold']], on='i', how='left')\n        else:\n            meta['fold'] = -3\n        pseudo_weight = config.pseudo_weight or 1\n        meta['ws'] = 1 / meta.groupby('filename').transform('size') * pseudo_weight\n\n        return meta, pseudo_labels\n\n\n    def get_catalog_pseudo(meta, pseudo_labels, filename_to_path):\n        \"\"\" build catalog_pseudo: Dict[bird, list(Audios)], \n            where Audios: list(tuple(path, list(tuple(order, pseudo_label)))\n        \"\"\"\n        catalog_pseudo = dict()\n        filename = None\n        pseudo_labels = torch.tensor(pseudo_labels, dtype=torch.float32)\n        for row, pseudo_label in zip(meta.to_dict(orient='records'), pseudo_labels):\n            row = dotdict(row)\n            bird = BIRDS.index(row.primary_label)\n            if bird not in catalog_pseudo:\n                catalog_pseudo[bird] = dict()\n\n            path = filename_to_path[row.filename]\n            if path not in catalog_pseudo[bird]:\n                catalog_pseudo[bird][path] = []\n            catalog_pseudo[bird][path].append((row.order, pseudo_label))\n\n        # convert dict[filename, list(tuple(order, label)) into list[(filename, list(tuple(order, label))]\n        for bird in catalog_pseudo:\n            d_path_orders_labels = catalog_pseudo[bird]\n            lst = [(path, orders_labels) for path, orders_labels in d_path_orders_labels.items()]\n            catalog_pseudo[bird] = lst\n        \n        return catalog_pseudo     \n\n\n    def get_labels_primary(meta, num_classes=N_BIRDS):\n        \"\"\" get one_hot encoded labels for primary class only \"\"\"\n        labels = meta.primary_label.map(label_to_num).values.astype(np.int64)\n        labels = np.eye(num_classes)[labels].astype(np.float32)\n        return labels\n\n\n    def get_labels(meta, num_classes=N_BIRDS):\n        \"\"\" get one_hot encoded labels for all classes \"\"\"\n        labels_str = meta.labels\n        labels = np.zeros((meta.shape[0], num_classes), dtype=np.float32)\n        for i, ls in enumerate(labels_str):\n            for l in ls:\n                idx = label_to_num.get(l, -1)\n                if idx >= 0:\n                    labels[i, idx] = 1.\n        return labels\n    \n\n    def memory_usage_cap(meta, cap_seconds=30, dim=224):\n        \"\"\" Memory needed to old spectrograms with size (dim,dim) if capped at `cap_seconds` \"\"\"\n        caped_seconds = np.where(meta.dim < sr*cap_seconds, meta.dim/sr, cap_seconds)\n        size = caped_seconds.sum() / 5 * dim*dim * 4 / 1e9\n        print(f'<={cap_seconds}s: {size:.1f}MB')\n\n\n    def get_bird_groups():\n        \"\"\" group birds that show together in an audio \"\"\"\n        bird_to_group = dict()\n        next_id = 0\n        groups = dict()\n        for birds in meta.labels.values:\n            g = set()\n\n            # check for existing groups\n            for bird in birds:\n                if bird in bird_to_group:\n                    g.add(bird_to_group[bird])\n\n            if len(g) == 0:  # new group\n                gid = next_id\n                next_id += 1\n                groups[gid] = set()\n            elif len(g) == 1:  # existing group\n                gid = next(iter(g))\n            if len(g) > 1:  # merge groups\n                gid = next_id\n                next_id += 1\n                new_group = set()\n                for g_ in g:\n                    new_group = new_group.union(groups[g_])\n                    for b in groups[g_]:\n                        bird_to_group[b] = gid\n                    del groups[g_]\n                groups[gid] = new_group\n\n            # set group for each bird\n            for bird in birds:\n                bird_to_group[bird] = gid\n                groups[gid].add(bird)\n            # print(gid, birds, groups[gid])\n\n        return groups\n\n    # from https://www.kaggle.com/code/robbynevels/bc24-duplicate-audio-files/\n\n    def set_dataset(files, config):\n        new_files = []\n        for x in files:\n            folders = x.split('/')\n            if 'unlabeled' not in folders[3] and 'previous' not in folders[3]:\n                folders[3] = f'b5-data-{config.dataset}'\n            new_files.append('/'.join(folders))\n        return np.array(new_files)\n\n\n    def read(cfg):\n        if cfg.meta:\n            meta = load_obj(f'{ROOT}/b5-cache/meta {cfg.meta}')\n            meta.filename = [x.replace('_s', '').replace(f'{ROOT}/b5-data-raw-previous-2x7s/', '')\n                             for x in meta.filename]\n        else:\n            previous_dataset = 'additional-h5-curated-3'\n            meta = get_merged_meta(additional_data_path=f'{ROOT}/b5-data-{cfg.previous_dataset}-l-0/meta')\n            if not cfg.additional_data:\n                meta = meta[meta.source != 'other']\n                meta['counter'] = meta.groupby('primary_label')['filename'].transform('count')\n            print(meta.source.value_counts(), '\\n')\n            \n            if 1:  # split author in case there is only one per primary_label\n                meta['count_unique_group'] = meta.groupby('primary_label')['group'].transform('nunique')\n                meta['order'] = meta.groupby('primary_label').cumcount().astype('str')\n                meta['group'] = np.where(meta.count_unique_group == 1, meta.group + meta.order, meta.group)\n                meta = meta.drop(['order'], axis='columns')\n\n            # split data\n            meta['split_candidate'] = meta.counter >= 10\n            folds = meta.loc[meta.split_candidate, ['i', 'primary_label', 'group']]\n            folds = add_mlstratified_folds(folds, n_folds=cfg.n_folds, random_state=cfg.seed, group='group', verbose=0)\n            meta = pd.merge(meta, folds[['i', 'fold']], on='i', how='left')\n            meta.loc[meta.fold.isnull(), 'fold'] = -2\n\n        if 'ws' not in meta.columns:\n            meta['ws'] = 1\n        meta['order'] = -1\n\n        if cfg.single_label:\n            meta = meta[(meta.fold == 0) | (meta.n_labels == 1)]\n\n        datasets = [cfg.dataset, cfg.previous_dataset, cfg.train_unlabeled_dataset]\n        filename_to_path = get_filename_to_path(datasets)\n\n        # add unlabeled with pseudo soft labels\n        if cfg.train_unlabeled_dataset is not None:\n            meta_pseudo, pseudo_labels = get_meta_pseudo()\n            print(f'Equivalent pseudo audios: {meta_pseudo.ws.sum():,.0f}')\n            if cfg.sample_pseudo_prob == 0:\n                meta = pd.concat([meta, meta_pseudo], axis=0)\n                catalog_pseudo = None\n            else:\n                catalog_pseudo = get_catalog_pseudo(meta_pseudo, pseudo_labels, filename_to_path)\n\n        meta['path'] = meta.filename.map(filename_to_path)\n        print(meta.groupby('fold')['filename'].size())\n        meta.to_csv('meta.csv')\n        save_obj(meta, 'meta')\n\n        files = meta.path.values\n        if 'noise_path' in globals():\n            files_background_noise = glob.glob(f'{noise_path}/ff1010bird_nocall/nocall/*.ogg')\n        else:\n            files_background_noise = glob.glob(f'{ROOT}/birdclef2021-background-noise/ff1010bird_nocall/nocall/*.ogg')\n        assert meta.primary_label.isnull().sum() == 0, 'No support for null primary labels'\n        birds = meta.primary_label.map(label_to_num).values\n        labels = get_labels_primary(meta, N_BIRDS) if cfg.labels =='p' else get_labels(meta, N_BIRDS)\n        if cfg.train_unlabeled_dataset is not None and catalog_pseudo is None:\n            labels[-len(meta_pseudo):] = pseudo_labels  # use soft pseudo labels\n\n        print(f'\\nmeta shape: {meta.shape}  files: {len(files):,.0f}  labels: {len(labels):,.0f}  '\n              f'background noise: {len(files_background_noise):,.0f}')\n        \n        return meta, files, birds, labels, files_background_noise, catalog_pseudo\n    \n\n    # ----------------------------------------------------------------------------------------------------\n    # 0.5 L2\n    # ----------------------------------------------------------------------------------------------------\n\n    def build_bird_catalog(meta, min_duration):\n        \"\"\" build dict of lists of audios per bird. Each entry of the list contains the\n            filename, size of the audio and position in meta (this must be unique across all folds \n            as it's used to index all files and all labels):\n            - bird_catalog[bird_idx] = [(name, size, i)*].\n            audios muts have a minimum duration of `min_duration` seconds.\n        \"\"\"\n        bird_catalog = {k: [] for k in RANGE_BIRDS}\n        min_size = min_duration * 32_000\n        for row in meta.to_dict(orient='records'):\n            row = dotdict(row)\n            if row.dim < min_size:\n                continue\n            bird_idx = label_to_num[row.primary_label]\n            filename = row.filename\n            bird_catalog[bird_idx].append((filename, row.dim, int(row.i)))\n        return bird_catalog\n        \n\n    def get_usable_birds(bird_catalog, min_count=20, min_total_duration=4*4*60, verbose=True):\n        \"\"\" takes a bird_catalog dict and returns two arrays per bird:\n            birds_aux: birds with qualifying audio\n            birds_main: birds with at least `min_count` audios and whose total duration of \n                all audios adds to `min_total_duration` seconds.\n        \"\"\"\n        birds_aux, birds_main = set(), set()\n        for bird, audios in bird_catalog.items():\n            if len(audios) > 0:\n                birds_aux.add(bird)\n                total_duration = sum((x[1] for x in audios))\n                if len(audios) > min_count and total_duration > min_total_duration:\n                    birds_main.add(bird)\n        birds_aux = list(birds_aux)\n        birds_main = list(birds_main)\n        if verbose:\n            print(f'birds aux: {len(birds_aux)} - main: {len(birds_main)}')\n        return birds_aux, birds_main\n    \n\n    # ----------------------------------------------------------------------------------------------------\n    # 0.5 samples\n    # ----------------------------------------------------------------------------------------------------\n\n    class GroupKFoldShuffle(_BaseKFold):\n        \"\"\" random shuffle in K folds with unique groups in each fold.\n            Attempts to equalize the size of each fold as follows: groups are shuffled \n            and one at a time placed in the fold with the smallest size.\n        \"\"\"\n       \n        def __init__(self, n_splits=5, *, random_state=None):\n            super().__init__(n_splits=n_splits, shuffle=True, random_state=random_state)\n\n        def split(self, x, groups=None):\n            # assign each group to a fold\n            unique_groups, counts = np.unique(groups, return_counts=True)\n            group_counter = {g: c for g, c in zip(unique_groups, counts)}\n            rng = np.random.default_rng(seed=self.random_state)\n            unique_groups = rng.permutation(unique_groups)\n            groups_per_fold = [[] for x in range(self.n_splits)]\n            fold_size = [0] * self.n_splits\n            for g in unique_groups:\n                fold = np.argmin(fold_size)\n                groups_per_fold[fold].append(g)\n                fold_size[fold] += group_counter[g]\n            \n            # return fold indices\n            for test_group_ids in groups_per_fold:\n                test_mask = np.isin(groups, test_group_ids)\n                # train_mask = ~test_mask\n                # train_idx = np.where(train_mask)[0]\n                train_idx = None\n                test_idx = np.where(test_mask)[0]\n                yield train_idx, test_idx\n\n        def test(self, meta):\n            \"\"\" \"\"\"\n            df = meta.groupby(['patient_id', 'eeg_id'])['label_id'].count().reset_index()\n            x = df.index\n            g = df.patient_id.values\n            targets = []\n            ref = []\n            for i in range(20):\n                kf = GroupKFoldShuffle(n_splits=5, random_state=i)\n                l, idxs = [], []\n                for j, (train_index, test_index) in enumerate(kf.split(x, groups=g)):\n                    if len(set(g[train_index]).intersection(set(g[test_index]))) > 0:\n                        print('error')\n                    l.append(len(test_index))\n                    idxs.append(test_index)\n                diff = max(l) - min(l)\n                if i == 0:\n                    ref = idxs\n                else:\n                    for i1 in range(5):\n                        print('\\t', i1, end=':  ')\n                        for i2 in range(5):\n                            print(f'{len(set(idxs[i1]).intersection(set(ref[i2])))*100/len(idxs[i1]):.1f}', end='  ')\n                        print()\n                if diff < len(x)*.02:\n                    targets.append((i, idxs))\n                print(i, diff, f'{diff/len(x):.4f}')\n\n\n    def add_folds(meta, random_state=42, n_folds=5, verbose=0):\n        \"\"\" simple split. Return folds \"\"\"\n        if n_folds is None:\n            n_folds = config.n_folds\n        kf = KFold(n_splits=n_folds, shuffle=True, random_state=random_state)\n        index = meta.index\n        folds = np.arange(len(index))\n        for fold, (_, val_idx) in enumerate(kf.split(index)):\n            folds[val_idx] = fold\n        meta['fold'] = folds\n        if verbose:\n            print(meta.fold.value_counts())\n    \n\n    def add_group_folds(meta, group_col, random_state=42, n_folds=5, verbose=0):\n        \"\"\" split per group return folds as a dict \"\"\"\n        kf = GroupKFoldShuffle(n_splits=n_folds, random_state=random_state)\n        index = meta.index\n        groups = meta[group_col].values\n        folds = pd.Series(0, index=index)\n        for fold, (_, val_idx) in enumerate(kf.split(index, groups=groups)):\n            folds.iloc[val_idx] = fold\n        meta['fold'] = folds\n        # return folds.to_dict()\n\n\n    def add_stratified_folds(meta, strat_col, random_state=42, n_folds=5, verbose=0):\n        \"\"\" split per strat \"\"\"\n        skf = StratifiedKFold(n_splits=n_folds, random_state=random_state, shuffle=True)\n        strats = meta[strat_col]\n        index = np.arange(meta.shape[0])\n\n        folds = pd.Series(0, index=index)\n        for fold, (_, val_idx) in enumerate(skf.split(index, strats)):\n            folds.iloc[val_idx] = fold\n        meta['fold'] = folds\n            \n        if verbose:\n            values_col = 'ID'  # column that has no null values\n            labels_per_fold = meta.pivot_table(\n                index=strat_col, columns='fold', values=values_col, aggfunc=['count'])\n            print(labels_per_fold)\n            # labels_per_fold.to_csv('labels_per_fold.csv')\n            # labels_per_fold[labels_per_fold.isnull().any(axis=1)]\n\n   \n    def add_folds_strat(meta, strat_col='primary_label', group_col='author', seed=42, n_folds=5, verbose=0):\n        \"\"\" add folds to meta using stratification (split equally between folds) and groups\n            (each group in a single fold). Use add_mlstratified_folds because it's better.\n        \"\"\"\n\n        def _get_folds_strat(index, groups, strats, random_state=42, n_folds=5):\n            if n_folds is None:\n                n_folds = config.n_folds\n            sgkf = StratifiedGroupKFold(n_splits=n_folds, shuffle=True, random_state=random_state)\n            folds = pd.Series(0, index=index)\n            for fold, (_, val_idx) in enumerate(sgkf.split(index, y=strats, groups=groups)):\n                folds.iloc[val_idx] = fold\n            return folds\n\n        groups = meta[group_col] if group_col is not None else None\n        strats = meta[strat_col]\n        index = np.arange(meta.shape[0])\n        meta['fold'] = _get_folds_strat(index, groups, strats, random_state=seed, n_folds=n_folds)\n        \n        if verbose:\n            print(meta.fold.value_counts().sort_index())\n\n        # return folds\n\n\n    def IterativeStratification(labels, r, random_state):\n        \"\"\"This function implements the Iterative Stratification algorithm described\n        in the following paper:\n        Sechidis K., Tsoumakas G., Vlahavas I. (2011) On the Stratification of\n        Multi-Label Data. In: Gunopulos D., Hofmann T., Malerba D., Vazirgiannis M.\n        (eds) Machine Learning and Knowledge Discovery in Databases. ECML PKDD\n        2011. Lecture Notes in Computer Science, vol 6913. Springer, Berlin,\n        Heidelberg.\n        \"\"\"\n\n        n_samples = labels.shape[0]\n        test_folds = np.zeros(n_samples, dtype=int)\n\n        # Calculate the desired number of examples at each subset\n        c_folds = r * n_samples\n\n        # Calculate the desired number of examples of each label at each subset\n        c_folds_labels = np.outer(r, labels.sum(axis=0))\n\n        labels_not_processed_mask = np.ones(n_samples, dtype=bool)\n\n        while np.any(labels_not_processed_mask):\n            # Find the label with the fewest (but at least one) remaining examples,\n            # breaking ties randomly\n            num_labels = labels[labels_not_processed_mask].sum(axis=0)\n\n            # Handle case where only all-zero labels are left by distributing\n            # across all folds as evenly as possible (not in original algorithm but\n            # mentioned in the text). (By handling this case separately, some\n            # code redundancy is introduced; however, this approach allows for\n            # decreased execution time when there are a relatively large number\n            # of all-zero labels.)\n            if num_labels.sum() == 0:\n                sample_idxs = np.where(labels_not_processed_mask)[0]\n\n                for sample_idx in sample_idxs:\n                    fold_idx = np.where(c_folds == c_folds.max())[0]\n\n                    if fold_idx.shape[0] > 1:\n                        fold_idx = fold_idx[random_state.choice(fold_idx.shape[0])]\n\n                    test_folds[sample_idx] = fold_idx\n                    c_folds[fold_idx] -= 1\n\n                break\n\n            label_idx = np.where(num_labels == num_labels[np.nonzero(num_labels)].min())[0]\n            if label_idx.shape[0] > 1:\n                label_idx = label_idx[random_state.choice(label_idx.shape[0])]\n\n            sample_idxs = np.where(np.logical_and(labels[:, label_idx].flatten(), labels_not_processed_mask))[0]\n\n            for sample_idx in sample_idxs:\n                # Find the subset(s) with the largest number of desired examples\n                # for this label, breaking ties by considering the largest number\n                # of desired examples, breaking further ties randomly\n                label_folds = c_folds_labels[:, label_idx]\n                fold_idx = np.where(label_folds == label_folds.max())[0]\n\n                if fold_idx.shape[0] > 1:\n                    temp_fold_idx = np.where(c_folds[fold_idx] ==\n                                             c_folds[fold_idx].max())[0]\n                    fold_idx = fold_idx[temp_fold_idx]\n\n                    if temp_fold_idx.shape[0] > 1:\n                        fold_idx = fold_idx[random_state.choice(temp_fold_idx.shape[0])]\n\n                test_folds[sample_idx] = fold_idx\n                labels_not_processed_mask[sample_idx] = False\n\n                # Update desired number of examples\n                c_folds_labels[fold_idx, labels[sample_idx]] -= 1\n                c_folds[fold_idx] -= 1\n\n        return test_folds\n        \n    \n    class MultilabelStratifiedKFold(_BaseKFold):\n        \"\"\"Multilabel stratified K-Folds cross-validator\n        Provides train/test indices to split multilabel data into train/test sets.\n        This cross-validation object is a variation of KFold that returns\n        stratified folds for multilabel data. The folds are made by preserving\n        the percentage of samples for each label.\n        Parameters\n        ----------\n        n_splits : int, default=3\n            Number of folds. Must be at least 2.\n        shuffle : boolean, optional\n            Whether to shuffle each stratification of the data before splitting\n            into batches.\n        random_state : int, RandomState instance or None, optional, default=None\n            If int, random_state is the seed used by the random number generator;\n            If RandomState instance, random_state is the random number generator;\n            If None, the random number generator is the RandomState instance used\n            by `np.random`. Unlike StratifiedKFold that only uses random_state\n            when ``shuffle`` == True, this multilabel implementation\n            always uses the random_state since the iterative stratification\n            algorithm breaks ties randomly.\n        Examples\n        --------\n        >>> from iterstrat.ml_stratifiers import MultilabelStratifiedKFold\n        >>> import numpy as np\n        >>> X = np.array([[1,2], [3,4], [1,2], [3,4], [1,2], [3,4], [1,2], [3,4]])\n        >>> y = np.array([[0,0], [0,0], [0,1], [0,1], [1,1], [1,1], [1,0], [1,0]])\n        >>> mskf = MultilabelStratifiedKFold(n_splits=2, random_state=0)\n        >>> mskf.get_n_splits(X, y)\n        2\n        >>> print(mskf)  # doctest: +NORMALIZE_WHITESPACE\n        MultilabelStratifiedKFold(n_splits=2, random_state=0, shuffle=False)\n        >>> for train_index, test_index in mskf.split(X, y):\n        ...    print(\"TRAIN:\", train_index, \"TEST:\", test_index)\n        ...    X_train, X_test = X[train_index], X[test_index]\n        ...    y_train, y_test = y[train_index], y[test_index]\n        TRAIN: [0 3 4 6] TEST: [1 2 5 7]\n        TRAIN: [1 2 5 7] TEST: [0 3 4 6]\n        Notes\n        -----\n        Train and test sizes may be slightly different in each fold.\n        See also\n        --------\n        RepeatedMultilabelStratifiedKFold: Repeats Multilabel Stratified K-Fold\n        n times.\n        \"\"\"\n\n        def __init__(self, n_splits=3, shuffle=False, random_state=None):\n            super(MultilabelStratifiedKFold, self).__init__(n_splits, shuffle=shuffle, random_state=random_state)\n\n        def _make_test_folds(self, X, y):\n            y = np.asarray(y, dtype=bool)\n            type_of_target_y = type_of_target(y)\n\n            if type_of_target_y != 'multilabel-indicator':\n                raise ValueError(\n                    'Supported target type is: multilabel-indicator. Got {!r} instead.'.format(type_of_target_y))\n\n            num_samples = y.shape[0]\n\n            rng = check_random_state(self.random_state)\n            indices = np.arange(num_samples)\n\n            if self.shuffle:\n                rng.shuffle(indices)\n                y = y[indices]\n\n            r = np.asarray([1 / self.n_splits] * self.n_splits)\n\n            test_folds = IterativeStratification(labels=y, r=r, random_state=rng)\n\n            return test_folds[np.argsort(indices)]\n\n        def _iter_test_masks(self, X=None, y=None, groups=None):\n            test_folds = self._make_test_folds(X, y)\n            for i in range(self.n_splits):\n                yield test_folds == i\n\n        def split(self, X, y, groups=None):\n            \"\"\"Generate indices to split data into training and test set.\n            Parameters\n            ----------\n            X : array-like, shape (n_samples, n_features)\n                Training data, where n_samples is the number of samples\n                and n_features is the number of features.\n                Note that providing ``y`` is sufficient to generate the splits and\n                hence ``np.zeros(n_samples)`` may be used as a placeholder for\n                ``X`` instead of actual training data.\n            y : array-like, shape (n_samples, n_labels)\n                The target variable for supervised learning problems.\n                Multilabel stratification is done based on the y labels.\n            groups : object\n                Always ignored, exists for compatibility.\n            Returns\n            -------\n            train : ndarray\n                The training set indices for that split.\n            test : ndarray\n                The testing set indices for that split.\n            Notes\n            -----\n            Randomized CV splitters may return different results for each call of\n            split. You can make the results identical by setting ``random_state``\n            to an integer.\n            \"\"\"\n            y = check_array(y, ensure_2d=False, dtype=None)\n            return super(MultilabelStratifiedKFold, self).split(X, y, groups)\n\n        \n    def add_mlstratified_folds(meta, n_folds=5, group='author', random_state=42, verbose=True):\n        \"\"\" splits instances of score equally while keeping all instances of each anchor in the same fold \"\"\"\n        group_col = group\n        strat_col = 'primary_label'\n        dfx = pd.get_dummies(meta, columns=[strat_col]).groupby([group_col], as_index=False).sum()\n        cols = [c for c in dfx.columns if c.startswith('primary_label_') or c == group_col]\n        dfx = dfx[cols]\n\n        mskf = MultilabelStratifiedKFold(n_splits=n_folds, shuffle=True, random_state=random_state)\n        labels = [c for c in dfx.columns if c != group_col]\n        dfx_labels = dfx[labels]\n        dfx['fold'] = -1\n\n        for fold, (trn_, val_) in enumerate(mskf.split(dfx, dfx_labels)):\n            dfx.loc[val_, 'fold'] = fold\n\n        if 'fold' in meta:\n            del meta['fold']\n        meta = meta.merge(dfx[[group_col, 'fold']], on=group_col, how='left')\n        if verbose:\n            print(meta.fold.value_counts().sort_index())\n            if verbose == 2:\n                labels_per_fold = meta.pivot_table(index='primary_label', columns='fold', values='author', aggfunc=['count'])\n                print(labels_per_fold)\n                labels_per_fold.to_csv('labels_per_fold.csv')\n                # labels_per_fold[labels_per_fold.isnull().any(axis=1)]\n        return meta\n        \n    \n    def adjust_fold_small_samples(meta, th=10, fold=0):\n        \"\"\" moves small samples from fold `fold` to fold -1 \"\"\"\n        count_labels = meta.groupby('primary_label')['primary_label'].transform('count')\n        f_small = count_labels < th\n        meta.loc[f_small, 'fold'] = -1\n    \n    # ----------------------------------------------------------------------------------------------------\n    # 0.4 audios\n    # ----------------------------------------------------------------------------------------------------\n\n    def butter_lowpass_filter(data, cutoff_freq=20, sampling_rate=200, order=5):\n        nyquist = 0.5 * sampling_rate\n        normal_cutoff = cutoff_freq / nyquist  # to express cutoff in range 0..1 = cutoff * 2 / sampling_rate\n        b, a = butter(order, normal_cutoff, btype='low', analog=False)\n        filtered_data = lfilter(b, a, data, axis=-1)\n        return filtered_data\n\n\n    def maddest(d, axis=None):\n        return np.mean(np.absolute(d - np.mean(d, axis)), axis)\n\n\n    def wavelet_denoise(x, wavelet='haar', level=1):\n        \"\"\" set `wavelet` to the desired mode, e.g, db8, db2, dmey. \"\"\"\n        coeff = pywt.wavedec(x, wavelet, mode=\"per\")\n        sigma = (1/0.6745) * maddest(coeff[-level])\n\n        uthresh = sigma * np.sqrt(2*np.log(len(x)))\n        coeff[1:] = (pywt.threshold(i, value=uthresh, mode='hard') for i in coeff[1:])\n\n        ret = pywt.waverec(coeff, wavelet, mode='per')\n        return ret\n\n\n    def butter_bandpass_filter(data, lowcut=.5, highcut=20, fs=200, order=5):\n        b, a = butter(order, [lowcut, highcut], fs=fs, btype='band')\n        y = lfilter(b, a, data, axis=-1)\n        return y\n\n\n    def read_audios(meta, parallel=True):\n        folder = f'{ROOT}/birdclef-2025/train_audio'\n        \n        def _process(file):\n            # audio_data, sr = librosa.load(f'{folder}/{file}', sr=None)\n            audio_data, sr = torchaudio.load(file)\n            f_nan = np.isnan(audio_data)\n            if f_nan.mean() < 1:\n                m = np.nanmean(val)\n                audio_data[f_nan] = m\n            else:\n                audio_data[:] = 0\n            return audio_data.astype(np.float32)\n        \n        files = meta.filename.values\n        if parallel:\n            data = Parallel(n_jobs=-1)(delayed(_process)(file) for file in tqdm(files))\n        else:\n            data = []\n            for file in tqdm(files):\n                data.append(_process(file))\n\n        audios = dict()\n        for file, audio in zip(files, data):\n            audios[file] = audio\n        return audios\n\n    # ----------------------------------------------------------------------------------------------------\n    # 0.4 raw audio\n    # ----------------------------------------------------------------------------------------------------\n\n    import soundfile as sf\n\n    def build_raw(files, duration=5, filter=None, parallel=True):\n        \"\"\" read raw audios, limits size and saves as file \"\"\"\n\n        _chunk_size = 5*32_000\n\n        def _process(file):\n            \"\"\" returns single img \"\"\"\n            audio, sr = torchaudio.load(file)\n            audio = audio[:1, :sr*duration]\n            path = file.split('/')\n            label, filename = path[-2:]\n            if not os.path.exists(label):\n                try:\n                    os.makedirs(label)\n                except:\n                    pass\n            if filter is not None:\n                name = '/'.join(file.split('/')[-2:]).split('.')[0]\n                start, end = filter(name, audio)\n                if start is not None:\n                    audio = audio[:, start : end]\n            \n            audio = audio[:1].numpy().T\n            # following doesn't work for files > 60*32_000 and is slower otherwise\n            # sf.write(f'{label}/{filename}', audio.numpy().T, sr, format='ogg', subtype='vorbis') # or\n            # torchaudio.save('test.ogg', audios, format='ogg', sample_rate=sr, backend='soundfile')\n            if filename[-4:] != '.ogg':\n                filename += '.ogg'\n            with sf.SoundFile(f'{label}/{filename}', mode='w', samplerate=sr,\n                              channels=1, format='ogg', subtype='vorbis') as file:\n                for start in range(0, audio.shape[0], _chunk_size):\n                    end = start + _chunk_size\n                    file.write(audio[start:end])\n\n        if parallel:\n            Parallel(n_jobs=-1)(delayed(_process)(file) for file in tqdm(files))\n        else:\n            for file in tqdm(files):\n                _process(file)\n\n\n    def build_raw_edges(files, duration=7, min_duration=0, filter=None, parallel=True):\n        \"\"\" read raw audios, limits size and saves as file \"\"\"\n        \n        _chunk_size = 5*32_000\n        \n        def _process(file):\n            \"\"\" returns single img \"\"\"\n            audio, sr = torchaudio.load(file)\n            len_ = audio.shape[-1]\n            if min_duration and len_ < min_duration * sr:\n                return\n            path = file.split('/')\n            label, filename = path[-2:]\n            if not os.path.exists(label):\n                try:\n                    os.makedirs(label)\n                except:\n                    pass\n            if filter is not None:\n                name = '/'.join(file.split('/')[-2:]).split('.')[0]\n                start, end = filter(name, audio)\n                if start is not None:\n                    audio = audio[:, start : end]\n                    # print(name, raw_start/SR, raw_end/SR, audio.shape)\n            else:\n                start, end = None, None\n\n            audio = audio[:1].numpy().T\n            audios = audio[:sr*duration]\n            # following doesn't work for files > 60*32_000 and is slower otherwise\n            # sf.write(f'{label}/{filename}_s', audios, sr, format='ogg', subtype='vorbis')\n            if filename[-4:] != '.ogg':\n                filename += '.ogg'\n            with sf.SoundFile(f'{label}/{filename}_s', mode='w', samplerate=sr,\n                              channels=1, format='ogg', subtype='vorbis') as file:\n                for s in range(0, audios.shape[0], _chunk_size):\n                    e = s + _chunk_size\n                    file.write(audios[s : e])\n            \n            if audio.shape[0] > duration * sr:\n                audioe = audio[-sr*duration:]\n                # print('audioe', duration, audioe.shape)\n                # sf.write(f'{label}/{filename}_e', audioe, sr, format='ogg', subtype='vorbis')\n                with sf.SoundFile(f'{label}/{filename}_e', mode='w', samplerate=sr,\n                                  channels=1, format='ogg', subtype='vorbis') as file:\n                    for s in range(0, audioe.shape[0], _chunk_size):\n                        e = s + _chunk_size\n                        file.write(audioe[s : e])\n            \n            return start if start !=0 else None, end if end != len_ else None\n\n        if parallel:\n            edges = Parallel(n_jobs=-1)(delayed(_process)(file) for file in tqdm(files))\n        else:\n            edges = []\n            for file in tqdm(files):\n                edge = _process(file)\n                edges.append(edge)\n        starts, ends = zip(*edges)\n        files = ['/'.join(file.split('/')[-2:]) for file in files]\n        edges_df = pd.DataFrame(dict(start=starts, end=ends), index=files)\n        edges_df.to_csv('edges.csv')\n\n    \n    def load_audios(files, duration=5, parallel=True):\n        \"\"\" read raw audios, limits size and saves as file \"\"\"\n\n        def _process(file):\n            \"\"\" returns single img \"\"\"\n            audio, sr = torchaudio.load(file)\n            if duration is not None:\n                audio[:, :sr*duration]\n            return audio\n\n        if parallel:\n            audios = Parallel(n_jobs=-1)(delayed(_process)(file) for file in tqdm(files))\n        else:\n            for file in tqdm(files):\n                _process(file)\n        return audios\n\n\n    def size_audios(files, parallel=True):\n        \"\"\" read raw audios, limits size and saves as file \"\"\"\n\n        def _process(file):\n            \"\"\" returns single img \"\"\"\n            audio, sr = torchaudio.load(file)\n            #if sr != 32_000:\n            #    print(file)\n            return audio.shape[-1]\n\n        if parallel:\n            sizes = Parallel(n_jobs=-1)(delayed(_process)(file) for file in tqdm(files))\n        else:\n            for file in tqdm(files):\n                _process(file)\n        return sizes\n\n\n    def build_unlabeled(meta, duration=5, parallel=True):\n        \"\"\" read raw audios, limits size and saves as file \"\"\"\n\n        _chunk_size = 5*32_000\n\n        def _process(row):\n            \"\"\" returns single img \"\"\"\n            row = dotdict(row)\n            audio_file = f'{ROOT}/birdclef-2025/train_soundscapes/{row.audio}.ogg'\n            audio, sr = torchaudio.load(audio_file)\n            len_ = audio.shape[1]\n            if row.order == 0:  # first 5s\n                start = 0\n            elif row.order == 11:  # last 5s\n                start = len_ - sr * duration\n            else:\n                band = (duration - 5) // 2\n                start = (row.order - band) * sr\n            audio = audio[:1, start : start + sr * duration]\n            folder = row.primary_label\n            if not os.path.exists(folder):\n                try:\n                    os.makedirs(folder)\n                except:\n                    pass\n            \n            audio = audio.numpy().T\n            # following doesn't work for files > 60*32_000 and is slower otherwise\n            # sf.write(f'{label}/{filename}', audio.numpy().T, sr, format='ogg', subtype='vorbis') # or\n            # torchaudio.save('test.ogg', audios, format='ogg', sample_rate=sr, backend='soundfile')\n            with sf.SoundFile(f'{row.filename}', mode='w', samplerate=sr,\n                              channels=1, format='ogg', subtype='vorbis') as file:\n                for start in range(0, audio.shape[0], _chunk_size):\n                    end = start + _chunk_size\n                    file.write(audio[start:end])\n\n        rows = meta.to_dict(orient='records')\n        if parallel:\n            Parallel(n_jobs=-1)(delayed(_process)(row) for row in tqdm(rows))\n        else:\n            for row in tqdm(rows):\n                _process(row)\n\n\n    def build_unlabeled_h5(files, parallel=True):\n        \"\"\" read raw audios, limits size and saves as file \"\"\"\n\n        files_edges = dict()\n\n        def _process(file):\n            \"\"\" returns single img \"\"\"\n            audio, sr = torchaudio.load(file)\n            audio = audio[0]\n            start, end = 0, len(audio)\n            filename = file.split('/')[-1]\n            id_ = filename.split('.')[0]\n\n            files_edges[filename] = (start, end)  # saved audio file starts at 0\n        \n            # save h5 file\n            with h5py.File(f'{id_}.h5', 'w') as data_file:\n                data_file.create_dataset('raw', data=audio.numpy())\n\n        if parallel:\n            Parallel(n_jobs=-1)(delayed(_process)(file) for file in tqdm(files))\n        else:\n            for row in tqdm(files):\n                _process(file)\n        save_obj(files_edges, 'files_edges')\n\n\n    # ----------------------------------------------------------------------------------------------------\n    # 0.5 h5 audio\n    # ----------------------------------------------------------------------------------------------------\n    \n    def build_h5_files(files, folder, custom_filter={}, speech_filter=None):\n        \"\"\" save audio in h5 format \"\"\"\n\n        files_edges = dict()\n        for file in tqdm(files):\n            \"\"\" returns single img \"\"\"\n            path = f'{folder}/{file}' if folder is not None else file\n            audio, sr = torchaudio.load(path)\n            if sr != SR:\n                resampler = torchaudio.transforms.Resample(orig_freq=sr, new_freq=SR)\n                audio = resampler(audio)\n            audio = audio[0]\n            id_ = '/'.join(file.split('/')[-2:]).split('.')[0]\n            sections = custom_filter.get(id_, None)\n            if sections is not None:\n                start = int(sections[1][0] * SR)\n                if sections[-1][1] is not None:\n                    end = min(int(sections[-1][1] * SR), len(audio))\n                else:\n                    end = len(audio)\n            elif speech_filter is not None:\n                start, end = speech_filter(audio)\n                end = min(end, len(audio))\n            else:\n                start, end = 0, len(audio)\n                    \n            if start is not None:\n                audio = audio[start : end]\n        \n            # create folder if needed\n            label, filename = file.split('/')[-2:]\n            if not os.path.exists(label):\n                try:\n                    os.makedirs(label)\n                except:\n                    pass\n        \n            # save h5 file\n            with h5py.File(f'{id_}.h5', 'w') as data_file:\n                data_file.create_dataset('raw', data=audio.numpy())\n\n            k = '/'.join(file.split('/')[-2:])\n            if '.' in k:\n                k = k.split('.')[0]\n            k = k + '.ogg'\n            files_edges[k] = (start, end)\n        \n        save_obj(files_edges, 'files_edges')\n\n    # ----------------------------------------------------------------------------------------------------\n    # 0.5 google drive\n    # ----------------------------------------------------------------------------------------------------\n\n    from pydrive2.auth import GoogleAuth\n    from pydrive2.drive import GoogleDrive\n    from google.colab import auth\n    from oauth2client.client import GoogleCredentials\n\n    def get_or_create_folder(folder_name, parent_id=None):\n        \"\"\" given a folder name and a parent_id for google drive, creates the folder if it doesn't exist \n            and return its id.\n        \"\"\"\n        query = f\"title='{folder_name}' and mimeType='application/vnd.google-apps.folder' and trashed=false\"\n        if parent_id:\n            query += f\" and '{parent_id}' in parents\"\n        \n        file_list = drive.ListFile({'q': query}).GetList()\n        \n        if file_list:\n            return file_list[0]['id']\n        \n        # Create folder\n        metadata = {'title': folder_name,\n                    'mimeType': 'application/vnd.google-apps.folder'}\n        if parent_id:\n            metadata['parents'] = [{'id': parent_id}]\n        \n        folder = drive.CreateFile(metadata)\n        folder.Upload()\n        return folder['id']\n\n\n    def get_folder_id_from_path(path, parent_id='root'):\n        \"\"\" get the folder_id for a path in google drive \"\"\"\n        folders = path.strip('/').split('/')\n        current_id = parent_id\n        for folder in folders:\n            file_list = drive.ListFile({\n                \"q\": f\"'{current_id}' in parents and title='{folder}' and mimeType='application/vnd.google-apps.folder' and trashed=false\"\n            }).GetList()\n            if not file_list:\n                return None\n            current_id = file_list[0]['id']\n        return current_id\n        \n\n    def get_drive_subfolders(parent_id):\n        \"\"\"Get all subfolders of a parent folder.\"\"\"\n        query = f\"'{parent_id}' in parents and mimeType = 'application/vnd.google-apps.folder' and trashed=false\"\n        return drive.ListFile({'q': query}).GetList()\n\n\n    def get_drive_subfolders_dict(parent_id):\n        query = f\"'{parent_id}' in parents and mimeType='application/vnd.google-apps.folder' and trashed=false\"\n        return {f['title']: f for f in drive.ListFile({'q': query}).GetList()}\n\n\n    def get_drive_files(folder_id):\n        \"\"\"Get all files in a folder.\"\"\"\n        query = f\"'{folder_id}' in parents and trashed=false\"\n        return drive.ListFile({'q': query}).GetList()\n\n\n    def get_drive_files_dict(folder_id):\n        query = f\"'{folder_id}' in parents and trashed=false\"\n        return {f['title']: f for f in drive.ListFile({'q': query}).GetList()}\n\n\n    def upload_file(local_path, parent_id):\n        file = drive.CreateFile({'title': os.path.basename(local_path), 'parents': [{'id': parent_id}]})\n        file.SetContentFile(local_path)\n        file.Upload()\n        print(f\"Uploaded: {local_path}\")\n\n\n    def merge_duplicates(parent_folder_id):\n        subfolders = get_drive_subfolders(parent_folder_id)\n\n        # Group by normalized names\n        name_to_folders = defaultdict(list)\n        for folder in subfolders:\n            base_name = normalize_name(folder['title'])\n            name_to_folders[base_name].append(folder)\n\n        for name, folders in name_to_folders.items():\n            if len(folders) <= 1:\n                continue  # No duplicates\n\n            # Sort folders to find the \"primary\" one (usually the one without (n))\n            folders.sort(key=lambda f: (re.search(r'\\(\\d+\\)$', f['title']) is not None, f['title']))\n            primary_folder = folders[0]\n            primary_id = primary_folder['id']\n\n            print(f\"\\nMerging duplicates into: {primary_folder['title']}\")\n            for dup_folder in folders[1:]:\n                print(f\"  → Moving from: {dup_folder['title']}\")\n                for file in get_drive_files(dup_folder['id']):\n                    file['parents'] = [{'id': primary_id}]\n                    file.Upload()\n                dup_folder.Delete()\n        \n\n    def build_catalog_pseudo(cfg):\n        \"\"\" build array `pseudo_audios` with each audio segment with a valid pseudo label and corresponding \n            catalog[file]=List[idx]. The catalog contains a list of the indices for each 5s segment. Segments\n            with invalid pseudo labels use index `None`.\n        \"\"\"\n\n    def get_catalog_pseudo(cfg):\n        \"\"\" load array `pseudo_audios` and corresponding `pseudo_catalog`. \"\"\"\n        \n        \n    # ----------------------------------------------------------------------------------------------------\n    # 0.6 audio filters\n    # ----------------------------------------------------------------------------------------------------\n\n    def ts_to_s(speech_timestamps):\n        \"\"\" convert audio index positions with sampling rate `SR` to seconds \"\"\"\n        out = []\n        for dt in speech_timestamps:\n            ds = dict(start=dt['start']/SR, end=dt['end']/SR)\n            out.append(ds)\n        return out\n\n\n    def complete_audio_filters(raw_audio_filters):\n        \"\"\" finalize a dict of audio dilters by converting audio hits into audio sections (start, end),\n            merging them when applicable (if the distance to the previous hit <= 5s).\n            section for a hit is defined a (hit - BAND, hit + BAND)\n            The list sections is prefixed with the type of curation:\n                - 'm': the bird is vocalizing in every 5s segment of every section\n                - 'a': the bird is not guaranted to vocalize in every 5s segment of every section\n        \"\"\"\n        BAND = 4\n        \n        audio_filters = dict()\n        for id_, raw_sections in raw_audio_filters.items():\n            prior_hit = -100\n            if raw_sections[0] == 'm':\n                curation = 'm'\n                start = 1\n            else:\n                curation = 'a'\n                start = 0\n            sections = [curation]\n            for sec in raw_sections[start:]:\n                if isinstance(sec, tuple):\n                    sections.append(sec)\n                    if sec[1] is not None:\n                        prior_hit = sec[1] - 4\n                else:\n                    if prior_hit + 5 >= sec:  # merge with previous section\n                        sections[-1] = (sections[-1][0], sec + BAND)\n                    else:\n                        start = max(sec - BAND, 0)\n                        end = 5 if start == 0 else sec + BAND\n                        sections.append((start, end))\n                    prior_hit = sec\n                        \n            audio_filters[id_] = sections\n        return audio_filters   \n    \n\n    def build_custom_filter(meta):\n        \"\"\" build list `custom_filter` with audio sections, i.e., tuple(start, end), for \n            the files in meta in the same order.\n            The original start and end values are adjusted to account for the fact that the\n            audio file skips the beginning of the recording up to the start of the first section.\n            The first item in a section list i steh type of curation: (m)anual, (a)utomatic.\n        \"\"\"\n        # complete train filters and get file edges\n        train_audio_filters = complete_audio_filters(TRAIN_AUDIO_FILTERS)\n        additional_audio_filters = complete_audio_filters(ADDITIONAL_AUDIO_FILTERS)\n        raw_audio_filters = train_audio_filters | additional_audio_filters\n        datasets = [config.dataset, config.previous_dataset, config.train_unlabeled_dataset]\n        files_edges = get_files_edges(datasets)\n        \n        # prepare custom filter format\n        custom_filter = []\n        for row in meta.to_dict(orient='records'):\n            row = dotdict(row)\n            id_ = row.filename[:-4]\n            \n            if row.order != -1:  # unlabeled with pseudo labels\n                sections = ['a', (row.order * 5, (row.order + 1) * 5)]\n                sections_prob = [1]\n            else:\n                sections = raw_audio_filters.get(id_, None)\n                if sections is None:\n                    # if no files edges default to 7s (it should only apply to unlabeled audios)\n                    section_start, section_end = files_edges.get(row.filename, (0, 7*SR))\n                    sections = ['a', (section_start/SR, section_end/SR)]\n                elif sections[-1][-1] is None:  # end of audio\n                    sections[-1] = (sections[-1][0], files_edges[row.filename][1] / SR)\n\n                # make sure that end of last section is not larger than the length of the audio file\n                sections[-1] = (sections[-1][0], min(sections[-1][1], files_edges[row.filename][1] / SR))\n\n                # adjust section start and end based on audio start\n                start = sections[1][0]\n                sections = sections[:1] + [(s - start, e - start) for s, e in sections[1:]]\n                \n                # define section probability based on size\n                sections_prob = [(end - start) for start, end in sections[1:]]\n\n            custom_filter.append(dotdict(sections=sections, sections_prob=sections_prob))\n        return custom_filter\n\n\n    def get_lot_per_species(meta):\n        dim_per_species = meta.groupby('primary_label')['dim'].sum()\n        lot_per_species = dict()\n        lot = 0\n        lot_size = 0\n        max_dim_per_lot = 19.5e9 / 32_000 / 4\n        for species, dim in zip(dim_per_species.index, dim_per_species.values):\n            if lot_size + dim > max_dim_per_lot:\n                lot += 1\n                lot_size = 0\n            lot_size += dim\n            lot_per_species[species] = lot\n        return lot_per_species\n\n\n    def get_files_for_lot(meta, lot, filename=None):\n        \"\"\" split files in lots that fit in 19GB \"\"\"\n        lot_per_species = get_lot_per_species(meta)\n        files = []\n        for i, row in enumerate(meta.to_dict(orient='records')):\n            row = dotdict(row)\n            if lot_per_species[row.primary_label] != lot:\n                continue\n            if filename is not None:\n                files.append(filename[i])\n            else:\n                files.append(row.filename)\n        return files    \n    \n\n    def get_filename_to_lot(dataset, verbose=0):\n        filename_to_lot = dict()\n        for lot in range(7):\n            lot_files = glob.glob(f'{ROOT}/b5-data-{dataset}-l-{lot}/*/*.h5')\n            if len(lot_files) == 0:\n                break\n            for file in lot_files:\n                filename = '/'.join(file.split('/')[-2:])[:-3] + '.ogg'\n                filename_to_lot[filename] = file\n            if verbose:\n                print(lot, len(lot_files))\n        return filename_to_lot\n\n\n    def get_filename_to_path(datasets):\n        filename_to_path = dict()\n        for dataset in datasets:\n            if dataset is None:\n                continue\n            if 'unlabeled' in dataset:\n                files = glob.glob(f'{ROOT}/b5-data-{dataset}*/*.h5')\n                id_offset = 1\n            else:\n                files = glob.glob(f'{ROOT}/b5-data-{dataset}*/*/*.h5')\n                id_offset = 2\n            if len(files) == 0:\n                break\n            for file in files:\n                filename = '/'.join(file.split('/')[-id_offset:])[:-3] + '.ogg'\n                filename_to_path[filename] = file\n        return filename_to_path\n    \n    \n    def get_files_edges(datasets):\n        \"\"\" gest the start and end indices of the audio (ignoring speech and non vocalization sections) \"\"\"\n        files_edges = dict()\n        \n        n_files = 0\n        for dataset in datasets:\n            if dataset is None:\n                continue\n            files = glob.glob(f'{ROOT}/b5-data-{dataset}*/files_edges.pkl')\n            for file in files:\n                temp = load_obj(file)\n                # cleanup keys\n                cleaned = dict()\n                for k, v in temp.items():\n                    k = '/'.join(k.split('/')[-2:])\n                    if '.' in k:\n                        k = k.split('.')[0]\n                    k = k + '.ogg'\n                    cleaned[k] = v\n                files_edges = files_edges | cleaned\n                n_files += 1\n        print(f'Loaded files_edges: {n_files}')\n        return files_edges\n        \n    \n    SPEECH_DB_TH = -50\n\n    class Speech_filter:\n        \"\"\" Class that takes an audio and returns the start and end indices of the species\n            singing mostly by cutting off human speech.\n        \n            It identifies human speech using a combination of two mechanisms:\n            1) We accumulate the audio power over chunks of .1s and identify points in which it intersects\n                -50 db, which is considered silence. Usually a silent period occurs before the author adds\n                speech comments to the audio.\n            2) We use a model to identify the segments (start / end) of speech in the audio.\n            \n            We use the following parameters:\n                - sing_min_duration: min duration of singing between two periods of silence (power=-50db).\n                - speech_notes_time: if silence initiates (power=-50db) after this many seconds we cut the\n                    audio (likely speech follows).\n                - spech_merge_th: two consecutive speeches with an interval smaller than `spech_merge_th`\n                    are merged into a single speech\n                - speech_start_th: spech that occurs before `speech_start_th` seconds is considered a false\n                    positive.\n                - speech_min_duration: any speech that last less than this number of seconds is ignored\n                    \n            And implement the following rules (firt 1a/1b and then 2):\n            1a) If power crosses -50 db up and then down, for a duration >= sing_min_duration then we mark\n                those two points as the start and end of the audio.\n            1b) otherwise, if power crosses -50 down after speech_notes_time seconds, we mark the audio as\n                starting at 0 and ending at that point.\n            \n            2) We then use a model to identify periods of speech and traverse those periods:\n                - if a speech period initiates within spech_merge_th seconds of the previous one, we merge\n                    them.\n                - if the duration of the current speech is >= speech_min_duration:\n                    - if the speech start time <= speech_start_th, we assume it as false positive and ignore\n                        further speech periods.\n                    - otherwise we cut the audio at the start of the speech and return\n                - if the duration of the current speech > .5 and speech start time > 30 seconds, we cut the \n                    audio at the start of the speech and return.\n                - otherwise we move to the following speech period.\n        \"\"\"\n\n        def __init__(self, sing_min_duration=2, speech_notes_time=7, spech_merge_th=.3, speech_min_duration=2, \n                     speech_start_th=8, th=.5, threads=1):\n            self.sing_min_duration = sing_min_duration\n            self.speech_notes_time = speech_notes_time\n            self.spech_merge_th = spech_merge_th\n            self.speech_min_duration = speech_min_duration\n            self.speech_start_th = speech_start_th\n            self.th = th\n            \n            # interval variables\n            self.chunk_len = 0.1\n            self.chunk = int(self.chunk_len * SR)\n            torch.set_num_threads(threads)\n            self.model, (self.get_speech_timestamps, _, read_audio, _, _) = \\\n                torch.hub.load(repo_or_dir='snakers4/silero-vad', model='silero_vad')\n\n                \n        def __call__(self, audio, th=None):\n            if len(audio.shape) > 1:\n                audio = audio[0]\n            len_ = len(audio)\n            \n            # power based detection\n            chunk = int(self.chunk_len * SR)\n            power = audio ** 2\n            pad = int(np.ceil(len(power) / chunk) * chunk - len(power))\n            power = np.pad(power, (0, pad))\n            power = power.reshape((-1, chunk)).sum(axis=1)\n            power_dB = 10 * np.log10(power)\n            x = power_dB - SPEECH_DB_TH\n            start, end = 0, len_\n            intersections = np.where(x[:-1] * x[1:] < 0)[0]\n            for s, e in zip(intersections[:-1], intersections[1:]):\n                if x[s] < x[s+1] and (e - s) * self.chunk_len >= self.sing_min_duration:\n                    start, end = s * chunk, e * chunk\n                    break\n                elif x[s] > x[s+1] and s * self.chunk_len > self.speech_notes_time:\n                    start, end = 0, s * chunk\n                    break\n\n            # model based detection\n            threshold = th if th is not None else self.th\n            speech_timestamps = self.get_speech_timestamps(\n                audio[start : end], self.model, sampling_rate=SR, threshold=threshold)\n            \n            if len(speech_timestamps) > 0:\n                s, e = -1e6, -1e6\n                # print(ts_to_s(speech_timestamps))\n                for ts in speech_timestamps:\n                    if ts['start'] - e < self.spech_merge_th * SR:  # merge\n                        e = ts['end']\n                    else:\n                        s, e = ts['start'], ts['end']\n                    duration = (e - s) / SR\n                    start_s = (start + s) / SR\n                    if duration >= self.speech_min_duration or (duration > .5 and start_s >= 30):\n                        if start_s <= self.speech_start_th:\n                            break  # probably a false positive\n                        start, end = start, start + s\n                        break\n            return start, end        \n\n    # ----------------------------------------------------------------------------------------------------\n    # 0.6 specs val\n    # ----------------------------------------------------------------------------------------------------\n\n    def build_data_montage_val_fold(\n            files, labels, montages, norm=False, dim=(128,256),\n            duration=5, bits=8, snipet=False, spec=None,\n            unlabeled_selected=None, unlabeled_reference=None, offset=None, norm_audio=False,\n            parallel=True):\n        transform_to_spec_ = transform_to_spec(\n            dim=dim, norm=norm, duration=duration, snipet=snipet, spec=spec, bits=bits)\n\n        def _process(idx):\n            \"\"\" creates single img and return label \"\"\"\n            file, sample, y = files[idx], montages[idx], labels[idx]\n            audio, sr = torchaudio.load(file)\n            audio = audio[:, :duration * sr]\n            gap = len_audio_5s - audio.shape[-1]\n\n            if gap > 0:\n                pad_before = gap // 2\n                pad_after = gap - pad_before\n                audio = F.pad(audio, pad=(pad_before, pad_after), mode='constant', value=0)\n            \n            if len(sample) > 0:\n                for idx_m in sample:\n                    secondary_audio, sr = torchaudio.load(files[idx_m])\n                    secondary_audio = secondary_audio[:, :duration * sr]\n                    gap = len_audio_5s - secondary_audio.shape[-1]\n                    if gap > 0:\n                        pad_before = gap // 2\n                        pad_after = gap - pad_before\n                        secondary_audio = F.pad(secondary_audio,\n                            pad=(pad_before, pad_after), mode='constant', value=0)\n                    audio = audio + secondary_audio\n                    y = y + labels[idx_m]\n\n            if unlabeled_selected is not None and unlabeled_selected[idx] is not None:\n                audio = add_unlabeled(\n                    audio, unlabeled_selected[idx], offset[idx], unlabeled_reference)\n                \n            if norm_audio:\n                audio /= torch.abs(audio).max()\n            spec = transform_to_spec_(audio)\n            return spec, y\n\n        if parallel:\n            out = Parallel(n_jobs=-1)(delayed(_process)(idx) for idx in tqdm(range(len(files))))\n        else:\n            out = []\n            for idx in tqdm(range(len(files))):\n                out.append(_process(idx))\n        x, y = zip(*out)\n        x, y = np.stack(x), np.stack(y)\n        return x, y\n\n\n    def build_data_montage_val_fold_h5(\n            files, labels, montages, norm=False, dim=(128,256),\n            duration=5, bits=8, snipet=False, spec=None,\n            unlabeled_selected=None, unlabeled_reference=None, offset=None, norm_audio=False,\n            parallel=True):\n        transform_to_spec_ = transform_to_spec(\n            dim=dim, norm=norm, duration=duration, snipet=snipet, spec=spec, bits=bits)\n\n        def _process(idx):\n            \"\"\" creates single img and return label \"\"\"\n            file, sample, y = files[idx], montages[idx], labels[idx]\n            with h5py.File(file, 'r') as f:\n                audio = f['raw'][0 : duration * SR]\n            audio = torch.as_tensor(audio)\n            gap = duration * SR - audio.shape[-1]\n\n            if gap > 0:\n                pad_before = gap // 2\n                pad_after = gap - pad_before\n                audio = F.pad(audio, pad=(pad_before, pad_after), mode='constant', value=0)\n            \n            if len(sample) > 0:\n                for idx_m in sample:\n                    with h5py.File(files[idx_m], 'r') as f:\n                        secondary_audio = f['raw'][0 : duration * SR]\n                    secondary_audio = torch.as_tensor(secondary_audio)\n                    gap = len_audio_5s - secondary_audio.shape[-1]\n                    if gap > 0:\n                        pad_before = gap // 2\n                        pad_after = gap - pad_before\n                        secondary_audio = F.pad(secondary_audio,\n                            pad=(pad_before, pad_after), mode='constant', value=0)\n                    audio = audio + secondary_audio\n                    y = y + labels[idx_m]\n            audio = audio.unsqueeze(0)\n            \n            if unlabeled_selected is not None and unlabeled_selected[idx] is not None:\n                audio = add_unlabeled(\n                    audio, unlabeled_selected[idx], offset[idx], unlabeled_reference)\n                \n            if norm_audio:\n                audio /= torch.abs(audio).max()\n            spec = transform_to_spec_(audio)\n            return spec, y\n\n        if parallel:\n            out = Parallel(n_jobs=-1)(delayed(_process)(idx) for idx in tqdm(range(len(files))))\n        else:\n            out = []\n            for idx in tqdm(range(len(files))):\n                out.append(_process(idx))\n        x, y = zip(*out)\n        x, y = np.stack(x), np.stack(y)\n        return x, y\n\n\n    def add_unlabeled(x, unlabeled_file, offset, unlabeled_reference):\n        \"\"\" add unlabeled segment from file `unlabeled_selected` to x \"\"\"\n        duration = 5\n        \n        unlabeled_audio, sr = torchaudio.load(unlabeled_file)\n        gap = duration * sr - unlabeled_audio.shape[-1]\n        if gap > 0:\n            pad_before = random.randint(0, gap - 1)\n            pad_after = gap - pad_before\n            unlabeled_audio = F.pad(unlabeled_audio, pad=(pad_before, pad_after), mode='constant', value=0)\n        elif gap < 0:\n            start = int(-gap * offset)\n            unlabeled_audio = unlabeled_audio[:, start : start + duration * sr]\n\n        if unlabeled_reference:\n            max_x = np.abs(x).max()\n            max_unlabeled = np.abs(unlabeled_audio).max()\n            x = x / max_x * max_unlabeled + unlabeled_audio\n        else:\n            x = x + unlabeled_audio\n\n        return x\n\n\n    def build_catalog_birds(meta_val):\n        \"\"\" adds to catalog bird instances with no secondary labels \"\"\"\n        catalog = {i: [] for i in RANGE_BIRDS}\n        \n        for i, row in enumerate(meta_val.to_dict(orient='records')):\n            row = dotdict(row)\n            if row.n_labels == 1:\n                bird = label_to_num[row.primary_label]\n                catalog[bird].append(i)\n        return catalog\n        \n\n    def get_montages(catalog, birds, apply_montage, seed=10):\n        \"\"\" gets montages \"\"\"\n        random.seed(seed)\n        montages = []\n        for i, apply in enumerate(apply_montage):\n            if not apply:\n                sample_birds = []\n            else:\n                sample_birds = [x for x in random.sample(RANGE_BIRDS, random.randint(*config.montage)) if x != birds[i]]\n            sample = []\n            for bird in sample_birds:\n                options = catalog[bird]\n                if len(options) > 0:\n                    idx = random.randint(0, len(options) - 1)\n                    sample.append(options[idx])\n            montages.append(sample)\n        random.seed()\n        return montages\n\n\n    def get_unlabeled(birds, files_unlabeled, unlabeled_prob, seed=10):\n        \"\"\" gets list of unlabeled audios and random relative offset \"\"\"\n        random.seed(seed)\n        n_birds, n_files = len(birds), len(files_unlabeled)\n        unlabeled_idx = np.random.randint(0, n_files, n_birds)\n        probs = np.random.rand(n_birds)\n        offset = np.random.rand(n_birds)\n        unlabeled_selected = [files_unlabeled[unlabeled_idx[i]] if probs[i] < unlabeled_prob\n                              else None for i in range(n_birds)]\n        random.seed()\n        return unlabeled_selected, offset\n\n\n    def build_data_val(meta, labels, norm=False, dim=(128,256), duration=5, bits=8,\n                       folds=None, seed=10, folder=None, snipet=False, spec=None, \n                       unlabeled_prob=.9, files_unlabeled=None, unlabeled_reference=None,\n                       norm_audio=False, parallel=True):\n        \"\"\" build validation montage \"\"\"\n        if unlabeled_reference is None:\n            unlabeled_reference = True if unlabeled_prob else False\n        for fold in folds:\n            val_idx = meta.fold.values == fold\n            meta_val = meta[val_idx]\n            catalog_birds = build_catalog_birds(meta_val)\n            files = [f'{folder}/{x}' for x in meta_val.filename]\n            montages = get_montages(catalog_birds, birds[val_idx], meta_val.n_labels == 1, seed=10)\n            if unlabeled_prob:\n                unlabeled_selected, offset = \\\n                    get_unlabeled(birds, files_unlabeled, unlabeled_prob, seed=10)\n            else:\n                unlabeled_selected, offset = None, None\n            x, y = build_data_montage_val_fold(files, labels[val_idx], montages, norm=norm, dim=dim,\n                    duration=duration, bits=bits, snipet=snipet, spec=spec,\n                    unlabeled_selected=unlabeled_selected, unlabeled_reference=unlabeled_reference,\n                    offset=offset, norm_audio=norm_audio, parallel=parallel)\n            \n            assert y.max() == 1\n            save_obj(x, f'x_f{fold}')\n            save_obj(y, f'y_f{fold}')\n\n\n    def build_data_val_h5(meta, labels, norm=False, dim=(128,256), duration=5, bits=8,\n                       folds=None, seed=10, folder=None, snipet=False, spec=None, \n                       unlabeled_prob=.9, files_unlabeled=None, unlabeled_reference=None,\n                       norm_audio=False, parallel=True):\n        \"\"\" build validation montage \"\"\"\n        if unlabeled_reference is None:\n            unlabeled_reference = True if unlabeled_prob else False\n        for fold in folds:\n            val_idx = meta.fold.values == fold\n            meta_val = meta[val_idx]\n            catalog_birds = build_catalog_birds(meta_val)\n            files = meta_val.path.values\n            montages = get_montages(catalog_birds, birds[val_idx], meta_val.n_labels == 1, seed=10)\n            if unlabeled_prob:\n                unlabeled_selected, offset = \\\n                    get_unlabeled(birds, files_unlabeled, unlabeled_prob, seed=10)\n            else:\n                unlabeled_selected, offset = None, None\n            x, y = build_data_montage_val_fold_h5(files, labels[val_idx], montages, norm=norm, dim=dim,\n                    duration=duration, bits=bits, snipet=snipet, spec=spec,\n                    unlabeled_selected=unlabeled_selected, unlabeled_reference=unlabeled_reference,\n                    offset=offset, norm_audio=norm_audio, parallel=parallel)\n            \n            assert y.max() == 1\n            save_obj(x, f'x_f{fold}')\n            save_obj(y, f'y_f{fold}')\n\n\n    def build_data_train(meta, labels, norm=False, dim=(128,256), duration=5, bits=8,\n                       folds=None, seed=10, folder=None, snipet=False, spec=None, \n                       unlabeled_prob=.9, files_unlabeled=None, unlabeled_reference=None,\n                       norm_audio=False, parallel=True):\n        \"\"\" build validation montage \"\"\"\n        if unlabeled_reference is None:\n            unlabeled_reference = True if unlabeled_prob else False\n        for fold in folds:\n            idx = meta.fold.values != fold\n            meta_val = meta[idx]\n            catalog_birds = build_catalog_birds(meta_val)\n            files = [f'{folder}/{x}' for x in meta_val.filename]\n            montages = get_montages(catalog_birds, birds[idx], meta_val.n_labels == 1, seed=10)\n            if unlabeled_prob:\n                unlabeled_selected, offset = \\\n                    get_unlabeled(birds, files_unlabeled, unlabeled_prob, seed=10)\n            else:\n                unlabeled_selected, offset = None, None\n            x, y = build_data_montage_val_fold(files, labels[idx], montages, norm=norm, dim=dim,\n                    duration=duration, bits=bits, snipet=snipet, spec=spec,\n                    unlabeled_selected=unlabeled_selected, unlabeled_reference=unlabeled_reference,\n                    offset=offset, norm_audio=norm_audio, parallel=parallel)\n            \n            assert y.max() == 1\n            save_obj(x, f'x_f{fold}')\n            save_obj(y, f'y_f{fold}')\n\n\n    # ----------------------------------------------------------------------------------------------------\n    # 0.11 optuna\n    # ----------------------------------------------------------------------------------------------------\n\n    try:\n        from optuna.exceptions import ExperimentalWarning\n        warnings.filterwarnings(\"ignore\", category=ExperimentalWarning)\n    except:\n        pass\n\n    class Objective_optuna():\n        \"\"\" wrapper for function to be optimized by Optuna, whose purpose is to\n            customize the parameters generated by Trial, including handling exceptions to\n            parameters (fixed parameters for a trial).\n        \"\"\"\n        def __init__(self, study, hp_fixed, hp_optimize, hp_exceptions, optimize_fn):\n            self.study = study\n            self.hp_fixed = hp_fixed\n            self.hp_optimize = hp_optimize\n            self.hp_exceptions = hp_exceptions\n            self.optimize_fn = optimize_fn\n\n        def __call__(self, trial):\n            \"\"\" build hp and invoke training optimize_fn \"\"\"\n            study = self.study\n            hp_fixed = self.hp_fixed\n            hp_optimize = self.hp_optimize\n            hp_exceptions = self.hp_exceptions\n            optimize_fn = self.optimize_fn\n\n            # placeholder for code that implements specific rules, e.g., if encoder!='n/a' set norm=True\n\n            study.sampler = optuna.samplers.TPESampler()\n            if len(hp_exceptions) > 0:\n                fixed_params = self.hp_exceptions\n                partial_sampler = optuna.samplers.PartialFixedSampler(fixed_params, study.sampler)\n                study.sampler = partial_sampler\n\n            hp = {**hp_fixed, **hp_optimize(trial)}\n\n            return optimize_fn(hp)\n\n\n    def train_optuna(scenario, hp_fixed, hp_optimize, coach, previous_scenario='',\n                     duration=8, n_trials=100, trial_lot=10, hp_exceptions={}):\n        \"\"\" run various optuna trials for a model\n            Example:\n                >>> train(coach, previous_scenario='previous_scenario',\n                >>>       HP_FIXED_LINEAR, HP_OPTIMIZED_LINEAR,\n                >>>       duration=2, n_trials=100, trial_lot=10)\n        \"\"\"\n        SLEEP = 0\n        timer_ = Timer()\n        timer_.start(3600 * duration)\n        optuna.logging.set_verbosity(optuna.logging.WARNING)\n        if hasattr(coach, 'study'):\n            study = coach.study\n        else:\n            study = optuna.create_study(study_name=scenario, direction='maximize')\n            iteration = 0\n            if previous_scenario != '':\n                try:\n                    study = joblib.load(previous_scenario)\n                    iteration = len(study.trials)\n                except:\n                    print(f'Unable to load scenario {previous_scenario}')\n\n            coach.study = study\n            coach.iteration += iteration\n        objective = Objective_optuna(study, hp_fixed, hp_optimize, hp_exceptions, coach)\n        for i in range(n_trials // trial_lot):\n            study.optimize(objective, n_trials=trial_lot, timeout=3600, gc_after_trial=True)\n            gc.collect()\n\n            study.trials_dataframe().to_csv(f'./trials {scenario}.csv')\n            joblib.dump(study, f'./study {scenario}.dump'.format(scenario))\n            coach.save(scenario)\n            if not timer_.lap():\n                break\n            time.sleep(SLEEP)\n\n        try:\n            optuna.visualization.plot_optimization_history(study).show()\n            optuna.visualization.plot_param_importances(study).show()\n            optuna.visualization.plot_slice(study).show()\n        except Exception:\n            pass\n\n\n    # ----------------------------------------------------------------------------------------------------\n    # 0.12 feature importance\n    # ----------------------------------------------------------------------------------------------------\n\n    def predict_folds_(models, x, folds):\n        \"\"\" predict multiple folds and consolidate \"\"\"\n        y_pred = np.empty(folds.shape[0])\n        for fold in range(len(models)):\n            filter_fold = folds == fold\n            y_pred_fold = models[fold](x[fold]).squeeze()\n            y_pred[filter_fold] = y_pred_fold\n        return y_pred\n\n\n    def feature_importance_permutation(model_spec, x, y_true, folds, scorer_fn=None, n=1):  # FG: pending\n        \"\"\" supports only one model or ensemble at a time \"\"\"\n        if not isinstance(model_spec, list):\n            model_spec = [model_spec]\n        if not isinstance(model_spec[0][SPEC_MODEL_FOLDER], str):\n            model_spec = [(x[SPEC_MODEL_NAME], get_folder(x[SPEC_MODEL_NAME], x[1]), *x[SPEC_MODEL_NAME + 1:]) for x in model_spec]\n        results = []\n        print('Calculate feature importance...')\n        models = load_models(model_spec)\n        y_pred = predict_folds_(models, x, folds)\n        baseline = scorer_fn(y_true, y_pred)\n        print('baseline:', baseline)\n\n        for feature in tqdm(range(len(x_columns))):\n            saved_feature = x[feature].copy()\n            for i in range(n):\n                np.random.shuffle(x[feature])\n                y_pred = predict_folds_(models, x, folds)\n                score = scorer_fn(y_true, y_pred)\n                results.append({'feature': x_columns[feature], 'iteration': i, 'diff': score - baseline})\n            x[feature] = saved_feature\n            gc.collect()\n\n        result = pd.DataFrame(results).set_index(['feature', 'iteration']).unstack('iteration')\n        columns = [f'mae_{i}' for i in range(n)]\n        result.columns = columns\n        result['mean'] = result[columns].mean(axis=1)\n        result['share'] = result['mean'] / result['mean'].sum() * 100\n        result['std'] = result[columns].std(axis=1)\n        result['var 95%'] = result['std'] * 1.96 / result['mean'] * 100\n        result = result.sort_values(by='mean', ascending=False)\n        return result\n\n    # ----------------------------------------------------------------------------------------------------\n    # 0.13 analysis\n    # ----------------------------------------------------------------------------------------------------\n\n    def print_scores(scores, name='', precision=4, note=''):\n        target_scores = [x for x in scores]\n        m = np.mean(target_scores)\n        detail = ', '.join([f'{x:.{precision}f}' for x in target_scores])\n        print(f'{name:<7}:  {m:.{precision}f} [{detail}] {note}')\n\n\n    def copy_oofs(names, in_folder, val_samples):\n        \"\"\" copy oofs from one predict folder to another, inlcuidng only some models \"\"\"\n        for val_sample in val_samples:\n            for name in names:\n                path = f'{in_folder}/{val_sample}/oof_{name}_f0.pkl'\n                if os.path.exists(path):\n                    if not os.path.exists(f'{val_sample}'):\n                        os.makedirs(f'{val_sample}')\n                    shutil.copy(path, f'{val_sample}/oof_{name}_f0.pkl')\n\n    # ----------------------------------------------------------------------------------------------------\n    # 0.14 ensembles\n    # ----------------------------------------------------------------------------------------------------\n\n    t_cols = [f't{i}' for i in RANGE_BIRDS]\n    p_cols = [f'p{i}' for i in RANGE_BIRDS]\n\n    def names_to_idxs(selected, names):\n        \"\"\" return idx for selected in the order present in selected \"\"\"\n        name_to_idx = dict(zip(names, range(len(names))))\n        return np.array([name_to_idx[x] for x in selected if x in name_to_idx])\n\n    def get_idx_selected(selected, all):\n        \"\"\" returns the indices of `all` that are in selected \"\"\"\n        return np.array([i for i, item in enumerate(all) if item in selected])\n\n\n    def ensemble_oof(coef, oof, names):\n        w, idxs = [], []\n        if isinstance(coef, list):\n            coef = {x:1 for x in coef}\n        for name, factor in coef.items():\n            w.append(factor)\n            idxs.append(names.index(name))\n        w = np.array(w).reshape((-1,1,1))\n        e = (oof[idxs] * w).sum(axis=0) / w.sum(axis=0)\n        # e = e / e.sum(axis=-1, keepdims=True)\n        return e\n\n\n    def scores_ensemble(coef, oof, names, target):\n        e_best = ensemble_oof(coef, oof, names)\n        score_best = auc_multi(e_best, target)\n        idxs = names_to_idxs(coef.keys(), names)\n        e_mean = oof[idxs].mean(axis=0)\n        score_mean = auc_multi(e_mean, target)\n        if 'scores' in globals():\n            score_base = np.mean([scores[x] for x in coef.keys()])\n        else:\n            score_base = 0\n        return score_best, score_mean, score_base\n\n\n    class Ensemble_params:\n        def __init__(self, df, n=0, kappa=True):\n            self.results = []\n            self.count = n\n            self.best = np.inf\n            self.reference = 1.006589\n\n            # prepare cartesians\n            self.az = df.az.values\n            self.ze = df.ze.values\n\n            models = [x[4:] for x in df.columns if 'paz_' in x]\n            xyz, wk = [], []\n            for m in models:\n                xyz_model = polars_to_cartesians(df[[f'paz_{m}', f'pze_{m}']].values)\n                if kappa:\n                    w_model = df[f'kappa_{m}'].values.reshape((-1, 1))\n                    xyz_model = xyz_model * w_model\n                    wk.append(w_model)\n                #cols += [f'x_{m}', f'y_{m}', f'z_{m}']\n                xyz.append(xyz_model)\n            self.xyz = np.stack(xyz, axis=1)  # n, models, 3\n            self.wk = np.concatenate(wk, axis=1)    # n, models, 1\n\n\n        def __call__(self, hp):\n            self.count += 1\n\n            # ensemble pred\n            w = np.array(list(hp.values())).reshape((1, -1))\n            w = self.wk * w\n            w_sum = np.sum(w, axis=1)\n            w = w / w_sum.reshape((-1,1))\n            w_3d = np.repeat(np.expand_dims(w, 2), 3, axis=2)\n            xyz = (self.xyz * w_3d).sum(axis=1)\n            paz, pze = cartesians_to_polars(xyz)\n\n            # calculate score\n            score = ad_polar(self.az, self.ze, paz, pze)\n            delta = score - self.reference\n            self.results.append(dict(**hp, ad=score, delta=delta))\n            if score < self.best:\n                str_hp = [k for k, v in hp.items() if v != 0]\n                print(f'{self.count:>3}: {delta:8>.6f} -> {score:.6f} \\t {str_hp}')\n                self.best = score\n            return score\n\n        @property\n        def result(self):\n            return pd.DataFrame(self.results).sort_values('ad', ascending=True)\n\n\n    def gaussian_kernel(l=5, sig=1.):\n        \"\"\" https://stackoverflow.com/questions/29731726/how-to-calculate-a-gaussian-kernel-matrix-efficiently-in-numpy\n            creates 1D gaussian kernel with side length `l` and a sigma of `sig`\n        \"\"\"\n        ax = np.linspace(-(l - 1) / 2., (l - 1) / 2., l)\n        gauss = np.exp(-0.5 * np.square(ax) / np.square(sig))\n        return gauss / np.sum(gauss)\n\n    # -------------------\n    ## means\n    # -------------------\n\n    def ensemble_mean(selected, names, oof, target, ws=None, temperature=None):\n        \"\"\" Simple mean of various models. Also generates mean per fold \"\"\"\n        idxs = names_to_idxs(selected, names)\n        if not temperature:\n            oof_ensemble = oof[idxs].mean(axis=0)\n        else:\n            if temperature > 1 and oof[idxs].min() < 0:\n                oof = sigmoid(oof)\n            oof_ensemble = np.power(np.power(oof[idxs], temperature).mean(axis=0), 1/temperature)\n        score = auc_multi(oof_ensemble, y_true)\n        print(f'{\"mean\":<6}: {score:.4f}')\n        return score\n\n    # -------------------\n    ## best\n    # -------------------\n\n    from itertools import combinations\n\n    def ensemble_best(names, scores_detail, n=5):  # PENDING\n        \"\"\" ensembles up to `n` models by picking the best score for each row.\n            Use to pick models for L2 model.\n        \"\"\"\n        all_options, best = dict(), dict()\n        best_option = None\n        best_score = np.inf\n        scores_detail_a = scores_detail.values\n        names_a = np.array(names)\n        for i in range(2, n+1):\n            options = list(combinations(range(len(names)), i))\n            for option in tqdm(options):\n                cols = list(option)\n                score = auc_multi(oof, y_true)\n                score = scores_detail_a[:, cols].min(axis=1).mean()\n                option_names = names_a[cols]\n                all_options[tuple(option_names)] = score\n                if score < best_score:\n                    best_score = score\n                    best_option_i = option_names\n                    print(f'{score:.4f} - {list(option_names)}')\n            best[i] = (best_score, best_option_i)\n        return best, all_options\n\n    # -------------------\n    ## logistic regression\n    # -------------------\n\n    def ensemble_regression(names, oof, y_true, fit_intercept=True, n=5):\n        \"\"\" ensemble using regression. Reports 3 scores: \n            - best based on correlation\n            - mean of all models used in correlation\n            - mean of models indicated by correlation\n        \"\"\"\n        LIMIT_ENSEMBLE = 2\n        COEF_THRESHOLD_MAX = 1e7\n        selected_models = names\n       \n        start_time = time.time()\n        coefs, results = [], []\n        best_base_score = baseline_score\n\n        # run regression\n        names_ = selected_models.copy()\n        print(f'{names_}: x{len(names_)}')\n        best_score = np.inf\n        names_idx = None\n        start_names = names_\n        M, B, C = oof.shape\n        yt = y_true.reshape(-1)\n        yp = oof.reshape((M, -1)).T\n        for step in range(len(start_names)):\n            reg = LinearRegression(fit_intercept=fit_intercept, positive=True)\n            if names_idx is not None:\n                yp = yp[:, names_idx]\n            reg.fit(yp, yt)\n            gc.collect()\n            coef, coef_all, skip = dict(), [], []\n            # print(reg.coef_)\n            for i, name in enumerate(names_):\n                if reg.coef_[i] > COEF_THRESHOLD_MAX:\n                    skip.append(name)\n                if reg.coef_[i] != 0:\n                    coef[name] = reg.coef_[i]\n                coef_all.append(reg.coef_[i])\n            count = len(coef)\n            coef_s = pd.Series(coef)\n            coef_s = coef_s.sort_index(key=lambda x: x.str[:4] + x.str[4:].str.zfill(4))\n            coef_all = np.array(coef_all).reshape((-1, 1))\n            coef['intercept'] = reg.intercept_\n            #print(f'--- {coef}')\n\n            pred_b = (yp @ coef_all).squeeze() + coef['intercept']\n            pred_ma = np.mean(yp, axis=1)\n            idxs = get_idx_selected(coef, names_)\n            pred_m = np.mean(yp[:, idxs], axis=1)\n            \n            ensemble_best = auc_multi(pred_b.reshape((-1, B, C)), y_true)\n            ensemble_mean = auc_multi(pred_m.reshape((-1, B, C)), y_true)\n            ensemble_meana = auc_multi(pred_ma.reshape((-1, B, C)), y_true)\n            results.append([step, count, ensemble_best, ensemble_mean, ensemble_meana,\n                            ensemble_best - best_base_score, ensemble_mean - best_base_score, \n                            ensemble_meana - best_base_score, \n                            coef, reg.intercept_, None, list(coef_s.index)])\n            coefs.append({'best': ensemble_best, 'mean': ensemble_mean, **coef})\n\n            if ensemble_best <= best_score:\n                best_score = ensemble_best\n                best_idx = len(results) - 1\n\n            if len(skip) > 0:\n                to_remove = skip\n            else:\n                # remove model with smallest absolute weight\n                to_remove = coef_s.sort_values().index[0]\n            print(f'{step:>2} ({count:>2})  {ensemble_best:.6f}, {ensemble_mean:>9.6f}, {ensemble_meana:>9.6f} -> {to_remove}')\n            names_idx = [names_.index(x) for x in coef_s.index if x != to_remove]\n            names_ = [x for x in coef_s.index if x != to_remove]\n            if count <= LIMIT_ENSEMBLE:\n                break\n\n        column_names = ['step', 'count', 'best', 'mean', 'meana', 'delta_best', \n                        'delta_mean', 'delta_meana', 'coefs', 'intercept', 'selected', 'models']\n        results_df = pd.DataFrame(results, columns=column_names)\n        coef_df = pd.DataFrame(coefs).T\n        coef_df.to_csv(f'coef_df reg.csv')\n        results_df.to_csv(f'results_df reg.csv')\n\n        # summarize results\n        best_idxs = results_df.best.argsort().values\n        for idx in best_idxs[:n]:\n            show_ensemble_result(results_df, idx, len(start_names), best_base_score)\n        \n        return results_df, coef_df\n\n\n    def show_ensemble_result(results_df, idx, initial_len='', best_base_score=None):\n        if best_base_score is None:\n            best_base_score = baseline_score if 'baseline_score' in globals() else 0.30\n        results_df.iloc[idx][-2] = 1\n        coef = results_df.iloc[idx][-4]\n        coef_s = pd.Series(coef)\n        final_names = [x for x in coef.keys() if x != 'intercept']\n\n        ensemble_best, ensemble_mean, ensemble_meana = results_df.iloc[idx][2:5]\n        end_time = time.time()\n        print(f\"\\n{'-'*20} x{len(final_names)} / {initial_len} = \"\n              f\"{ensemble_best:.6f} {ensemble_mean:>.6f} {ensemble_meana:>.6f}  |  \"\n              f\"{ensemble_best - best_base_score:.6f} {ensemble_mean - best_base_score:>.6f}\"\n              f\"{ensemble_mean - best_base_score:>.6f}\")\n        coef_s = coef_s.sort_values(ascending=False)\n        print(coef_s)\n        print()\n        print(sort_names(final_names))\n        print()\n        print(coef_s.to_dict())\n\n    # -------------------\n    ## Optimize\n    # -------------------\n\n    def auc_multi(preds, labels):\n        \"\"\" Version of macro-averaged ROC-AUC score that ignores all classes that have no true positive labels. \"\"\"\n        scored_cols = labels.sum(axis=0) > 0\n        return sklearn.metrics.roc_auc_score(\n            labels[:, scored_cols], preds[:, scored_cols], average='macro')\n\n\n    def opt_ensemble_models_(w_in):\n        \"\"\" optimization function. In addition to weights, which are provided as input,\n            uses global variables to compute metric.\n        \"\"\"\n        wt = w_in.sum()\n        w = w_in.reshape((-1, 1, 1))\n        pred = (oof_ensemble * w).sum(axis=0) / wt\n        score = auc_multi(pred, _opt_target)\n        # print(score, oof_ensemble.shape, w.shape)\n        return -score\n\n    def ensemble_optimize(selected, names, oof, target, n=5, ws=None, method=None):\n        \"\"\" find best combination of models using scipy optimize function. \n            Once a solution is found, removes the model with the lowest coeficient and\n            tries to find a better solution without it.\n            Tends to be better than optuna and worse than regression. It can be used in\n            cases in which regression can't be used (the scoring metric desn't align with\n            regression).\n        \"\"\"\n        global oof_ensemble, _opt_target\n        \n        best_base_score, best, best_score =  baseline_score, [], -np.inf\n        step, coefs, results = 0, [], []\n        _opt_target = target\n        _opt_weights = ws\n        \n        initial_names = selected.copy()\n        print(f'{initial_names}: x{len(initial_names)}')\n        while len(selected) > 0:\n            t1 = time.time()\n            # prepare data\n            idxs = names_to_idxs(selected, names)\n            oof_ensemble = oof[idxs]\n\n            # run optimization\n            dim = oof_ensemble.shape[0]\n            guess = np.ones(dim)\n            bounds = [(0., 1.)] * dim\n            opt = optimize.minimize(opt_ensemble_models_, guess, bounds=bounds, method=method)\n            score = -opt.fun\n            # print results\n            coef, coef_all = dict(), []\n            for name, coef_ in zip(selected, opt.x):\n                if coef_ > 0:\n                    coef[name] = coef_\n                coef_all.append(coef_)\n            count = len(coef)\n            coef_all = np.array(coef_all).reshape((-1, 1, 1))\n            coef_s = pd.Series(coef)\n            if score > best_score:\n                best_score = score\n                best = coef\n            \n            # recalculate scores\n            #print(oof_ensemble.shape, coef_all.shape)\n            pred_b = (oof_ensemble * coef_all.reshape((-1, 1, 1))).sum(axis=0)\n            pred_ma = np.mean(oof_ensemble, axis=0)\n            idxs = get_idx_selected(coef, selected)\n            pred_m = np.mean(oof_ensemble[idxs], axis=0)\n            \n            # score = auc_multi(oof, y_true)\n            ensemble_best = auc_multi(pred_b, _opt_target)\n            if np.abs(score - ensemble_best) > 1e-6:\n                print('Score in opt != auc:', score, ensemble_best)\n            ensemble_mean = auc_multi(pred_m, _opt_target)\n            ensemble_meana = auc_multi(pred_ma, _opt_target)\n            results.append([step, count, ensemble_best, ensemble_mean, ensemble_meana,\n                            ensemble_best - best_base_score, ensemble_mean - best_base_score, \n                            ensemble_meana - best_base_score, \n                            coef, None, None, list(coef.keys())])\n            coefs.append({'best': ensemble_best, 'mean': ensemble_mean, **coef})\n            \n            # reduce number of selected models and try again. Remove models with insignificant coeficients.\n            # if none exists, remove the model with the lowest coeficient.\n            size = len(coef)\n            coef_s = coef_s.sort_values(ascending=False)\n            coef_s = coef_s[coef_s > 1e-4]\n            selected = list(coef_s.keys())\n            if len(selected) == size:\n                selected = selected[:-1]\n            duration = time.time() - t1\n            print(f'{step:>2} ({count:>2})  {ensemble_best:.6f}, {ensemble_mean:>9.6f}, {ensemble_meana:>9.6f}   {duration:,.0f} s')\n            step += 1\n\n        column_names = ['step', 'count', 'best', 'mean', 'meana', 'delta_best', \n                        'delta_mean', 'delta_meana', 'coefs', 'intercept', 'selected', 'models']\n        results_df = pd.DataFrame(results, columns=column_names)\n        coef_df = pd.DataFrame(coefs).T\n        coef_df.to_csv(f'coef_df opt.csv')\n        results_df.to_csv(f'results_df opt.csv')\n\n        # summarize results\n        best_idxs = results_df.best.argsort().values[::-1]\n        for idx in best_idxs[:n]:\n            show_ensemble_result(results_df, idx, len(names), best_base_score)\n\n        # print()\n        # print('-'*40, 'best', len(best), '-'*40)\n        # print(f'{len(best):>2}: {best_score:.6f}  {best_score - baseline_score:.6f}')\n        # print(f'{sorted(list(best.keys()))} \\n{best}')\n        return results_df, coef_df\n    # -------------------\n    ## Hill climb\n    # -------------------\n\n    def is_large_model(name):\n        if 'eca' in name:\n            return True\n        return False\n        \n\n    def count_large_models(names):\n        return sum(1 for x in names if is_large_model(x))\n\n\n    def ensemble_hill_climb(candidates, names, oof, target, temperature=None, fixed=None, n=12):\n        \"\"\" start by best model and add a model at a time (the one with the best improvement)\n            until we stop improving.\n        \"\"\"\n        \n        # get candidates\n        scores_ = {name: score_ for name, score_ in scores.items() if name in candidates}\n        temp = pd.Series(scores_, name='score').to_frame().reset_index().sort_values('score', ascending=False)\n        temp.rename(columns={'index': 'name'}, inplace=True)\n        best = temp.name.values[0]\n        if fixed is None:\n            selected = [best]\n            best_score = temp.score.values[0]\n        else:\n            selected = fixed\n            best_score = np.max([scores_[x] for x in fixed])\n\n        if temperature > 1 and oof.min() < 0:\n            oof = sigmoid(oof)\n        # select\n        candidates = [x for x in candidates if x not in selected]\n        while len(selected) < n:\n            has_improved = False\n            for candidate in candidates:\n                idxs = names_to_idxs(selected + [candidate], names)\n                if not temperature:\n                    oof_ensemble = oof[idxs].mean(axis=0)\n                else:\n                    oof_ensemble = np.power(np.power(oof[idxs], temperature).mean(axis=0), 1/temperature)\n                score = auc_multi(oof_ensemble, target)\n                if score > best_score:\n                    best_candidate = candidate\n                    best_score = score\n                    has_improved = True\n            if not has_improved:\n                print('finished improving')\n                break\n            selected.append(best_candidate)\n            sorted_selected = sort_names(selected)\n            print(f'{len(selected):>2}: {best_score:.6f}  {best_score - baseline_score:.6f}  - {sorted_selected}')\n            \n        # print('-' * 40)\n        # print(f'{len(selected):>2}: {best_score:.6f}  {best_score - baseline_score:.6f}  - {sorted_selected}')\n        return selected\n\n\n    # -------------------\n    ## Optuna\n    # -------------------\n\n    class Objective_weights:\n        def __init__(self, target, oof, hp, reference, n=0, short_progress=0):\n            self.results = []\n            self.count = n\n            self.best_ = -np.inf\n            self.target = target\n            self.oof = oof\n            self.hp = hp\n            self.reference = reference\n            self.short_progress = short_progress\n\n        def __call__(self, trial):\n            self.count += 1\n\n            # ensemble pred\n            params = self.hp(trial)\n            w = np.array(list(params.values())).reshape((-1, 1, 1))\n            pred = (self.oof * w).sum(axis=0)\n            score = auc_multi(pred, self.target)\n            delta = score - self.reference\n            self.results.append(dict(**params, c=score, delta=delta))\n            if score > self.best_:\n                str_hp = [k for k, v in params.items() if v != 0]\n                if self.short_progress and len(str_hp) > self.short_progress:\n                    str_hp = len(str_hp)\n                sys.stdout.write(f'{\" \"*80}\\r')\n                sys.stdout.flush()\n                print(f'{self.count:>3}: {delta:8>.6f} -> {score:.6f} \\t {str_hp}')\n                self.best_ = score\n            else:\n                text = f'\\r{self.count:>3}: {delta:8>.6f} -> {score:.6f}        '\n                sys.stdout.write(text)\n            return score\n\n        @property\n        def result(self):\n            return pd.DataFrame(self.results).sort_values('c', ascending=False)\n            \n        def best(self):\n            coef = self.result.sort_values('c', ascending=False).head(1).to_dict(orient='records')[0]\n            coef = {k: v for k, v in coef.items() if v != 0}\n            score = coef.pop('c')\n            delta = coef.pop('delta')\n            return coef, score\n\n\n    def ensemble_optuna(candidates, names, target, oof, hp, reference, initial=None, sampler=None,\n                        scenario='optuna', n_results=5, study=None, previous_study=None, \n                        n_trials=2_000, duration=3600, short_progress=0):\n        optuna.logging.set_verbosity(optuna.logging.WARNING)\n        \n        # create or load study\n        if not study:\n            if previous_study is not None:\n                try:\n                    study = joblib.load(previous_study)\n                except:\n                    print(f\"Unable to load scenario '{previous_study}'\")\n                    study = optuna.create_study(study_name=scenario, direction='maximize')\n            else:\n                study = optuna.create_study(study_name=scenario, direction='maximize')\n\n        # handle fixed values\n        study.sampler = sampler if sampler is not None else optuna.samplers.TPESampler()\n        if initial is not None:\n            study.enqueue_trial(initial)   \n\n        # optimize\n        idx_candidates = names_to_idxs(candidates, names)\n        oof_candidates = oof[idx_candidates]\n        objective = Objective_weights(\n            target, oof_candidates, hp, reference, n=len(study.trials), short_progress=short_progress)\n        try:\n            study.optimize(objective, n_trials=n_trials, timeout=duration, gc_after_trial=False)\n        except KeyboardInterrupt:\n            print(f'\\nInterrupted by request')\n\n        best_coef, best_score = objective.best()\n        \n        # output results\n        objective.result.to_csv(f'./results {scenario}.csv')\n        study.trials_dataframe().to_csv(f'./trials {scenario}.csv')\n        joblib.dump(study, f'study {scenario}.dump')\n        \n        score_best, score_mean, score_base = scores_ensemble(best_coef, oof, names, target)\n        delta = best_score - reference\n        print(f'\\n*** best: {delta:8>.6f} -> {score_best:.6f}, {score_mean:.6f}, {score_base:.6f}')\n        print(f'\\n{best_coef}')\n\n        if n_results > 0:\n            # visualize and store results\n            print('best:\\n')\n            display(objective.result.head(n_results))\n        \n        return study, objective, best_coef, best_score\n\n    # ----------------------------------------------------------------------------------------------------\n    # 0.20 Visualize\n    # ----------------------------------------------------------------------------------------------------\n\n    def plot_audio_speech(file, play_audio=True, th=.5, no_speech=False):\n        global speech_filter, model, get_speech_timestamps\n        \n        if 'get_speech_timestamps' not in globals():\n            torch.set_num_threads(1)\n            model, (get_speech_timestamps, _, read_audio, _, _) = \\\n                torch.hub.load(repo_or_dir='snakers4/silero-vad', model='silero_vad')\n            if not no_speech:\n                speech_filter = Speech_filter()        \n        \n        audio, sr = torchaudio.load(file)\n        if sr != SR:\n            resampler = torchaudio.transforms.Resample(orig_freq=sr, new_freq=SR)\n            audio = resampler(audio)\n        audio = audio[0]\n        \n        name = '/'.join(file.split('/')[-2:]).split('.')[0]\n        if not no_speech:\n            start, end = speech_filter(audio, th=th)\n            speech_timestamps = get_speech_timestamps(audio, model, sampling_rate=SR, threshold=th)\n            segmentation = np.zeros_like(audio)\n            for st in speech_timestamps:\n                segmentation[st['start']: st['end']] = 20\n        else:\n            start, end = 0, len(audio)\n        if start is not None:\n            print(file, start/SR, end/SR)\n        else:\n            print(file)\n        #if mode != 2: continue\n        chunk_len = 0.1\n        power = audio ** 2\n        chunk = int(chunk_len * SR)\n        pad = int(np.ceil(len(power) / chunk) * chunk - len(power))\n        power = np.pad(power, (0, pad))\n        power = power.reshape((-1, chunk)).sum(axis=1)\n        \n        fig = plt.figure(figsize=(24, 3))\n        fig.suptitle(f'{file}')\n        \n        t = np.arange(len(power)) * chunk_len\n        plt.plot(t, 10 * np.log10(power), 'b')\n\n        if start is not None:\n            plt.axvline(x=start/SR, color='g')\n        if end is not None:\n            plt.axvline(x=end/SR, color='g')\n\n        # plot time lines\n        ax1 = plt.gca()\n        duration = int(t.shape[0]*.1)\n        interval = 5 if duration <=  5 * 50 else 20\n        xi = [x for x in range(0, duration, interval)]\n        ax1.set_xticks(xi)\n        ax1.set_xticklabels(xi, rotation=90)\n        for x in xi:\n            ax1.axvline(x=x, color='black', linestyle='--', linewidth=1)\n\n        if not no_speech:\n            t = np.arange(len(segmentation)) / SR\n            plt.plot(t, segmentation, 'r')        \n        plt.show()\n\n        if play_audio:\n            display(Audio(audio, rate=SR, autoplay=True))","metadata":{"trusted":true,"_kg_hide-input":true,"jupyter":{"source_hidden":true},"execution":{"iopub.status.busy":"2025-06-03T01:30:04.04034Z","iopub.status.idle":"2025-06-03T01:30:04.040611Z","shell.execute_reply.started":"2025-06-03T01:30:04.04049Z","shell.execute_reply":"2025-06-03T01:30:04.0405Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"models_config = [\n('ebs.426', [0,1,2,3,4], 'b5-models-1-8', {'accumulation_steps': 1, 'add_position': False, 'adv_eps': 0.01, 'adv_lr': 0.005, 'adv_th': 0.3, 'aug': {'CoarseDropout': (0.375, 0.375, 1, 0.7), 'Flip': 0.5, 'audio': False, 'mixup': 1, 'volume': (0.3333333333333333, 3)}, 'awp': False, 'background_noise_cache_size': 1000, 'background_noise_max_usage': 6, 'background_noise_prob': 0.5, 'background_noise_reference': True, 'betas': (0.9, 0.999), 'check_val_every_n_epoch': 1, 'curation_mode': 'speech', 'dataset': 'h5-6', 'dataset_val': 'val-128-256-m1-4', 'drop_path_rate': None, 'duration': {'hours': 11, 'minutes': 55}, 'encoder': 'tf_efficientnetv2_s', 'end_lr': 1e-06, 'epochs': 50, 'eps': 1e-08, 'extension': '.ogg', 'files_edges': '/kaggle/input/b5-cache/files_edges speech 5', 'filter': (False, 0), 'fold_col': 'author', 'gem_p': 1.8, 'head_dropout': 0.0, 'img_dim': (128, 256), 'img_duration': 5, 'in_chans': 1, 'label_smoothing': 0.0, 'log_every_n_steps': 50, 'loss_schedule': {0: 'focal_volodymyr'}, 'lr': [(1e-05, 0.001), (3, 0.00025)], 'mask_non_labels': True, 'max_grad_norm': 10, 'mode': 'max', 'model_name': 'SpecNetImg', 'monitor': 'val_auc', 'montage': (0, 3), 'montage_cache_size': 10, 'n_folds': 5, 'n_preload_species': 15, 'norm_audio': False, 'num_workers': 6, 'optimizer': 'AdamW', 'out_indices': 2, 'precision_schedule': {0: 'float32'}, 'prefetch_factor': 2, 'previous_dataset': 'add-h5-6', 'pseudo_label_zero_th': 0.1, 'sample_unlabeled_prob': 0.0, 'saved_optimizer': None, 'saved_scheduler': None, 'schedule_type': 'multi_lr', 'seed': 10, 'spec': None, 'spec_norm': False, 'start_lr': None, 'step_scheduler_after': 'step', 'swa': None, 'train_batch_size': 64, 'train_mode': 'random', 'train_primary_th': 0.5, 'train_size': 28000, 'train_unlabeled_dataset': 'h5-unlabeled-2', 'unlabeled_primary_th': 0.5, 'unlabeled_weight': 1, 'use_mask': False, 'val_batch_size': 64, 'warmup_share': 0.02, 'weight_decay': 1e-06, 'ws_power': 0.5, 'name': 'ebs.1', 'folds': [4], 'unlabeled_pseudo_preds': ['b5-data-pseudo-pred-v3/unlabeled pseudo pred v2'], 'montage_bird_prob': False, 'pretrained_encoder': '/kaggle/input/b5-pretrained-weights/tf_efficientnetv2_s_in21k_Pretrainversion1.pth'}),\n]\n\nROOT = '/kaggle/input'\nmodels_config, models_config_dict = clean_config(models_config)\n\nebs_426 = {'ebs.426_f0': 1/5, 'ebs.426_f1': 1/5, 'ebs.426_f2': 1/5, 'ebs.426_f3': 1/5, 'ebs.426_f4': 1/5}\n\nclass ARGS:\n    engine = 'onnx'  # opt onnx vino\n    lot_size = 90\n    num_workers = 3\n    batch_size = 4\n    ws = ebs_426\n    ensemble_logit = False\n    norm_pred = False\n        \nif ARGS.ws is not None:\n    ARGS.ws = {f'{x}_f0'if '_' not in x else x : w for x, w in ARGS.ws.items()}\n    ws_names = set(x.split('_')[0] for x in ARGS.ws.keys())\n    models_config = [x for x in models_config if x[0] in ws_names]","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-06-03T01:30:04.042367Z","iopub.status.idle":"2025-06-03T01:30:04.042763Z","shell.execute_reply.started":"2025-06-03T01:30:04.042571Z","shell.execute_reply":"2025-06-03T01:30:04.04259Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"test_path = '/kaggle/input/birdclef-2025/test_soundscapes/'\ntrain_path = '/kaggle/input/birdclef-2025/train_soundscapes'\ntest_files_ = glob.glob(f'{test_path}*')\nfolder = test_path if len(test_files_) > 1 else train_path\nrepeated_files = ['_'.join(x.split('_')[:-1]) for x in test_pred_df.row_id]\nfiles = []\nfor x in repeated_files:\n    if len(files) == 0 or x != files[-1]:\n        files.append(x)\nfiles = [f'{folder}/{x}.ogg' for x in files]\n\nprint(f'files: {len(files):,.0f}')\n\nif not IS_TEST:\n    import cpuinfo\n    print(cpuinfo.get_cpu_info()['brand_raw'])","metadata":{"trusted":true,"_kg_hide-input":true,"execution":{"iopub.status.busy":"2025-06-03T01:30:04.044039Z","iopub.status.idle":"2025-06-03T01:30:04.044465Z","shell.execute_reply.started":"2025-06-03T01:30:04.044254Z","shell.execute_reply":"2025-06-03T01:30:04.044271Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"if not SKIP:  # predict\n    t1 = time.time()\n    file_lots = [files[i : i+ARGS.lot_size] for i in range(0, len(files), ARGS.lot_size)]\n    models_dict, predict_fn = prepare_predict()\n    predict_fn = predict_fn_onnx_timeout_bits\n    spec_configs = split_spec_configs(ARGS.ws, models_config_dict)\n    print(spec_configs)\n    names = []\n    oof = []\n    for (spec, img_dim, img_duration, spec_norm, snipet), spec_model_names in spec_configs.items():\n        print('-'*20, 'spec', spec, img_dim, img_duration, spec_norm, snipet)\n        models = [models_dict[f'{name}'] for name in spec_model_names]\n        all_bits = [models_config_dict[name.split('_')[0]].bits for name in spec_model_names]\n        names += spec_model_names\n        preds = []\n        for lot in file_lots:\n            specs = prepare_data(lot, 4, verbose=0, norm=spec_norm, dim=img_dim, duration=img_duration,\n                                 snipet=snipet, spec=spec, bits=32)\n            chunk_size = math.ceil(len(specs) / ARGS.num_workers)\n            chunks = [specs[i : i+chunk_size] for i in range(0, len(specs), chunk_size)]\n            with concurrent.futures.ThreadPoolExecutor(max_workers=ARGS.num_workers) as executor:\n                chunk_preds = list(executor.map(predict_fn, chunks))\n            chunk_preds = np.concatenate(chunk_preds, axis=1)\n            print('chunk_preds', chunk_preds.shape)\n            preds.append(chunk_preds)\n        oof_spec = np.concatenate(preds, axis=1)\n        oof.append(oof_spec)\n\n    oof = np.concatenate(oof, axis=0)\n    duration = time.time() - t1\n    print(f'time: {duration:,.0f} s for {len(files):,.0f} files  -  all: {duration * 700 / len(files) / 60:,.0f} m')\n    print(oof.shape)\n    # del specs\n    gc.collect()\n\n    # ensemble\n    oof = sigmoid(oof)\n    fernando_probs = ensemble_oof(ARGS.ws, oof, names)","metadata":{"trusted":true,"_kg_hide-input":true,"jupyter":{"source_hidden":true},"execution":{"iopub.status.busy":"2025-06-03T01:30:04.045823Z","iopub.status.idle":"2025-06-03T01:30:04.046254Z","shell.execute_reply.started":"2025-06-03T01:30:04.046024Z","shell.execute_reply":"2025-06-03T01:30:04.046041Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# merge Volodymyr with Fernando\nvolodymyr_probs = test_pred_df[BIRDS].values.astype(np.float32)\nN_, F_ = volodymyr_probs.shape\nif ARGS.norm_pred:\n    print(\"Normalizing Preds by STD\")\n    volodymyr_probs = volodymyr_probs / volodymyr_probs.std(axis=1, keepdims=1)\n    fernando_probs = fernando_probs / fernando_probs.std(axis=1, keepdims=1)\nonly_probs = volodymyr_probs * 2/3 + fernando_probs * 1/3\ntest_pred_df[BIRDS] = only_probs.reshape((N_, F_))\nif not IS_TEST:\n    corr_vol_fer = np.corrcoef(volodymyr_probs.flatten(), fernando_probs.flatten())[0,1]\n    corr_vol_ensemble = np.corrcoef(volodymyr_probs.flatten(), only_probs.flatten())[0,1]\n    print(f'correlation Volodymyr / Fernando: {corr_vol_fer:.3f}')\n    print(f'correlation Volodymyr / ensemble: {corr_vol_ensemble:.3f}')\n    print(f'std Volodymyr: {volodymyr_probs.std():.4f}')\n    print(f'std Fernando:  {fernando_probs.std():.4f}')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-06-03T01:30:04.047564Z","iopub.status.idle":"2025-06-03T01:30:04.047957Z","shell.execute_reply.started":"2025-06-03T01:30:04.047763Z","shell.execute_reply":"2025-06-03T01:30:04.047781Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# -- Fernando end","metadata":{}},{"cell_type":"code","source":"def postprocessing(input_df, top=6):\n    print(\"Top\", top)\n    only_probs = input_df.iloc[:, 1:].values\n    N, F = only_probs.shape\n    only_probs = only_probs.reshape((N//12, 12, F))\n    # mean_ = np.mean(np.sort(only_probs, axis=1)[:, 8:], axis=1, keepdims=True)\n    mean_ = np.mean(np.sort(only_probs, axis=1)[:, -top:], axis=1, keepdims=True)\n    only_probs *= mean_\n    input_df.iloc[:, 1:] = only_probs.reshape((N, F))\n    return input_df","metadata":{"execution":{"iopub.status.busy":"2025-06-03T01:30:04.049825Z","iopub.status.idle":"2025-06-03T01:30:04.050254Z","shell.execute_reply.started":"2025-06-03T01:30:04.050024Z","shell.execute_reply":"2025-06-03T01:30:04.050043Z"},"papermill":{"duration":0.018099,"end_time":"2025-04-16T20:13:44.871809","exception":false,"start_time":"2025-04-16T20:13:44.85371","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"test_pred_df = postprocessing(test_pred_df, top=1)\n\nif not IS_TEST:\n    display(test_pred_df)\n    print(test_pred_df.iloc[:,1:].values.max(), test_pred_df.iloc[:,1:].values.min())\n    print(test_pred_df.iloc[:,1:].sum(axis=1))\n    plt.hist(test_pred_df.iloc[:,1:].values.max(axis=1), bins=30)\n    plt.show()","metadata":{"execution":{"iopub.status.busy":"2025-06-03T01:30:04.051921Z","iopub.status.idle":"2025-06-03T01:30:04.052315Z","shell.execute_reply.started":"2025-06-03T01:30:04.052111Z","shell.execute_reply":"2025-06-03T01:30:04.052142Z"},"papermill":{"duration":0.016011,"end_time":"2025-04-16T20:13:44.897044","exception":false,"start_time":"2025-04-16T20:13:44.881033","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Save Prediction","metadata":{"papermill":{"duration":0.009274,"end_time":"2025-04-16T20:13:45.312816","exception":false,"start_time":"2025-04-16T20:13:45.303542","status":"completed"},"tags":[],"_kg_hide-input":true}},{"cell_type":"code","source":"# sample_submission = pd.read_csv(\"/kaggle/input/birdclef-2023/sample_submission.csv\")\nassert set(sample_submission.columns) == set(test_pred_df.columns)\ntest_pred_df = test_pred_df[sample_submission.columns]\ntest_pred_df.to_csv(\"submission.csv\", index=False)","metadata":{"execution":{"iopub.status.busy":"2025-06-03T01:30:04.052906Z","iopub.status.idle":"2025-06-03T01:30:04.053211Z","shell.execute_reply.started":"2025-06-03T01:30:04.053036Z","shell.execute_reply":"2025-06-03T01:30:04.053047Z"},"papermill":{"duration":0.066031,"end_time":"2025-04-16T20:13:45.388411","exception":false,"start_time":"2025-04-16T20:13:45.32238","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null}]}