{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.11.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":91844,"databundleVersionId":11361821,"sourceType":"competition"},{"sourceId":14183563,"sourceType":"datasetVersion","datasetId":9042475},{"sourceId":687815,"sourceType":"modelInstanceVersion","isSourceIdPinned":true,"modelInstanceId":519622,"modelId":534001}],"dockerImageVersionId":31192,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"#submission notebook doesn't have internet access, so I simply copy pasted code: https://huggingface.co/bird-of-paradise/deepseek-mla\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom typing import Tuple\n\nimport math\n\n\n\ndef precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0):\n    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))\n    t = torch.arange(end, device=freqs.device)  # type: ignore\n    freqs = torch.outer(t, freqs).float()  # type: ignore\n    freqs_cis = torch.polar(torch.ones_like(freqs), freqs)  # complex64\n    return freqs_cis\n\ndef reshape_for_broadcast(freqs_cis: torch.Tensor, x: torch.Tensor):\n    ndim = x.ndim\n    assert 0 <= 1 < ndim\n    assert freqs_cis.shape == (x.shape[1], x.shape[-1])\n    shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)]\n    return freqs_cis.view(*shape)\n\n\ndef apply_rotary_emb(\n    xq: torch.Tensor,\n    xk: torch.Tensor,\n    freqs_cis: torch.Tensor,\n) -> Tuple[torch.Tensor, torch.Tensor]:\n    # Validate input dimensions\n    assert xq.shape[-1] == xk.shape[-1], \"Query and Key must have same embedding dimension\"\n    assert xq.shape[-1] % 2 == 0, \"Embedding dimension must be even\"\n\n    # Get sequence lengths\n    q_len = xq.shape[1]\n    k_len = xk.shape[1]\n    \n    # Use appropriate part of freqs_cis for each sequence\n    q_freqs = freqs_cis[:q_len]\n    k_freqs = freqs_cis[:k_len]\n    \n    # Apply rotary embeddings separately\n    # split last dimention to [xq.shape[:-1]/2, 2]\n    xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))\n    xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))\n    \n \n    # Reshape freqs for each\n    q_freqs = reshape_for_broadcast(q_freqs, xq_)\n    k_freqs = reshape_for_broadcast(k_freqs, xk_)\n    \n    # Works for both [bsz, seqlen, n_heads*head_dim] and [bsz, seqlen, n_heads, head_dim]\n    xq_out = torch.view_as_real(xq_ * q_freqs).flatten(xq.ndim-1) \n    xk_out = torch.view_as_real(xk_ * k_freqs).flatten(xk.ndim-1)\n\n    return xq_out.type_as(xq), xk_out.type_as(xk)\n\n\n\n\nclass MultiHeadLatentAttention(nn.Module):\n    \"\"\"\n        Multi-Head Latent Attention(MLA) Module As in DeepSeek_V2 pape\n        Key innovation from standard MHA:\n             1. Low-Rank Key-Value Joint Compression \n             2. Decoupled Rotary Position Embedding\n             \n    Args:\n        d_model:  Total dimension of the model.\n        num_head: Number of attention heads.\n        d_embed:  Embedding dimension\n        d_c:      K/V compression dimension\n        d_c1:     Q compression dimension\n        d_rotate: Dimension for Rotary Position Embedding\n        dropout:  Dropout rate for attention scores.\n        bias:     Whether to include bias in linear projections.\n        d_head:   Inferred from d_model//num_head\n    Inputs:\n        sequence: input sequence for self-attention and the query for cross-attention\n        key_value_state: input for the key, values for cross-attention\n    \"\"\"\n    def __init__(\n        self, \n        d_model,             # Infer d_head from d_model\n        num_head, \n        d_embed, \n        d_c, \n        d_c1, \n        d_rotate, \n        dropout=0.1, \n        bias=True,\n        max_batch_size=32,   # For KV cache sizing\n        max_seq_len=2048     # For KV cache sizing \n        ):\n        super().__init__()\n        \n        assert d_model % num_head == 0, \"d_model must be divisible by num_head\"\n        assert d_c < d_embed, \"Compression dim should be smaller than embedding dim\"\n        assert d_c1 < d_embed, \"Query compression dim should be smaller than embedding dim\"\n        \n        self.d_model = d_model\n        self.num_head = num_head\n        # Verify dimensions match up\n        assert d_model % num_head == 0, f\"d_model ({d_model}) must be divisible by num_head ({num_head})\"\n        self.d_head=d_model//num_head\n        self.d_embed = d_embed\n        self.d_c = d_c\n        self.d_c1 = d_c1\n        self.d_rotate = d_rotate\n        self.dropout_rate = dropout  # Store dropout rate separately\n\n        # Linear down-projection(compression) transformations\n        self.DKV_proj = nn.Linear(d_embed, d_c, bias=bias)\n        self.DQ_proj = nn.Linear(d_embed, d_c1, bias=bias)\n        \n        # linear up-projection transformations\n        self.UQ_proj = nn.Linear(d_c1, d_model, bias=bias)\n        self.UK_proj = nn.Linear(d_c, d_model, bias=bias)\n        self.UV_proj = nn.Linear(d_c, d_model, bias=bias)\n\n        # Linear RoPE-projection\n        self.RQ_proj = nn.Linear(d_c1, num_head*d_rotate, bias=bias)\n        self.RK_proj = nn.Linear(d_embed, d_rotate, bias=bias)\n        \n        # linear output transformations\n        self.output_proj = nn.Linear( d_model, d_model, bias=bias)\n\n        # Dropout layer\n        self.dropout = nn.Dropout(p=dropout)\n\n        # Initiialize scaler\n        self.scaler = float(1.0 / math.sqrt(self.d_head + d_rotate)) # Store as float in initialization\n\n        # Initialize C_KV and R_K cache for inference\n        self.cache_kv = torch.zeros(\n            (max_batch_size, max_seq_len, d_c)\n        )\n        self.cache_rk = torch.zeros(\n            (max_batch_size, max_seq_len, d_rotate)\n        )\n\n        # Initialize freqs_cis for RoPE\n        self.freqs_cis = precompute_freqs_cis(\n            d_rotate, max_seq_len * 2\n        )\n    \n\n    def forward(\n        self, \n        sequence, \n        key_value_states = None, \n        att_mask=None,\n        use_cache=False,\n        start_pos: int = 0\n    ):\n\n        \"\"\"\n        Forward pass supporting both standard attention and cached inference\n        Input shape: [batch_size, seq_len, d_model=num_head * d_head]\n        Args:\n            sequence: Input sequence [batch_size, seq_len, d_model]\n            key_value_states: Optional states for cross-attention\n            att_mask: Optional attention mask\n            use_cache: Whether to use KV caching (for inference)\n            start_pos: Position in sequence when using KV cache\n        \"\"\"\n        batch_size, seq_len, model_dim = sequence.size()\n        # prepare for RoPE\n        self.freqs_cis = self.freqs_cis.to(sequence.device)\n        freqs_cis = self.freqs_cis[start_pos : ]\n\n        # Check only critical input dimensions\n        assert model_dim == self.d_model, f\"Input dimension {model_dim} doesn't match model dimension {self.d_model}\"\n        if key_value_states is not None:\n            assert key_value_states.size(-1) == self.d_model, \\\n            f\"Cross attention key/value dimension {key_value_states.size(-1)} doesn't match model dimension {self.d_model}\"\n\n        # if key_value_states are provided this layer is used as a cross-attention layer\n        # for the decoder\n        is_cross_attention = key_value_states is not None\n\n        # Determine kv_seq_len early\n        kv_seq_len = key_value_states.size(1) if is_cross_attention else seq_len\n        \n        # Linear projections and reshape for multi-head, in the order of Q, K/V\n        # Down and up projection for query\n        C_Q = self.DQ_proj(sequence)     #[batch_size, seq_len, d_c1]\n        Q_state = self.UQ_proj(C_Q)      #[batch_size, seq_len, d_model]\n        # Linear projection for query RoPE pathway\n        Q_rotate = self.RQ_proj(C_Q)      #[batch_size, seq_len, num_head*d_rotate]\n\n\n        if use_cache:\n            #Equation (41) in DeepSeek-v2 paper: cache c^{KV}_t\n            self.cache_kv = self.cache_kv.to(sequence.device)\n\n            # Get current compressed KV states\n            current_kv = self.DKV_proj(key_value_states if is_cross_attention else sequence) #[batch_size, kv_seq_len, d_c]\n            # Update cache using kv_seq_len instead of seq_len\n            self.cache_kv[:batch_size, start_pos:start_pos + kv_seq_len] = current_kv\n            # Use cached compressed KV up to current position\n            C_KV = self.cache_kv[:batch_size, :start_pos + kv_seq_len]\n\n            #Equation (43) in DeepSeek-v2 paper: cache the RoPE pathwway for shared key k^R_t\n            assert self.cache_rk.size(-1) == self.d_rotate, \"RoPE cache dimension mismatch\"\n            self.cache_rk = self.cache_rk.to(sequence.device)\n            # Get current RoPE key\n            current_K_rotate = self.RK_proj(key_value_states if is_cross_attention else sequence) #[batch_size, kv_seq_len, d_rotate]\n            # Update cache using kv_seq_len instead of seq_len\n            self.cache_rk[:batch_size, start_pos:start_pos + kv_seq_len] = current_K_rotate\n            # Use cached RoPE key up to current position\n            K_rotate = self.cache_rk[:batch_size, :start_pos + kv_seq_len] #[batch_size, cached_len, d_rotate]\n            \n            \n            \"\"\"handling attention mask\"\"\"\n            if att_mask is not None:\n                # Get the original mask shape\n                mask_size = att_mask.size(-1)\n                cached_len = start_pos + kv_seq_len        # cached key_len, including previous key\n                assert C_KV.size(1) == cached_len, \\\n            f\"Cached key/value length {C_KV.size(1)} doesn't match theoretical length {cached_len}\"\n                \n                # Create new mask matching attention matrix shape\n                extended_mask = torch.zeros(\n                    (batch_size, 1, seq_len, cached_len),  # [batch, head, query_len, key_len]\n                    device=att_mask.device,\n                    dtype=att_mask.dtype\n                )\n                \n                # Fill in the mask appropriately - we need to be careful about the causality here\n                # For each query position, it should only attend to cached positions up to that point\n                for i in range(seq_len):\n                    extended_mask[:, :, i, :(start_pos + i + 1)] = 0  # Can attend\n                    extended_mask[:, :, i, (start_pos + i + 1):] = float('-inf')  # Cannot attend\n                    \n                att_mask = extended_mask\n        else:\n            # Compression projection for C_KV\n            C_KV = self.DKV_proj(key_value_states if is_cross_attention else sequence) #[batch_size, kv_seq_len, d_c]\\\n            # RoPE pathway for *shared* key\n            K_rotate = self.RK_proj(key_value_states if is_cross_attention else sequence)\n            \n\n        # Up projection for key and value\n        K_state = self.UK_proj(C_KV)               #[batch_size, kv_seq_len/cached_len, d_model]\n        V_state = self.UV_proj(C_KV)               #[batch_size, kv_seq_len/cached_len, d_model]\n\n        \n        Q_state = Q_state.view(batch_size, seq_len, self.num_head, self.d_head)\n\n        # After getting K_state from projection, get its actual sequence length\n        actual_kv_len = K_state.size(1)    # kv_seq_len or start_pos + kv_seq_len\n        # in cross-attention, key/value sequence length might be different from query sequence length\n        # Use actual_kv_len instead of kv_seq_len for reshaping\n        K_state = K_state.view(batch_size, actual_kv_len, self.num_head, self.d_head) \n        V_state = V_state.view(batch_size, actual_kv_len, self.num_head, self.d_head)\n\n\n        #Apply RoPE to query and shared key\n        Q_rotate = Q_rotate.view(batch_size, seq_len, self.num_head, self.d_rotate)\n        K_rotate = K_rotate.unsqueeze(2).expand(-1, -1, self.num_head, -1)  # [batch, cached_len, num_head, d_rotate]\n        Q_rotate, K_rotate = apply_rotary_emb(Q_rotate, K_rotate, freqs_cis=freqs_cis)\n\n\n        # Concatenate along head dimension\n        Q_state = torch.cat([Q_state, Q_rotate], dim=-1)  # [batch_size, seq_len, num_head, d_head + d_rotate]\n        K_state = torch.cat([K_state, K_rotate], dim=-1)  # [batch_size, actual_kv_len, num_head, d_head + d_rotate]\n\n\n        # Scale Q by 1/sqrt(d_k)\n        Q_state = Q_state * self.scaler\n        Q_state = Q_state.transpose(1, 2)  # [batch_size, num_head, seq_len, head_dim]\n        K_state = K_state.transpose(1, 2)  # [batch_size, num_head, actual_kv_len, head_dim]\n        V_state = V_state.transpose(1, 2)  # [batch_size, num_head, actual_kv_len, head_dim]\n\n    \n        # Compute attention matrix: QK^T\n        self.att_matrix = torch.matmul(Q_state, K_state.transpose(-1,-2)) \n    \n        # apply attention mask to attention matrix\n        if att_mask is not None and not isinstance(att_mask, torch.Tensor):\n            raise TypeError(\"att_mask must be a torch.Tensor\")\n\n        if att_mask is not None:\n            self.att_matrix = self.att_matrix + att_mask\n        \n        # apply softmax to the last dimension to get the attention score: softmax(QK^T)\n        att_score = F.softmax(self.att_matrix, dim = -1)\n    \n        # apply drop out to attention score\n        att_score = self.dropout(att_score)\n    \n        # get final output: softmax(QK^T)V\n        att_output = torch.matmul(att_score, V_state)\n        assert att_output.size(0) == batch_size, \"Batch size mismatch\"\n        assert att_output.size(2) == seq_len, \"Output sequence length should match query sequence length\"\n        \n            \n        # concatinate all attention heads\n        att_output = att_output.transpose(1, 2).contiguous().view(batch_size, seq_len, self.num_head*self.d_head) \n\n\n        # final linear transformation to the concatenated output\n        att_output = self.output_proj(att_output)\n\n        assert att_output.size() == (batch_size, seq_len, self.d_model), \\\n        f\"Final output shape {att_output.size()} incorrect\"\n\n        return att_output","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-19T08:22:36.114760Z","iopub.execute_input":"2025-12-19T08:22:36.115355Z","iopub.status.idle":"2025-12-19T08:22:39.472047Z","shell.execute_reply.started":"2025-12-19T08:22:36.115328Z","shell.execute_reply":"2025-12-19T08:22:39.471442Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import pandas as pd\n#import matplotlib.pyplot as plt\n#import seaborn as sns\nimport pickle\nimport os\nfrom kaggle_secrets import UserSecretsClient\nuser_secrets = UserSecretsClient()\nos.environ['WANDB_API_KEY'] = user_secrets.get_secret(\"WANDB_API_KEY\")\n\ntrain_df = pd.read_csv(\"/kaggle/input/birdclef-2025/train.csv\").drop(columns = ['url', 'license'])\ntaxonomy_df = pd.read_csv(\"/kaggle/input/birdclef-2025/taxonomy.csv\")\n\ntrain_df = pd.merge(\n                train_df,\n                taxonomy_df[['primary_label', 'class_name']],\n                how = 'left',\n                on = ['primary_label']\n            )","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-19T08:22:44.784686Z","iopub.execute_input":"2025-12-19T08:22:44.785068Z","iopub.status.idle":"2025-12-19T08:22:46.698192Z","shell.execute_reply.started":"2025-12-19T08:22:44.785046Z","shell.execute_reply":"2025-12-19T08:22:46.697563Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"Dataset curation","metadata":{}},{"cell_type":"code","source":"#Well, I think it's a good practice nowadays, to explicitly write, where LLMs were used\nfrom dataclasses import dataclass, field\nfrom typing import Callable, Optional, Any\n\n@dataclass\nclass AudioConfig:\n    \"\"\"\n    Central configuration for Audio Processing.\n    \"\"\"\n    sr: int = 32000\n    n_fft: int = 2048\n    win_length: int = n_fft\n    hop_length: int = win_length // 2\n    n_mels: int = 256\n    top_db: int = 80\n    \n    # Slicing/Striping params\n    stripe_width: int = sr // (n_fft // 4)\n    stripe_overlap: int = stripe_width // 5 \n    \n    # Helper to inject external slicing logic\n    slicing_func: Any = None\n\n    prior_cut_sec: int = 5","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-19T08:22:46.699381Z","iopub.execute_input":"2025-12-19T08:22:46.699608Z","iopub.status.idle":"2025-12-19T08:22:46.705276Z","shell.execute_reply.started":"2025-12-19T08:22:46.699589Z","shell.execute_reply":"2025-12-19T08:22:46.704708Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from sklearn.model_selection import train_test_split\nfrom sklearn.preprocessing import LabelEncoder\nimport numpy as np\n#todo statify k-fold\nle = LabelEncoder().fit(train_df.primary_label)\n#==full version\n# train_idx, test_idx = train_test_split(\n#     np.arange(len(train_df)), \n#     test_size = 0.2, \n#     random_state = 32, \n#     stratify = train_df['primary_label']\n# )\n\n#==small version\n# train_idx, test_idx = train_test_split(\n#     np.arange(len(train_df)),\n#     #train_size = 2_000,\n#     test_size = 0.2, #5_00, \n#     random_state = 32, \n#     stratify = train_df['primary_label']\n# )\n\ntrain_idx, small_test_idx = train_test_split(\n    np.arange(len(train_df)),\n    train_size = 0.8,\n    test_size = 0.2*0.2, \n    random_state = 32, \n    stratify = train_df['primary_label']\n)\n#==random sampling\n#np.random.seed(42)\n#df_idx = np.arange(len(train_df))\n#np.random.shuffle(df_idx)\n#test_size = int(len(train_df)*0.2)\n#test_idx = df_idx[:test_size]\n#train_idxs =  [np.random.choice(df_idx[test_size:], test_size) for _ in range(10)]\n#train_idx = df_idx[test_size:]","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-19T08:22:46.705953Z","iopub.execute_input":"2025-12-19T08:22:46.706267Z","iopub.status.idle":"2025-12-19T08:22:48.049975Z","shell.execute_reply.started":"2025-12-19T08:22:46.706238Z","shell.execute_reply":"2025-12-19T08:22:48.049376Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import torch\nimport torchaudio.transforms as transforms\nimport torchaudio\n#import librosa\nfrom joblib import Parallel, delayed\nimport ast\n\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n\nclass BirdDataset(torch.utils.data.Dataset):\n    def __init__(\n        self,\n        meta_df,\n        label_encoder,\n        config: AudioConfig,\n        group_mode:bool = True,\n        return_id: bool = True,\n        prediction_as_target: bool = False,\n        root_dir:str = '/kaggle/input/birdclef-2025/train_audio',\n        multitarget:bool = True,\n    ):\n        super().__init__()\n        self.cfg = config\n        #sklearn labelencoder\n        self.label_encoder = label_encoder\n        #files root dir\n        self.root_dir = root_dir\n        #flag to retrieve whole sequnces instead of just tokens\n        self.group_mode = group_mode\n        #flag to retrieve filenames to sequence of tensors\n        self.return_id = return_id\n        #includes the secondary labels as primary\n        self.multitarget = multitarget\n        #flag to retrive pseudo labels\n        self.prediction_as_target = prediction_as_target\n        #will serve as pseudo labels storage later\n        self.predictions = None\n        \n        #creating spectograms and preparing labels\n        results_with_index = self._parallel_prepare(meta_df)\n        \n        #extract spectograms and labels\n        self.spectrs, self.labels, self.idx = zip(*results_with_index)\n\n        if not group_mode:\n            #few row down we do item assignment\n            self.labels = list(self.labels)\n            \n            #expand labels with respect to token_num in each sample\n            for i, tokenized_tensor in enumerate(self.spectrs):\n                expand_size = tokenized_tensor.shape[-1]\n                self.labels[i] = np.tile(self.labels[i], (expand_size,1))\n                \n            #concatenate all tensor into 1 huge brick, because we doesn't care about their relations\n            self.spectrs = torch.cat(self.spectrs, dim = -1)\n            self.labels = [arr for sublist in self.labels for arr in sublist]\n\n        #label tensor init with n_classes\n        label_tensor = torch.zeros(len(self.labels), len(self.label_encoder.classes_)) #B, L\n        #ohe assigning \n        #todo it seems very slow, but I don't know better approach for multi label ohe assigning\n        #because once again, self.labels isn't a size=1 array\n        for i, label in enumerate(self.labels):\n            label_tensor[i, label] = 1\n\n        self.labels = label_tensor\n        self.idx = np.array(self.idx, dtype = object)\n        \n        self.token_h, self.token_w =  config.n_mels, config.stripe_width\n    \n    def preprocess_label_file(\n        self, \n        file_path, \n        label, \n        secondary_labels, \n        root_dir = None,\n    ):\n        \"\"\"\n        Reads audio -> STFT -> Mel/Mel**2/Linear Features -> Stacks them.\n        Reads id(str) -> id from label_encoder\n        \"\"\"\n        if root_dir is None:\n            root_dir = self.root_dir\n        path = os.path.join(root_dir, file_path)\n        waveform, sr = torchaudio.load(path, num_frames = self.cfg.sr * self.cfg.prior_cut_sec)\n        \n        if sr != self.cfg.sr:\n            waveform = transforms.Resample(sr, self.cfg.sr)(waveform)\n            \n        waveform = waveform[0]\n        \n        label_id = self.label_encoder.transform([label])\n\n        #secondary label processing\n        #str -> list\n        secondary_labels = ast.literal_eval(secondary_labels)\n        #if secondary labels actually exist:\n        if (secondary_labels != ['']) and self.multitarget:\n            secondary_labels = self.label_encoder.transform(secondary_labels)\n            label_id = np.concatenate((label_id, secondary_labels))\n\n        slices = self.cfg.slicing_func( #[C, H, W] -> [C, H, W, T]\n                    ar = waveform.unsqueeze(0).unsqueeze(0),\n                    stripe = self.cfg.stripe_width * self.cfg.hop_length, \n                    overlap = self.cfg.stripe_overlap * self.cfg.hop_length,\n                    pad_value = 0,\n                )\n      \n        return slices, label_id, file_path\n\n    def _parallel_prepare(self, df):\n        \"\"\"\n        parallel file processing using joblib\n        \"\"\"\n\n        columns_of_interest = ['filename', 'primary_label', 'secondary_labels']        \n        \n        results_with_index = Parallel(n_jobs=-1)(\n                    delayed(self.preprocess_label_file)(\n                        file_path = row.filename, \n                        label = row.primary_label,\n                        secondary_labels = row.secondary_labels,\n                    ) \n                    for _, row in df[columns_of_interest].iterrows()\n                )\n\n        return results_with_index\n\n    def update_pseudo_labels(self, obj:pd.Series):\n        \"\"\"\n        appends new pseudolabels into existing pseudolabels storage\n        \"\"\"\n        if self.predictions is None:\n            self.predictions = pd.Series()\n\n        self.predictions = pd.concat([self.predictions, obj])\n        \n    def __len__(self):\n        return len(self.labels) \n\n    def __getitem__(self, index):\n        \"\"\"\n        returns batch with ohe target\n        \"\"\"\n        if self.group_mode:\n            #it is the only option to fix dataloader select random samples via list problem\n            #we can't select tuple[list], so instead we have to itterate\n            #yes, it's the ONLY PLACE, we can't even use collate_fn\n            if isinstance(index, list) or isinstance(index, np.ndarray):\n                batch_tensor = [self.spectrs[i] for i in index]\n            else:\n                batch_tensor = self.spectrs[index]\n\n        else:\n            #because in not group mode we iterate through tokens and they are in the last dim\n            batch_tensor = self.spectrs[..., index]\n\n        str_id = self.idx[index]\n        #pseudo labels\n        if self.prediction_as_target:\n            label_tensor = self.predictions[str_id].values\n        else:\n            label_tensor = self.labels[index]\n\n        #in predict mode, I also want to retrieve a file name\n        if self.return_id:\n            return batch_tensor, label_tensor, str_id\n        else:\n            return batch_tensor, label_tensor","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-19T08:22:48.051457Z","iopub.execute_input":"2025-12-19T08:22:48.051805Z","iopub.status.idle":"2025-12-19T08:22:48.829528Z","shell.execute_reply.started":"2025-12-19T08:22:48.051785Z","shell.execute_reply":"2025-12-19T08:22:48.828895Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class OnlineDataset(torch.utils.data.Dataset):\n    \"\"\"\n    Essentially, just loads only needed part of default BirdDataset\n    \"\"\"\n    def __init__(        \n        self,\n        meta_df,\n        label_encoder,\n        config: AudioConfig,\n        group_mode:bool = True,\n        return_id: bool = True,\n        prediction_as_target: bool = False,\n        root_dir:str = '/kaggle/input/birdclef-2025/train_audio/',\n        multitarget:bool = True,\n    ):\n        super().__init__()\n        self.cfg = config\n        #sklearn labelencoder\n        self.label_encoder = label_encoder\n        #files root dir\n        self.root_dir = root_dir\n        #flag to retrieve whole sequnces instead of just tokens\n        self.group_mode = group_mode\n        #flag to retrieve filenames to sequence of tensors\n        self.return_id = return_id\n        #includes the secondary labels as primary\n        self.multitarget = multitarget\n        #flag to retrive pseudo labels\n        self.prediction_as_target = prediction_as_target\n        \n        #df that contains info about each file\n        self.meta_df = meta_df\n        self.predictions = None \n\n    def update_pseudo_labels(self, obj:pd.Series):\n        \"\"\"\n        appends new pseudolabels into existing pseudolabels storage\n        \"\"\"\n        if self.predictions is None:\n            self.predictions = pd.Series()\n\n        self.predictions = pd.concat([self.predictions, obj])\n\n    def __len__(self):\n        return len(self.meta_df)\n\n    def __getitem__(self, index):\n        temp_ds = BirdDataset(\n            meta_df = self.meta_df.iloc[index],\n            label_encoder = self.label_encoder,\n            config = self.cfg,\n            group_mode = self.group_mode,\n            return_id = self.return_id,\n            prediction_as_target = self.prediction_as_target,\n            root_dir = self.root_dir,\n            multitarget = self.multitarget,\n        )\n        if self.prediction_as_target:\n            temp_ds.predictions = self.predictions[index]\n            \n        return temp_ds[:]\n        ","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-19T08:22:48.830284Z","iopub.execute_input":"2025-12-19T08:22:48.830533Z","iopub.status.idle":"2025-12-19T08:22:48.837597Z","shell.execute_reply.started":"2025-12-19T08:22:48.830508Z","shell.execute_reply":"2025-12-19T08:22:48.836915Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"#AI generated (not slope)\nfrom torch.utils.data import Sampler\n\nclass IndexBatchSampler(Sampler):\n    def __init__(self, data_source, batch_size: int, shuffle: bool = True, drop_last: bool = False):\n        \"\"\"\n        Yields a list of indices at each iteration instead of a single index.\n        \n        Args:\n            data_source: The dataset (used to determine length).\n            batch_size: Number of indices to yield per iteration.\n            shuffle: Whether to shuffle indices before batching.\n            drop_last: Whether to drop the last incomplete batch.\n        \"\"\"\n        self.data_source = data_source\n        self.batch_size = batch_size\n        self.shuffle = shuffle\n        self.drop_last = drop_last\n\n    def __iter__(self):\n        # Create a list of all indices\n        indices = list(range(len(self.data_source)))\n        \n        if self.shuffle:\n            np.random.shuffle(indices)\n            \n        # Yield chunks (batches) of indices\n        batch = []\n        for idx in indices:\n            batch.append(idx)\n            if len(batch) == self.batch_size:\n                yield batch\n                batch = []\n        \n        # Handle the remaining items\n        if len(batch) > 0 and not self.drop_last:\n            yield batch\n\n    def __len__(self):\n        # Calculate how many batches this sampler will produce\n        if self.drop_last:\n            return len(self.data_source) // self.batch_size\n        else:\n            return (len(self.data_source) + self.batch_size - 1) // self.batch_size","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-19T08:22:48.838612Z","iopub.execute_input":"2025-12-19T08:22:48.838877Z","iopub.status.idle":"2025-12-19T08:22:48.858951Z","shell.execute_reply.started":"2025-12-19T08:22:48.838859Z","shell.execute_reply":"2025-12-19T08:22:48.858177Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from torch import nn\nimport timm \nclass Encoder(nn.Module):\n    def __init__(self, backbone_name:str, original_weights: bool = True):\n        super().__init__()\n        self.backbone = timm.create_model(\n            backbone_name, \n            pretrained=original_weights, \n            num_classes=0,\n        )\n        #init data-transform functions\n        self.stft_transform = None\n        self.db_transform = None\n        self.mel_transform = None\n        self.pooler = None\n\n    def set_audiopreprocessing(self, cfg):\n        self.stft_transform = transforms.Spectrogram(\n            n_fft=cfg.n_fft, \n            hop_length=cfg.hop_length,\n            power=1.0,\n        )\n        self.db_transform = transforms.AmplitudeToDB(\n            stype=\"magnitude\", \n            top_db=cfg.top_db\n        )\n        self.mel_transform = transforms.MelScale(\n            sample_rate=cfg.sr,\n            n_stft=cfg.n_fft // 2 + 1,\n            n_mels=cfg.n_mels,\n        )\n        self.pooler = nn.AdaptiveMaxPool1d(cfg.n_mels)\n        \n    def stripe_w_overlap(\n        self, \n        ar:torch.Tensor, \n        stripe:int, \n        overlap:int,\n        pad_value:float,\n    ):\n        \"\"\"\n        splits tensor into overlaping chunks along time dim\n        the last token is dropped\n        if the time < token_width => pad with constant value of pad_value\n        returns (B(depends on input), C, H, W, tokens)\n        \"\"\"\n        step = stripe - overlap\n        time_len = ar.shape[-1] #time\n\n        #if time < stripe => pad with zeros, because we can't stack None values\n        if time_len < stripe:\n            ar = torch.nn.functional.pad(\n                input = ar,\n                pad = (0, stripe - time_len), \n                mode = 'constant', \n                value = pad_value\n                \n            )\n            num_steps = 1\n        #if time < stripe+step means that there is possible only 1 step \n        elif time_len < stripe+step:\n            num_steps = 1\n        #number of full chunks\n        else:\n            num_steps = (time_len - stripe) // step\n            \n        try:\n            striped_tensor = torch.stack(\n                [\n                    ar[:,..., i * step : i * step + stripe] \n                     for i in range(num_steps)\n                ]\n            ) #(T, ...)\n        except: \n            raise RuntimeError(f\"error in stripes. here is the tensor shape: {ar.shape} ; the stripe: {stripe} ; and overlap: {overlap} ; and steps: {num_steps} \")\n\n        striped_tensor = torch.moveaxis(striped_tensor, 0, -1).contiguous() # (T, ...) -> (..., T)\n\n        return striped_tensor\n        \n    def forward(self, x, **kwargs):\n        #(B*T, C, H, W) -> (B*T, C_backbone)\n        \n        #==prerpcoessing raw waves into images\n        #(B*T, C, H, W) -> (B, C_spects, C_mel, W_stft)\n        \n        stft = self.stft_transform(x)\n        stft = stft[..., :320] #todo idk why, but I started to receive 321 windows after moving striping on raw waves\n        #conver to dB spectrs\n        mel_out_stft = self.db_transform(self.mel_transform(stft))\n        mel_out_stft_2 = self.db_transform(self.mel_transform(stft**2))\n        stft = self.db_transform(stft)\n        \n        #normalize stft to mel size\n        mega_B, H, W, C_stft, W_stft = stft.shape\n        stft = torch.transpose(stft, -1, -2) # mega_B, H, W, C_stft, W_stft -> mega_B, H, W, W_stft, C_stft\n\n        #flatten B, H, W, T into a single *B* dimension \n        stft = stft.reshape(-1, 1, C_stft) #(B,H,W,W_s,C) -> (*B*, 1, C)\n        pooled_lin = self.pooler(stft) # (*B*, 1, C_stft) -> (*B*, 1, C_mel) \n        pooled_lin = pooled_lin.reshape(mega_B, H, W, W_stft, -1) #(mega_B, H, W, W_stft, C_mel)\n        \n        x = torch.stack((mel_out_stft, mel_out_stft_2, pooled_lin), dim = 1) #B, C_spects, H, W, C_mel, W_stft \n        x = x.squeeze(2).squeeze(2) #B, C_spects, C_mel, W_stft \n        #==prerpcoessing raw waves into images\n        \n        return self.backbone(x) #(B*T, C, H, W) -> (B*T, C)\n\nclass Classifier(nn.Module):\n    def __init__(\n        self, \n        seq_mode = True, #to process instance wise or sequence wise\n        single_head:nn.Module = None,\n        multi_head:nn.Module = None,\n        single_activation:nn.Module = None,\n        multi_activation:nn.Module = None,\n\n    ):\n        super().__init__()\n        self.seq_mode = seq_mode\n        self.single_target_model = single_head\n        self.multi_target_model = multi_head\n        #activation fn in this class instead of token encoder, if I want to thrain only 1 SED head\n        self.single_activation = single_activation\n        self.multi_activation = multi_activation\n    \n    def forward(\n            self, \n            x, \n            multitarget_mask = None,\n            attention_mask:torch.Tensor = None,\n            return_NoF:bool = False,\n            pool:bool = True,\n        ):\n        \"\"\"\n        Pass padded vision embeddings in classifier with implementing additional logic:\n        * multi target model split(multitarget_mask)\n        * removes NoF token (return_NoF = False)\n        * returns prediction token wise (pool = False)\n\n        the classifier object's attribute seq_mode determines, whether to return list of token-wise predictions\n        or padded tensor of token-wise predictions\n        Args:\n            x: features from vision encoder.\n            multitarget_mask: list of booleans, where True means to use multitarget model.\n            return_NoF: whether to return NoF token.\n            pool: wheter to return output sequence wise or token wise.\n        \"\"\"\n\n        #B, T_max, Channels -> B, empty/T_max, Class_digit+NoF\n        #we have rotary embeddings, so we can't pass empty batches anymore\n        if torch.all(multitarget_mask) == True:\n            output = self.multi_target_model(x, return_NoF, attention_mask, pool)\n            output = self.multi_activation(output)\n        elif torch.any(multitarget_mask) == True:\n            multi_output = self.multi_target_model(x[multitarget_mask], return_NoF, attention_mask[multitarget_mask], pool)\n            multi_output = self.multi_activation(multi_output)\n            single_output = self.single_target_model(x[~multitarget_mask], return_NoF, attention_mask[~multitarget_mask], pool)\n            single_output = self.single_activation(single_output)\n            output = torch.cat([multi_output, single_output], dim = 0)\n        else:\n            output = self.single_target_model(x, return_NoF, attention_mask, pool)\n            output = self.single_activation(output)\n    \n        if not self.seq_mode:\n\n            masked_output = []\n            last_indx = torch.sum(attention_mask, dim = 1).int()\n            \n            #list of token predictions for each instance\n            for i, out in enumerate(output):\n                masked_output.append(out[:last_indx[i]].cpu().detach())\n\n            output = masked_output\n            \n        return output\n\n\nclass CLEFModel(nn.Module):\n    def __init__(\n        self, \n        encoder = None, \n        classifier = None,\n        padding_value = 0,\n        return_NoF = False,\n    ):\n        super().__init__()\n        self.encoder = encoder\n        self.classifier = classifier\n        self.padding_value = padding_value\n        self.return_NoF = return_NoF\n\n    def pad_list_to_tensor(self, tensor_list:list, padding_value:float = None):\n        \"\"\"\n        takes a list of embeddings [(T_i, C), ...] and creates out of them 1 padded tensor\n        \"\"\"\n        if padding_value is None:\n            padding_value = self.padding_value\n        #in batch maximum for padding (in tokens)\n        max_len = np.max([s.shape[0] for s in tensor_list])\n        \n        batch_list = []\n        attention_mask_list = []\n        for tensor in tensor_list:\n            #create attention mask, to prevent usage of padded values in the future\n            attention_mask = torch.ones(tensor.shape[0])\n            attention_mask = torch.nn.functional.pad(\n                attention_mask, \n                (0, max_len - tensor.shape[0]), #token wise padding\n                mode ='constant',\n                value = 0\n            )\n            \n            attention_mask_list.append(attention_mask)\n            \n            tensor = torch.nn.functional.pad(\n                            tensor, \n                            (0,0,0, max_len - tensor.shape[0]), #token wise padding\n                            mode ='constant', \n                            value = padding_value #equivavelnt of zero\n                        )\n        \n            batch_list.append(tensor)\n        \n        batch_tensor = torch.stack(batch_list, dim=0) # B, max(T_b), C\n        attention_mask = torch.stack(attention_mask_list, dim=0) # B, max(T_b)\n        \n        return batch_tensor, attention_mask\n\n    def forward(\n        self, \n        x, \n        multitarget_mask,\n        pool,\n    ):\n        \"\"\"\n        Pass list of sequences into vision encoder and classifier models.\n        \n        seq-wise:\n        Unioun list of tensor into huge one, to pass it through encoder,\n        The received features are padded, to create once again 1 huge tensor for classifier\n        \"\"\"\n        #list -> predict seq-wise\n        if not isinstance(x, torch.Tensor):\n            \n            #huge tensor for feature encoding \n            x_feature_tensor = torch.cat(x, dim = -1).permute(-1, 0, 1, 2).to(device) # (T*B, C, H, W)\n            try:\n                x_feature_tensor = self.encoder(x_feature_tensor) #(T*B, C)\n            except:\n                raise ValueError(f\"Most likely cuda memmory allocation, tensor shape: {x_feature_tensor.shape}\")\n            \n            #split back into seqs \n            prev_len = 0\n            embed = [] # (B, T_b, C)\n            for sub_x in x:\n                cur_len = sub_x.shape[-1]\n                embed.append(x_feature_tensor[prev_len : prev_len+cur_len])# (T_b, C)\n                prev_len = prev_len+cur_len\n                \n            #pad sequnces\n            embed, attention_mask = self.pad_list_to_tensor(embed)\n        else:        \n            embed = self.encoder(x.to(device))\n            attention_mask = torch.ones((embed.shape(0), embed.shape(1)))\n\n        probs = self.classifier(\n            x = embed, \n            attention_mask = attention_mask,\n            multitarget_mask = multitarget_mask,\n            return_NoF = self.return_NoF,\n            pool = pool\n        )\n        return probs","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-19T08:22:48.860095Z","iopub.execute_input":"2025-12-19T08:22:48.860409Z","iopub.status.idle":"2025-12-19T08:22:57.873261Z","shell.execute_reply.started":"2025-12-19T08:22:48.860383Z","shell.execute_reply":"2025-12-19T08:22:57.872479Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class ETransformerBlock(nn.Module):\n    def __init__(\n        self,\n        embedding_size,\n        num_head,\n        d_embed,\n        d_c,\n        d_c1,\n        d_rotate,\n    ):\n        super().__init__()\n        self.rms_n = nn.RMSNorm(embedding_size)\n        self.att = MultiHeadLatentAttention(\n            d_model = embedding_size, #after attention embedding_size\n            num_head = num_head, \n            d_embed = d_embed, #input to attention embedding_size\n            d_c = d_c, \n            d_c1 = d_c1, \n            d_rotate = d_rotate, \n            dropout=0.1, \n            bias=True,\n        )\n        self.liner = nn.Linear(embedding_size, embedding_size)\n        self.lin_act = nn.SiLU()\n\n    def forward(self, x, att_mask):\n        res = x\n        x = self.rms_n(x)\n        x = self.att(sequence = x, att_mask = att_mask)\n        x += res\n        res = x\n        x = self.rms_n(x)\n        x = self.liner(x)\n        x = self.lin_act(x)\n        x += res\n        return x","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-19T08:22:57.873987Z","iopub.execute_input":"2025-12-19T08:22:57.874236Z","iopub.status.idle":"2025-12-19T08:22:57.880205Z","shell.execute_reply.started":"2025-12-19T08:22:57.874217Z","shell.execute_reply":"2025-12-19T08:22:57.879439Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class TokenEncoder(nn.Module):\n    def __init__(\n        self, \n        embedding_size,\n        class_num,\n        num_head,\n        d_c,\n        d_c1,\n        d_rotate,\n        pool_type:str = 'cls',\n    ):\n        super().__init__()\n        #self.vembed2embed = nn.Linear(vembedding_size, embedding_size)\n        self.output_size = class_num+1\n        self.embedding_size = embedding_size\n        self.num_head = num_head\n        self.embed2embed = ETransformerBlock(\n            self.embedding_size,\n            num_head,\n            self.embedding_size,\n            d_c,\n            d_c1,\n            d_rotate,\n        )\n        self.embed2class = nn.Linear(embedding_size, self.output_size)\n        self.cls_token = nn.Parameter(torch.randn(1, 1, embedding_size))\n        self.pool_type = pool_type\n        \n    def forward(\n        self,\n        x,\n        return_NoF,\n        attention_mask:torch.Tensor = None,\n        pool:bool = True,\n    ):\n        \n        batch_size = x.shape[0]\n        cls = self.cls_token.expand(batch_size, -1, -1)\n        x = torch.cat((cls, x), dim=1) #B, T, E -> B, T+1, E\n        #todo attention mask expand\n        attention_mask = torch.cat(\n            (torch.ones(attention_mask.shape[0]).unsqueeze(-1), attention_mask), \n            dim = -1\n        ) #B, T -> B, T+1 {1, 0}\n        \n        q_mask = attention_mask.unsqueeze(2) \n        k_mask = attention_mask.unsqueeze(1)\n        attention_matrix = q_mask * k_mask # (B, T+1, T+1)\n        attention_matrix -= 1 #{0 ,-1}\n        attention_matrix *= 1e9 #{0, -inf}\n        attention_matrix = attention_matrix.unsqueeze(1).expand(-1, self.num_head, -1, -1).to(device) #(B, H, T+1, T+1)\n        output = self.embed2embed(x, att_mask = attention_matrix) #B, T+1, E -> B, T+1, E #todo probably bolleans to int in att_mask\n        #output = x #for simplicity in test setup\n\n        if pool:\n            #or in model cls_token pool\n            if self.pool_type == 'cls':\n                output = output[:, 0, :] #(B, 1, E)\n                \n            #or token_avg pool\n            elif self.pool_type == 'mean':\n                \n                output = output[:, 1:, :].to(device)\n                attention_mask = attention_mask[:, 1:].to(device)\n                last_indx = torch.sum(attention_mask, dim = 1).int().unsqueeze(-1).to(device) # (B, 1)\n                \n                #handle empty batch\n                if output.shape[0] == 0:\n                    output = torch.empty(0, self.embedding_size, dtype=output.dtype, device=output.device)\n                else:\n                    #because 'output' contains padded tokens (with zeros) calculated stats will also include them\n                    #to play it safe, I want to multiply 'output' on 'att_mask'\n                    output = (output * attention_mask.unsqueeze(-1)).sum(dim = 1) / last_indx #B, T, E -> B, E\n            else:\n                raise ValueError(f\"pool_type {self.pool_type} isn't implemented yet\")\n        else:\n            output = output[:, 1:, :]\n\n        output = self.embed2class(output) #B, ?, E -> B, ?, C+1\n    \n        if not return_NoF:\n            #last digit is NoF\n            output = output[..., :-1]\n                \n        return output","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-19T08:22:57.880965Z","iopub.execute_input":"2025-12-19T08:22:57.881197Z","iopub.status.idle":"2025-12-19T08:22:57.904080Z","shell.execute_reply.started":"2025-12-19T08:22:57.881180Z","shell.execute_reply":"2025-12-19T08:22:57.903398Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from torch.utils.data import DataLoader\nfrom tqdm.autonotebook import tqdm\nfrom sklearn.metrics import roc_auc_score\nimport wandb\n    \nclass CustomTrainer:\n    def __init__(\n        self,\n        batch_size,\n        optimizer = None,\n        scheduler = None,\n        loss = None,\n    ):\n        self.optimizer = optimizer\n        self.scheduler = scheduler\n        self.loss = loss\n        self.batch_size = batch_size\n        self._cached_model = None\n        self.skip_seq_len = None\n\n\n    def _prepare_dataloader(\n        self, \n        dataset, \n        shuffle,\n    ):\n        \"\"\"\n        additional dataloader transforming(used only for group_mode == True, because of tuple batches)\n        \"\"\"\n        if hasattr(dataset, 'group_mode') and dataset.group_mode:\n                            \n            batch_sampler = IndexBatchSampler(\n                dataset,\n                batch_size = self.batch_size,\n                shuffle = shuffle\n            )\n            dataloader = DataLoader(\n                dataset,\n                batch_size = None,\n                sampler = batch_sampler,\n            )\n        else:\n            dataloader = DataLoader(\n                train_ds, \n                batch_size = self.batch_size, \n                shuffle = shuffle\n            )\n        \n        return dataloader\n        \n    def _pass_batch(\n        self,\n        input,\n        target,\n        return_pred_target:bool = False,\n        only_single:bool = False,\n    ):\n        #todo. it's a very dumb solution for solving data allocation in cuda\n        if self.skip_seq_len is not None:\n            max_len = 0\n            if isinstance(input, list) and len(input) == 0:\n                return None\n                \n            if isinstance(input, list) and len(input) > 0:\n                max_len = torch.max(torch.tensor([i.shape[-1] for i in input]))\n            elif isinstance(input, torch.Tensor):\n                max_len = input.shape[-1]\n            \n            if max_len > self.skip_seq_len:\n                return None\n                \n        #if target is np array -> target is a list of tokens with different length -> pad+tensor\n        #the very same thing means, that we want to process token though model without any pooler\n        at_mask = 1\n        pool = True\n        if isinstance(target, np.ndarray):\n            pool = False\n            target, at_mask = self._cached_model.pad_list_to_tensor(target, padding_value = 0)\n            #because to multiply (b,n,c) on mask we should expand it from (b,n) to (b,n,1) for broadcasting\n            at_mask = at_mask.unsqueeze(-1).to(device)\n        \n        #multitarget split\n        if only_single:\n            multitarget_mask = torch.tensor([False] * len(target)) #base case\n        else:\n            multitarget_mask = torch.sum(target, dim = -1) != 1\n\n            #if we have targets as tokens, than recreate mask to batch idx\n            if multitarget_mask.dim() > 1:\n                multitarget_mask = torch.any(multitarget_mask, dim = 1)\n\n        prediction = self._cached_model(input, multitarget_mask, pool) * at_mask\n        #multitarget rearangment\n        target = torch.concat([target[multitarget_mask], target[~multitarget_mask]]).squeeze(1)\n        target = target.to(device)\n        \n        grad_output = self.loss(prediction, target)\n        \n        if return_pred_target:\n            return grad_output, prediction.cpu().detach(), target.cpu().detach()\n        else:\n            return grad_output\n\n    def _pass_batch_nano_vram(\n        self,\n        input,\n        target,\n        return_pred_target:bool = False,\n        only_single:bool = False,\n        micro_batch_size:int = 1,\n    ):\n        gradients_magn = []\n        predictions = []\n        targets = []\n        gradient_magn = 0\n        procesed_micro_batches = 0\n        for k in range(0, len(target), micro_batch_size):\n            i = input[k: k+micro_batch_size]\n            t = target[k: k+micro_batch_size]\n            res = self._pass_batch(\n                i,\n                t,\n                return_pred_target,\n                only_single\n            )\n            if res is None:\n                continue\n            procesed_micro_batches += 1\n            if return_pred_target:#TODO IT'S ONLY FOR EVAL\n                g, p, t = res\n                gradients_magn.append(g.item())\n                predictions.append(p.cpu().detach())\n                targets.append(t.cpu().detach())\n            else:#TODO IT'S ONLY FOR TRAIN\n                res = res *  micro_batch_size / len(target)\n                res.backward()\n                gradient_magn += res.item()\n\n\n        if procesed_micro_batches == 0:\n            return None\n        if gradient_magn != 0:\n            return gradient_magn, procesed_micro_batches\n        else:\n            return gradients_magn, predictions, targets, procesed_micro_batches\n        \n        \n    def train(\n        self,\n        model,\n        train_idx = None,\n        eval_idx = None,\n        ds_class = None,\n        train_ds = None,\n        eval_ds = None,\n        epochs:int = None,\n        steps:int = None,\n        #if steps are provided, than eval on each eval_freqth step, if epochs are provided eval on each (float part of epoch (e.g. each 0.3 epoch))\n        eval_freq:float = None,\n        #number of batches to process beefore updating\n        optim_freq:int = 1,\n        #the same as in eval_freq\n        save_freq:float = None,\n        skip_seq_len: int = None,\n        micro_batch_size: int = 1,\n    ):\n\n        train_dataloader = self._prepare_dataloader(train_ds, shuffle=True) \n        eval_dataloader = self._prepare_dataloader(eval_ds, shuffle=False)\n        \n        self.skip_seq_len = skip_seq_len\n        self._cached_model = model\n        if (epochs is None) and (steps is None):\n            raise ValueError(\"epochs and steps can not be Nones simultaneously\")\n        \n        if epochs is not None:\n            steps = epochs * len(train_dataloader)\n            eval_freq = int(eval_freq * len(train_dataloader))\n            save_freq = int(save_freq * len(train_dataloader))\n\n        cum_loss = 0\n        batches_processed = 0\n        train_iter = iter(train_dataloader)\n        progress_bar = tqdm(range(steps), desc=\"training\")\n        self._cached_model\n        self._cached_model.train()\n\n        \n        for i in progress_bar:\n            #continuously pulling from an iterator\n            try:\n                input, target = next(train_iter)\n                    \n            except StopIteration:\n                #reset iterator if it runs out before 'steps' is reached\n                train_iter = iter(train_dataloader)\n                input, target = next(train_iter)\n                \n            if micro_batch_size == 0:\n                output = self._pass_batch(\n                    input, \n                    target,\n                    return_pred_target = False, \n                    only_single = False, #todo set as parameter\n                )\n                if output is None:\n                    del input, target\n                    continue\n    \n                batches_processed += procesed_micro_batches\n                cum_loss += output.item()\n                output.backward()\n                \n            else:\n                output, procesed_micro_batches = self._pass_batch_nano_vram(\n                    input,\n                    target,\n                    return_pred_target = False,\n                    only_single = False,\n                    micro_batch_size = micro_batch_size,\n                )\n                if output is None:\n                    del input, target\n                    continue\n    \n                batches_processed += procesed_micro_batches\n                cum_loss += output#.item()\n \n            \n            if (i+1)%optim_freq == 0:\n                self.optimizer.step()\n                self.scheduler.step()\n                self.optimizer.zero_grad()\n            #(i+1)%freq <- e.g. i = 1 (2nd iter), freq = 2 (each 2 steps) -\\_( :) )_/-\n            if (i != 0) and ((i+1) % eval_freq == 0):\n                #eval loss per sample\n                eval_loss, roc_score = self.eval(eval_dataloader, micro_batch_size) \n                eval_loss /= len(eval_dataloader) \n                #train loss per sample\n                if micro_batch_size == 0:\n                    train_loss = cum_loss / (batches_processed * self.batch_size)\n                else:\n                    train_loss = cum_loss / (batches_processed * micro_batch_size)\n                    \n                log_dict = {\n                        \"train/loss\": train_loss,\n                        \"eval/loss\": eval_loss,\n                        \"eval/roc_auc\": roc_score,\n                        \"step\": i,\n                    }\n                \n                if self.scheduler:\n                    log_dict[\"train/lr\"] = self.optimizer.param_groups[0]['lr']\n                    \n                wandb.log(log_dict)\n                #print(f\"step: {i}\\ntrain_loss: {train_loss}\\neval_loss: {eval_loss}\\neval roc score: {roc_score}\")\n                cum_loss = 0\n                batches_processed = 0\n                self._cached_model.train()\n                \n            if (i+1)%save_freq == 0:\n                torch.save(model.state_dict(), f'dmodel_{i}.pt')\n\n        return self._cached_model\n\n    @torch.no_grad()\n    def predict(\n        self, \n        model,\n        dataset,\n        multitarget:bool = False\n    ):\n        \"\"\"\n        create predictions for each filename and store in dataset.predictions as pd.series\n        \"\"\"\n        dataloader = self._prepare_dataloader(dataset, shuffle=False)\n        model.eval()\n        for input, _, file_path in dataloader:\n            #input is a list of tensor\n            multitarget_mask = torch.tensor([multitarget] * len(input))\n            output = model(input, multitarget_mask, model.classifier.seq_mode)\n            \n            #if we run model in predicting only 1 vector for file, it remains to be a torch tensor (b,c)\n            #in that case we create a [np.array(c), ...(n-1)], to save predictions as a pd.series\n            if isinstance(output, torch.Tensor):\n                # for np.version > 2.1 should use simple .unstack()\n                output = np.split(output.cpu().detach().numpy(), output.shape[0], axis = 0)\n                \n            pretty_output = pd.Series(data = output, index = file_path)\n            dataset.update_pseudo_labels(pretty_output)\n            \n    \n    @torch.no_grad() #todo actually to inspect the magnitude of update isn't a bad idea, (remember TS)\n    def eval(\n        self,\n        eval_dataloader,\n        micro_batch_size,\n    ):\n        \"\"\"\n        evaluation implemented only in kaggle way (1 class)\n        \"\"\"\n        #todo handle multitarget target with roc auc\n        self._cached_model.eval()\n        \n        #todo we can rewrite it invo model.eval() method\n        return_NoF_buf = self._cached_model.return_NoF\n        if return_NoF_buf:\n            self._cached_model.return_NoF = False \n            \n        cum_loss = 0\n        y_true = []\n        y_pred = []\n        for input, target in eval_dataloader:\n            #todo throw an exception on micro_batch_size = 0\n            result = self._pass_batch_nano_vram(\n                input,\n                target,\n                return_pred_target = True,\n                only_single = False,\n                micro_batch_size = micro_batch_size,\n            )\n            if result is None:\n                del input, target\n                continue\n                \n            output, prediction, target, procesed_batches = result\n            y_true.extend(target)\n            y_pred.extend(prediction)\n            \n            cum_loss += torch.sum(torch.tensor(output)) / procesed_batches\n\n        #turning back on\n        self._cached_model.return_NoF = return_NoF_buf\n        \n        #roc calc\n        #datatypes transition\n        #also drop off the last digit, because the last digit is NoF\n        y_true = torch.concat(y_true, dim = 0).cpu().detach().numpy()\n        y_pred = torch.concat(y_pred, dim = 0).cpu().detach().numpy()\n        #roc mask, to include only presented classes\n        class_sums = y_true.sum(axis=0) > 0\n\n        y_true = y_true[:, class_sums]\n        y_pred = y_pred[:, class_sums]\n        \n        roc = roc_auc_score(\n            y_true, \n            y_pred, \n            average = 'macro', \n        ) \n        return cum_loss, roc","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-19T08:22:57.905934Z","iopub.execute_input":"2025-12-19T08:22:57.906211Z","iopub.status.idle":"2025-12-19T08:22:57.934315Z","shell.execute_reply.started":"2025-12-19T08:22:57.906194Z","shell.execute_reply":"2025-12-19T08:22:57.933474Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"base train eval dss","metadata":{}},{"cell_type":"code","source":"encoder = Encoder('efficientnet_b3a')\nconfig = AudioConfig(\n    n_fft = 400,\n    win_length = 400,\n    hop_length = 400 // 2,\n    slicing_func=encoder.stripe_w_overlap,\n    n_mels = 320,\n    stripe_width = 320,\n    stripe_overlap = 200,\n    prior_cut_sec = 5,\n)\nencoder.set_audiopreprocessing(config)\ntrain_ds = OnlineDataset(\n    meta_df = train_df.iloc[train_idx], #todo train size \n    label_encoder = le, \n    config = config,\n    group_mode = True,\n    return_id = False,\n)\neval_ds = OnlineDataset(\n    meta_df = train_df.iloc[small_test_idx], \n    label_encoder = le, \n    config = config,\n    group_mode = True,\n    return_id = False,\n    multitarget = False, #todo, how to properly handle multitarget in eval?! \n)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-19T08:22:57.935003Z","iopub.execute_input":"2025-12-19T08:22:57.935248Z","iopub.status.idle":"2025-12-19T08:23:00.657374Z","shell.execute_reply.started":"2025-12-19T08:22:57.935229Z","shell.execute_reply":"2025-12-19T08:23:00.656395Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import warnings\nwarnings.filterwarnings('ignore')\n\nsingle_target_head = TokenEncoder(\n    embedding_size = encoder.backbone.num_features,\n    class_num = len(le.classes_),\n    num_head = 4,\n    d_c = 256,\n    d_c1 = 256,\n    d_rotate = 16,\n    pool_type = 'cls',\n)\n# multi_target_head = TokenEncoder(\n#     embedding_size = encoder.backbone.num_features,\n#     class_num = len(le.classes_),\n# )\nclassifier = Classifier(\n    seq_mode = True,\n    single_head = single_target_head,\n    #multi_head = multi_target_head, #todo\n    multi_head = single_target_head, #todo same model is trained in different modes\n    single_activation = nn.Softmax(dim = -1), #dim = -1, because tokenwise softmax \n    multi_activation= nn.Sigmoid(),\n)\nmodel = CLEFModel(\n    encoder = encoder,\n    classifier = classifier,\n    padding_value = 0,\n).to(device)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-19T08:23:00.658347Z","iopub.execute_input":"2025-12-19T08:23:00.658656Z","iopub.status.idle":"2025-12-19T08:23:00.937616Z","shell.execute_reply.started":"2025-12-19T08:23:00.658636Z","shell.execute_reply":"2025-12-19T08:23:00.936695Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"train cell","metadata":{}},{"cell_type":"code","source":"batch_size = 64\noptimizer = torch.optim.AdamW(\n    params = model.parameters(), \n    lr = 5e-4\n)\nscheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(\n    optimizer = optimizer,\n    T_0 = int(len(train_ds)//batch_size*0.5),#todo change with respect to trainer\n    eta_min = 1e-4,\n)\nloss = nn.BCELoss(reduction = 'sum')\ntrainer = CustomTrainer(\n    optimizer = optimizer,\n    scheduler = scheduler,\n    loss = loss,\n    batch_size = batch_size,\n)\n\nwandb.init(project=\"my-project\", name=\"cls-short-1-epoch-fixed\")\n\nmodel = trainer.train(\n    model = model,\n    train_ds = train_ds,\n    eval_ds = eval_ds,\n    epochs = 1,\n    #steps = 10_000,\n    eval_freq = 0.2,\n    optim_freq = 1,\n    save_freq = 0.2,\n    skip_seq_len = 36,\n    micro_batch_size = 8,\n)\nwandb.finish()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-19T08:23:00.938507Z","iopub.execute_input":"2025-12-19T08:23:00.938749Z","execution_failed":"2025-12-19T10:59:58.405Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"#torch.save(model.state_dict(), 'lastfinetune.pt')\ntorch.save(trainer.optimizer.state_dict(), 'optimizer.pt')\ntorch.save(trainer.scheduler.state_dict(), 'scheduler.pt')\nwith open(\"audio_cfg.pkl\", 'wb') as f:\n    pickle.dump(config, f)\n\nwith open(\"label_encoder.pkl\", 'wb') as f:\n    pickle.dump(le, f)","metadata":{"trusted":true,"execution":{"execution_failed":"2025-12-19T10:59:58.405Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# import gc\n# gc.collect()\n# del model, train_ds, eval_ds, trainer ","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-15T19:41:19.638986Z","iopub.execute_input":"2025-12-15T19:41:19.639330Z","iopub.status.idle":"2025-12-15T19:41:20.686986Z","shell.execute_reply.started":"2025-12-15T19:41:19.639306Z","shell.execute_reply":"2025-12-15T19:41:20.685750Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# del model, train_ds, eval_ds, trainer \n# gc.collect()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-15T16:49:06.059909Z","iopub.status.idle":"2025-12-15T16:49:06.060126Z","shell.execute_reply.started":"2025-12-15T16:49:06.060020Z","shell.execute_reply":"2025-12-15T16:49:06.060031Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"pseudo-labels dataset creation","metadata":{}},{"cell_type":"code","source":"def label_radicaliser(prediction):\n    return 1/(1 + np.exp(-10*prediction+5))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-16T19:05:17.072260Z","iopub.execute_input":"2025-12-16T19:05:17.072964Z","iopub.status.idle":"2025-12-16T19:05:17.076276Z","shell.execute_reply.started":"2025-12-16T19:05:17.072943Z","shell.execute_reply":"2025-12-16T19:05:17.075696Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"with open(\"/kaggle/input/cltest/other/default/4/label_encoder.pkl\", 'rb') as f: #le path\n    le = pickle.load(f)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-17T09:42:23.693634Z","iopub.execute_input":"2025-12-17T09:42:23.694400Z","iopub.status.idle":"2025-12-17T09:42:23.705724Z","shell.execute_reply.started":"2025-12-17T09:42:23.694375Z","shell.execute_reply":"2025-12-17T09:42:23.705073Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import glob\nimport time\n#base ds into pseudo segment labels\n# weak_labels_ds = OnlineDataset(\n#     meta_df = train_df.iloc[train_idx[:1_000]], #todo train size \n#     label_encoder = le, \n#     config = config,\n#     group_mode = True,\n#     return_id = True,\n# )\n#no label ds into pseudo segment labels\nln_df = pd.DataFrame()\ndirectory = \"/kaggle/input/birdclef-2025/train_soundscapes\"\nfilename = glob.glob(os.path.join(directory, \"*ogg\"))    \nfilename = [os.path.basename(path) for path in filename]\ndummy_primaries = ['1139490'] * len(filename)\ndummy_sec = [\"['']\"] * len(filename)\nln_df['filename'] = filename\nln_df['primary_label'] = dummy_primaries\nln_df['secondary_labels'] = dummy_sec\n\nno_labels_ds = OnlineDataset(\n    meta_df = ln_df, #todo train size \n    label_encoder = le, \n    config = config,\n    group_mode = True,\n    return_id = True,\n    root_dir = directory,\n)\n\ntrainer = CustomTrainer(\n    batch_size = 8,\n)\n\nmodel.classifier.seq_mode = False\nmodel.return_NoF = False\n\n#model load\nmodel.load_state_dict(torch.load('/kaggle/input/cltest/other/default/4/dmodel_354.pt', map_location=torch.device('cpu'))) #model path\nmodel.eval()\n# start = time.time()\n# trainer.predict(\n#     model = model,\n#     dataset = weak_labels_ds,\n#     multitarget = True,\n# )\n# weak_labels_ds.predictions = pd.Series(data = np.apply_along_axis(label_radicaliser, 0, weak_labels_ds.predictions.values), index = weak_labels_ds.predictions.index)\n# end = time.time()\n# print(f'1st ds is done, speed: {len(weak_labels_ds) / (end - start)}')\nstart = time.time()\ntrainer.predict(\n    model = model,\n    dataset = no_labels_ds,\n    multitarget = True,\n)\nno_labels_ds.predictions = pd.Series(data = np.apply_along_axis(label_radicaliser, 0, no_labels_ds.predictions.values), index = no_labels_ds.predictions.index)\nend = time.time()\nprint(f'2nd ds is done, speed: {len(no_labels_ds) / (end - start)}')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-16T19:07:48.360580Z","iopub.execute_input":"2025-12-16T19:07:48.361314Z","iopub.status.idle":"2025-12-16T20:08:26.635797Z","shell.execute_reply.started":"2025-12-16T19:07:48.361290Z","shell.execute_reply":"2025-12-16T20:08:26.634956Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"no_labels_ds.predictions.to_csv('submit1.csv')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-16T20:08:54.778626Z","iopub.execute_input":"2025-12-16T20:08:54.779185Z","iopub.status.idle":"2025-12-16T20:09:00.277604Z","shell.execute_reply.started":"2025-12-16T20:08:54.779161Z","shell.execute_reply":"2025-12-16T20:09:00.277033Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"torch.save(no_labels_ds.predictions, 'no_labels_ds.pt')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-16T20:25:38.271702Z","iopub.execute_input":"2025-12-16T20:25:38.272283Z","iopub.status.idle":"2025-12-16T20:25:38.739404Z","shell.execute_reply.started":"2025-12-16T20:25:38.272261Z","shell.execute_reply":"2025-12-16T20:25:38.738586Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"load predictions","metadata":{}},{"cell_type":"code","source":"import glob\nln_df = pd.DataFrame()\ndirectory = \"/kaggle/input/birdclef-2025/train_soundscapes\"\nfilename = glob.glob(os.path.join(directory, \"*ogg\"))    \nfilename = [os.path.basename(path) for path in filename]\ndummy_primaries = ['1139490'] * len(filename)\ndummy_sec = [\"['']\"] * len(filename)\nln_df['filename'] = filename\nln_df['primary_label'] = dummy_primaries\nln_df['secondary_labels'] = dummy_sec\n\nno_labels_ds = OnlineDataset(\n    meta_df = ln_df, #todo train size \n    label_encoder = le, \n    config = config,\n    group_mode = True,\n    return_id = True,\n    root_dir = directory,\n)\nno_labels_ds.predictions = torch.load('/kaggle/input/pseudo-nl/no_labels_ds.pt', weights_only = False)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-17T13:38:18.479514Z","iopub.execute_input":"2025-12-17T13:38:18.479805Z","iopub.status.idle":"2025-12-17T13:38:21.113869Z","shell.execute_reply.started":"2025-12-17T13:38:18.479782Z","shell.execute_reply":"2025-12-17T13:38:21.113299Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"no_labels_ds.predictions[:5]","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-17T13:38:21.114806Z","iopub.execute_input":"2025-12-17T13:38:21.115018Z","iopub.status.idle":"2025-12-17T13:38:22.244774Z","shell.execute_reply.started":"2025-12-17T13:38:21.115002Z","shell.execute_reply":"2025-12-17T13:38:22.243993Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"create pseudo dataset","metadata":{}},{"cell_type":"code","source":"class PseudoOnlineDatasets(torch.utils.data.Dataset\n    \"\"\"\n    concatenates onlinedatasets with random sampling\n    \"\"\"\n    def __init__(self, datasets:list[OnlineDataset], group_mode:bool = True):\n        super().__init__()\n        self.datasets = datasets\n        self.lens = np.array([len(dataset.meta_df) for dataset in self.datasets])\n        self.lens_sum = np.sum(self.lens)\n        self.sample_p = self.lens / np.sum(self.lens)\n        self.group_mode = group_mode\n\n    def __len__(self):\n        return self.lens_sum\n    \n    def __getitem__(self, index):\n        ds_id = np.random.choice(len(self.datasets), p = self.sample_p)\n        ds = self.datasets[ds_id]\n        rescaled_index = np.floor(index / self.lens_sum * self.lens[ds_id]).astype(int)\n        return ds[rescaled_index]\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-17T13:38:26.509359Z","iopub.execute_input":"2025-12-17T13:38:26.510202Z","iopub.status.idle":"2025-12-17T13:38:26.516022Z","shell.execute_reply.started":"2025-12-17T13:38:26.510175Z","shell.execute_reply":"2025-12-17T13:38:26.515386Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"model.load_state_dict(torch.load('/kaggle/input/cltest/other/default/4/dmodel_354.pt', weights_only=True))#, map_location=torch.device('cpu'))) #model path","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# weak_labels_ds.prediction_as_target = True\nno_labels_ds.prediction_as_target = True\ntrain_ds.prediction_as_target = False\n\n# weak_labels_ds.return_id = False\nno_labels_ds.return_id = False\ntrain_ds.return_id = False\npo_dataset = PseudoOnlineDatasets([train_ds, no_labels_ds]) #weak_labels_ds\n\n\nbatch_size = 64\noptimizer = torch.optim.AdamW(\n    params = model.parameters(), \n    lr = 5e-4\n)\nscheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(\n    optimizer = optimizer,\n    T_0 = int(len(po_dataset)//batch_size*0.5),\n    eta_min = 1e-4,\n)\nloss = nn.BCELoss(reduction = 'sum')\ntrainer = CustomTrainer(\n    optimizer = optimizer,\n    scheduler = scheduler,\n    loss = loss,\n    batch_size = batch_size,\n)\n#turning back after creating pseudo_labels\nmodel.classifier.seq_mode = True","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-17T13:38:44.477469Z","iopub.execute_input":"2025-12-17T13:38:44.477736Z","iopub.status.idle":"2025-12-17T13:38:44.486223Z","shell.execute_reply.started":"2025-12-17T13:38:44.477714Z","shell.execute_reply":"2025-12-17T13:38:44.485714Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"len(po_dataset)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-16T21:28:20.260157Z","iopub.execute_input":"2025-12-16T21:28:20.260575Z","iopub.status.idle":"2025-12-16T21:28:20.266769Z","shell.execute_reply.started":"2025-12-16T21:28:20.260546Z","shell.execute_reply":"2025-12-16T21:28:20.265772Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"pseudolabels training","metadata":{}},{"cell_type":"code","source":"wandb.init(project=\"my-project\", name=\"pseudo-labels-cls-kl\")\nmodel = trainer.train(\n    model = model,\n    train_ds = po_dataset,\n    eval_ds = eval_ds,\n    epochs = 1,\n    eval_freq = 0.33,\n    optim_freq = 1,\n    save_freq = 0.33,\n    skip_seq_len = 12,\n    micro_batch_size = 1,\n)\nwandb.finish()","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}