Upload echo.ipynb

echo.ipynb

@@ -0,0 +1,1160 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import base64, gzip, math, os, functools, warnings, numpy as np, torch, transformers, aiohttp, torch.nn.functional as F, evaluate, json, random\n",
+    "from torch import Tensor, amp, optim, nn\n",
+    "from torch.utils.checkpoint import checkpoint\n",
+    "from torch.utils.tensorboard.writer import SummaryWriter\n",
+    "from threading import Thread\n",
+    "from typing import Dict, Optional, Tuple, Union, List, Any\n",
+    "from transformers.modeling_utils import PreTrainedModel\n",
+    "from dataclasses import dataclass\n",
+    "from transformers.optimization import Adafactor, AdafactorSchedule\n",
+    "from transformers import (Seq2SeqTrainer, Seq2SeqTrainingArguments, PretrainedConfig, TrainerCallback, WhisperProcessor, WhisperFeatureExtractor, WhisperTokenizerFast)\n",
+    "from torch.optim import Optimizer\n",
+    "import evaluate\n",
+    "from evaluate import module\n",
+    "from sklearn.metrics import accuracy_score, precision_score, f1_score, recall_score\n",
+    "from sklearn.model_selection import KFold, train_test_split\n",
+    "from datasets import load_dataset, Dataset, concatenate_datasets, IterableDatasetDict, Audio, load_from_disk\n",
+    "from torch.nn.functional import scaled_dot_product_attention\n",
+    "\n",
+    "from accelerate import Accelerator\n",
+    "import matplotlib.pyplot as plt\n",
+    "transformers.utils.logging.set_verbosity_error()\n",
+    "warnings.filterwarnings(action=\"ignore\")\n",
+    "warnings.warn = lambda *args, **kwargs: None\n",
+    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
+    "dtype = torch.float32\n",
+    "torch_dtype = torch.float32\n",
+    "torch.set_default_dtype(dtype)\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\n",
+    "class Linear(nn.Linear):\n",
+    "    def forward(self, x: Tensor) -> Tensor:# type: ignore\n",
+    "        return F.linear(x, self.weight.to(x.dtype),\n",
+    "                         None if self.bias is None else self.bias.to(x.dtype))\n",
+    "\n",
+    "class Conv1d(nn.Conv1d):\n",
+    "    def _conv_forward(self, x: Tensor, weight: Tensor, bias: Optional[Tensor] = None) -> Tensor:# type: ignore\n",
+    "        return super()._conv_forward(x, weight.to(x.dtype),\n",
+    "                                     None if bias is None else bias.to(x.dtype))\n",
+    "\n",
+    "class LayerNorm(nn.LayerNorm):\n",
+    "    def forward(self, x: Tensor) -> Tensor:  # type: ignore\n",
+    "        return super().forward(x.float()).type(x.dtype) "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class CombinedRotaryEmbedding(nn.Module):\n",
+    "    def __init__(self, base, dims, head, theta_learnable=True, rot_learnable=True,\n",
+    "                 matrix_learnable=False, freq_learnable=True):\n",
+    "        super(CombinedRotaryEmbedding, self).__init__()\n",
+    "\n",
+    "        self.base = base\n",
+    "        self.dims = dims\n",
+    "        self.head = head\n",
+    "\n",
+    "        self.h_dim = self.dims // self.head\n",
+    "        self.rot = (self.dims // self.head) // 2\n",
+    "\n",
+    "        self.thetas = nn.Parameter(torch.zeros(self.rot))\n",
+    "        self.r_pairs = nn.Parameter(data=torch.rand(self.rot, 2) * self.h_dim)\n",
+    "\n",
+    "        self.theta_scale = nn.Parameter(torch.ones(1), requires_grad=theta_learnable)\n",
+    "        self.rot_scale = nn.Parameter(torch.ones(1), requires_grad=rot_learnable)\n",
+    "\n",
+    "        self.r_matrix = nn.Parameter(torch.eye(n=self.h_dim), requires_grad=matrix_learnable)\n",
+    "\n",
+    "        freq_data = 1.0 / (self.base ** (torch.arange(start=0, end=self.h_dim, step=2).float() / self.h_dim))\n",
+    "        self.inv_freq = nn.Parameter(freq_data, requires_grad=freq_learnable)\n",
+    "\n",
+    "        self.orthogonal_reg_weight = 0.01\n",
+    "\n",
+    "    def blended_rotation_matrix(self, dims, i, j, theta):\n",
+    "        G = torch.eye(dims).to(theta.device)\n",
+    "        G[i, i] = torch.cos(theta)\n",
+    "        G[i, j] = -torch.sin(theta)\n",
+    "        G[j, i] = torch.sin(theta)\n",
+    "        G[j, j] = torch.cos(theta)\n",
+    "\n",
+    "        v = torch.zeros(dims).to(theta.device)\n",
+    "        v[i] = torch.cos(theta)\n",
+    "        v[j] = torch.sin(theta)\n",
+    "        H = torch.eye(dims).to(theta.device) - 2 * torch.outer(v, v) / torch.dot(v, v)\n",
+    "\n",
+    "        R = torch.eye(dims).to(theta.device)\n",
+    "        R[i, i] = torch.cos(theta)\n",
+    "        R[i, j] = -torch.sin(theta)\n",
+    "        R[j, i] = torch.sin(theta)\n",
+    "        R[j, j] = torch.cos(theta)\n",
+    "\n",
+    "        return (G + H + R) / 3\n",
+    "\n",
+    "    def apply_blended_rotation(self, x):\n",
+    "        adjusted_rot = int(torch.round(self.rot_scale * self.rot))\n",
+    "        for k in range(adjusted_rot):\n",
+    "            i, j = self.r_pairs[k].long()\n",
+    "            theta = self.thetas[k] * self.theta_scale\n",
+    "            B = self.blended_rotation_matrix(dims=self.h_dim, i=i, j=j, theta=theta)\n",
+    "            x = torch.matmul(input=x, other=B)\n",
+    "        return x\n",
+    "\n",
+    "    def update_base(self, new_base):\n",
+    "        if new_base is not None and new_base != self.base:\n",
+    "            self.base = new_base\n",
+    "            inv_freq = 1.0 / (self.base ** (torch.arange(start=0, end=self.h_dim, step=2).float() / self.h_dim))\n",
+    "            self.inv_freq.data.copy_(inv_freq)\n",
+    "            self.update_pairs()\n",
+    "\n",
+    "    def reset_parameters(self):\n",
+    "        nn.init.orthogonal_(self.r_matrix)\n",
+    "        nn.init.zeros_(self.thetas)\n",
+    "        nn.init.zeros_(self.r_pairs)\n",
+    "        nn.init.ones_(self.theta_scale)\n",
+    "        nn.init.ones_(self.rot_scale)\n",
+    "\n",
+    "    def orthogonal_regularization_term(self):\n",
+    "        loss = torch.tensor(0.0, device=self.r_matrix.device)\n",
+    "        if self.r_matrix.requires_grad:\n",
+    "            product = torch.matmul(self.r_matrix, self.r_matrix.t())\n",
+    "            identity = torch.eye(self.r_matrix.size(0)).to(self.r_matrix.device)\n",
+    "            loss = ((product - identity) ** 2).sum()\n",
+    "        return self.orthogonal_reg_weight * loss\n",
+    "\n",
+    "    def update_pairs(self):\n",
+    "        pairs = []\n",
+    "        while len(pairs) < self.rot:\n",
+    "            i, j = torch.randint(0, self.h_dim - 1, (2,))\n",
+    "            if i != j and (i, j) not in pairs and (j, i) not in pairs:\n",
+    "                pairs.append((i, j))\n",
+    "        self.r_pairs.data.copy_(torch.tensor(pairs, dtype=torch.float32))\n",
+    "\n",
+    "    def forward(self, x, global_step=None):\n",
+    "        if x.dim() not in [3, 4]:\n",
+    "            raise ValueError(f\"Expected input tensor to be 3D or 4D, but got {x.dim()}D\")\n",
+    "\n",
+    "        batch_size, seq_len, *rest = x.size()\n",
+    "\n",
+    "        if x.dim() == 3:\n",
+    "            dims = rest[0]\n",
+    "            if dims != self.head * self.h_dim:\n",
+    "                raise ValueError(f\"Expected dims ({dims}) to be compatible with head ({self.head}) * h_dim ({self.h_dim}={self.head * self.h_dim})\")\n",
+    "        else:\n",
+    "            head, h_dim = rest\n",
+    "            if head != self.head or h_dim != self.h_dim:\n",
+    "                raise ValueError(f\"For 4D input, expected head {self.head} and h_dim {self.h_dim}, but got head {head} and h_dim {h_dim}\")\n",
+    "\n",
+    "        x = x.view(batch_size, seq_len, self.head, self.h_dim)\n",
+    "        x = x.reshape(-1, self.h_dim)\n",
+    "\n",
+    "        x = self.apply_blended_rotation(x)\n",
+    "\n",
+    "        x = torch.matmul(input=x, other=self.r_matrix)\n",
+    "\n",
+    "        x = x.view(batch_size, seq_len, self.head, self.h_dim)\n",
+    "\n",
+    "        sinusoid_inp = torch.einsum('i, j -> i j', torch.arange(end=seq_len, device=x.device), self.inv_freq.to(device=x.device))\n",
+    "        sin = sinusoid_inp.sin()[None, :, None, :]\n",
+    "        cos = sinusoid_inp.cos()[None, :, None, :]\n",
+    "\n",
+    "        x1, x2 = x[..., ::2], x[..., 1::2]\n",
+    "        x = torch.cat(tensors=[x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1)\n",
+    "        x = x.view(batch_size, seq_len, self.dims)\n",
+    "\n",
+    "        return x\n",
+    "\n",
+    "class SinusoidalEmbedding(nn.Module):\n",
+    "    def __init__(self, n_ctx, dims, checkpoint):\n",
+    "        super().__init__()\n",
+    "        self.n_ctx = n_ctx\n",
+    "        self.dims = dims\n",
+    "        self.checkpoint = checkpoint\n",
+    "\n",
+    "        position = torch.arange(0, n_ctx, dtype=torch.float).unsqueeze(1)\n",
+    "        div_term = torch.exp(torch.arange(0, dims, 2).float() * -(math.log(10000.0) / dims))\n",
+    "        features = torch.zeros(n_ctx, dims)\n",
+    "        features[:, 0::2] = torch.sin(position * div_term)\n",
+    "        features[:, 1::2] = torch.cos(position * div_term)\n",
+    "        self.register_buffer('my_big_toe', features)\n",
+    "        self.pos_embeds = nn.Parameter(self.my_big_toe.clone())\n",
+    "\n",
+    "    def forward(self, positions):\n",
+    "        if self.checkpoint:\n",
+    "            position_embeddings = checkpoint(lambda x: self.pos_embeds[x], positions)\n",
+    "        else:\n",
+    "            position_embeddings = self.pos_embeds[positions]\n",
+    "        return F.normalize(position_embeddings, p=2, dim=-1) \n",
+    "\n",
+    "class CombinedPositionalEmbedding(nn.Module):\n",
+    "    def __init__(self, base, dims, head, n_ctx, theta_learnable=True, rot_learnable=True, \n",
+    "                 matrix_learnable=False, freq_learnable=True, checkpoint=False):\n",
+    "        super().__init__()\n",
+    "        self.rotary_embedding = CombinedRotaryEmbedding(base, dims, head, theta_learnable, \n",
+    "                                                        rot_learnable, matrix_learnable, freq_learnable)\n",
+    "        self.sinusoidal_embedding = SinusoidalEmbedding(n_ctx, dims, checkpoint)\n",
+    "\n",
+    "    def forward(self, x, positions, global_step=None):\n",
+    "        rotary_embed = self.rotary_embedding(x, global_step)\n",
+    "        sinusoidal_embed = self.sinusoidal_embedding(positions)\n",
+    "        \n",
+    "        combined_embedding = rotary_embed + sinusoidal_embed\n",
+    "        return combined_embedding"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\n",
+    "class MultiheadAttention(nn.Module):\n",
+    "    use_sdpa = True\n",
+    "\n",
+    "    def __init__(self, base, dims, head, max_dist):\n",
+    "        super().__init__()\n",
+    "        assert dims % head == 0, \"dims must be divisible by head\"\n",
+    "        self.head = head\n",
+    "        self.h_dim = dims // head\n",
+    "        assert self.h_dim % 2 == 0, \"Head dimension must be even for rotary embeddings\"\n",
+    "\n",
+    "        self.query = nn.Linear(dims, dims)\n",
+    "        self.key = nn.Linear(dims, dims, bias=False)\n",
+    "        self.value = nn.Linear(dims, dims)\n",
+    "        self.out = nn.Linear(dims, dims)\n",
+    "\n",
+    "    def forward(self, x, xa = None, mask = None, kv_cache = None):\n",
+    "\n",
+    "        q = self.query(x)\n",
+    "\n",
+    "        if kv_cache is None or xa is None or self.key not in kv_cache:\n",
+    "            k = self.key(x if xa is None else xa)\n",
+    "            v = self.value(x if xa is None else xa)\n",
+    "\n",
+    "        else:\n",
+    "            k = kv_cache[self.key]\n",
+    "            v = kv_cache[self.value]\n",
+    "        wv, qk = self.qkv_attention(q=q, k=k, v=v, mask=mask)\n",
+    "\n",
+    "        out = self.out(wv)\n",
+    "        return out, qk\n",
+    "    \n",
+    "    def qkv_attention(self, q: Tensor, k: Tensor, v: Tensor, mask: Optional[Tensor] = None):\n",
+    "        \n",
+    "        n_batch, n_ctx, dims = q.shape\n",
+    "        scale = (dims // self.head) ** -0.25\n",
+    "        q = q.view(*q.shape[:2], self.head, -1).permute(0, 2, 1, 3)\n",
+    "        k = k.view(*k.shape[:2], self.head, -1).permute(0, 2, 1, 3)\n",
+    "        v = v.view(*v.shape[:2], self.head, -1).permute(0, 2, 1, 3)\n",
+    "\n",
+    "        if MultiheadAttention.use_sdpa:\n",
+    "            a = scaled_dot_product_attention(query=q, key=k, value=v, is_causal=mask is not None and n_ctx > 1)\n",
+    "            out = a.permute(0, 2, 1, 3).flatten(start_dim=2)\n",
+    "            qk = None\n",
+    "        else:\n",
+    "            qk = (q * scale) @ (k * scale).transpose(-1, -2)\n",
+    "            if mask is not None:\n",
+    "                qk = qk + mask[:n_ctx, :n_ctx]\n",
+    "            qk = qk.float()\n",
+    "\n",
+    "            w = F.softmax(qk, dim=-1).to(dtype=q.dtype)\n",
+    "            out = (w @ v).permute(0, 2, 1, 3).flatten(start_dim=2)\n",
+    "            qk = qk.detach()\n",
+    "\n",
+    "        return out, qk\n",
+    "    "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\n",
+    "class AdaptiveSpanAttention(nn.Module):\n",
+    "    def __init__(self, base, dims, head, max_dist, sharpen, win_size, max_span, temp_scale=0.01):\n",
+    "        super().__init__()\n",
+    "        self.max_dist = max_dist\n",
+    "        self.win_size = win_size\n",
+    "        self.max_span = max_span\n",
+    "        self.temp_scale = temp_scale\n",
+    "        self.multihead_attn = MultiheadAttention(base=base, dims=dims, head=head, max_dist=max_dist)\n",
+    "        self.span_scale = nn.Parameter(torch.tensor(1.0))\n",
+    "        self.sharpen = sharpen\n",
+    "\n",
+    "    def forward(self, query, key, value, span_scale):\n",
+    "        span_len = int(self.max_span * span_scale.mean().item())\n",
+    "        span_len = min(span_len, query.shape[1], key.shape[1], value.shape[1])\n",
+    "        eff_span = min(span_len, self.max_dist)\n",
+    "\n",
+    "        q_span = query[:, :eff_span, :]\n",
+    "        k_span = key[:, :eff_span, :]\n",
+    "        v_span = value[:, :eff_span, :]\n",
+    "\n",
+    "        batch_size, _, dims = query.shape\n",
+    "        scale = (dims // self.multihead_attn.head) ** -0.25\n",
+    "\n",
+    "        q = q_span.view(q_span.shape[0], q_span.shape[1], self.multihead_attn.head, -1).permute(0, 2, 1, 3)\n",
+    "        k = k_span.view(k_span.shape[0], k_span.shape[1], self.multihead_attn.head, -1).permute(0, 2, 1, 3)\n",
+    "        v = v_span.view(v_span.shape[0], v_span.shape[1], self.multihead_attn.head, -1).permute(0, 2, 1, 3)\n",
+    "\n",
+    "        if self.sharpen:\n",
+    "            temperature = 1.0 + self.temp_scale * (1.0 - span_scale.mean().item())\n",
+    "        else:\n",
+    "            temperature = 0.5 + self.temp_scale * span_scale.mean().item()\n",
+    "\n",
+    "        attn_scores = torch.matmul(q, k.transpose(-2, -1))\n",
+    "        attn_weights = torch.softmax((attn_scores / temperature) * scale, dim=-1)\n",
+    "        attn_out = torch.matmul(attn_weights, v)\n",
+    "        attn_out = attn_out.permute(0, 2, 1, 3).flatten(start_dim=2)\n",
+    "        attn_out = attn_out.contiguous().view(batch_size, eff_span, dims)\n",
+    "\n",
+    "        return attn_out, attn_weights\n",
+    "\n",
+    "\n",
+    "class SpanPredictor(nn.Module):\n",
+    "    def __init__(self, dims):\n",
+    "        super().__init__()\n",
+    "        self.linear = nn.Linear(in_features=dims, out_features=1)\n",
+    "\n",
+    "    def forward(self, global_out):\n",
+    "        scale = torch.sigmoid(self.linear(global_out))\n",
+    "        return scale\n",
+    "\n",
+    "\n",
+    "class HybridAttention(nn.Module):\n",
+    "    def __init__(self, base, dims, head, max_dist, sharpen, win_size=32, max_span=32, slid_win=32):\n",
+    "        super().__init__()\n",
+    "        self.max_dist = max_dist\n",
+    "        self.win_size = win_size\n",
+    "        self.max_span = max_span\n",
+    "        self.slid_win = slid_win\n",
+    "\n",
+    "        self.span_pred = SpanPredictor(dims=dims)\n",
+    "        self.dist_local = max_dist\n",
+    "        self.dist_global = max_dist\n",
+    "\n",
+    "        self.attn_local = AdaptiveSpanAttention(base=base, dims=dims, head=head, max_dist=max_dist, sharpen=sharpen, win_size=win_size, max_span=max_span)\n",
+    "        self.attn_global = MultiheadAttention(base=base, dims=dims, head=head, max_dist=self.dist_global)\n",
+    "        self.ln_local = LayerNorm(normalized_shape=dims)\n",
+    "        self.ln_global = LayerNorm(normalized_shape=dims)\n",
+    "        self.projection = Linear(in_features=2 * dims, out_features=dims)\n",
+    "\n",
+    "    def forward(self, x, new_dist=None, new_base=None, xa=None, mask=None, kv_cache=None):\n",
+    "        local = self.ln_local(x)\n",
+    "        globe = self.ln_global(x)\n",
+    "\n",
+    "        globe_out, _ = self.attn_global(globe, globe, globe)\n",
+    "\n",
+    "        span_scale = self.span_pred(globe_out.mean(dim=1))\n",
+    "\n",
+    "        win_size = max(1, int(self.slid_win * span_scale.mean().item()))\n",
+    "        span_len = max(1, int(self.max_span * span_scale.mean().item()))\n",
+    "\n",
+    "        effective_max_dist = min(self.max_dist, local.size(1))\n",
+    "        local_max_dist = min(self.dist_local, span_len, win_size)\n",
+    "        globe_max_dist = effective_max_dist\n",
+    "\n",
+    "        self.attn_local.max_dist = local_max_dist\n",
+    "        self.attn_global.max_dist = globe_max_dist\n",
+    "\n",
+    "        local_out = self.slide_win(x=local, win_size=win_size, span_len=span_len, span_scale=span_scale)\n",
+    "\n",
+    "        combined = torch.cat(tensors=[local_out, globe_out], dim=-1)\n",
+    "        x = self.projection(combined)\n",
+    "\n",
+    "        return x\n",
+    "\n",
+    "    def slide_win(self, x, win_size, span_len, span_scale):\n",
+    "        batch_size, seq_len, dims = x.size()\n",
+    "        out = torch.zeros_like(x, device=x.device)\n",
+    "\n",
+    "        for i in range(0, seq_len, win_size):\n",
+    "            end = min(i + win_size, seq_len)\n",
+    "            query = x[:, i:end, :]\n",
+    "\n",
+    "            start = max(0, i - span_len + win_size)\n",
+    "            key = x[:, start:i + span_len, :]\n",
+    "            value = x[:, start:i + span_len, :]\n",
+    "            attn_out, _ = self.attn_local(query, key, value, span_scale)\n",
+    "            out[:, i:end, :] = attn_out\n",
+    "\n",
+    "        return out\n",
+    "\n",
+    "\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\n",
+    "class ResidualAttention(nn.Module):\n",
+    "    def __init__(self, base, dims, head, max_dist, win_size, max_span, hybrid, checkpoint, cross, sharpen):\n",
+    "        super().__init__()\n",
+    "\n",
+    "        if hybrid:\n",
+    "            # print(\"HybridDrive ON\")\n",
+    "            self.attn = HybridAttention(base=base, dims=dims, head=head, max_dist=max_dist, sharpen=sharpen)\n",
+    "            self.attn_ln = LayerNorm(normalized_shape=dims)\n",
+    "        else:\n",
+    "            self.attn = MultiheadAttention(base=base, dims=dims, head=head, max_dist=max_dist)\n",
+    "            self.attn_ln = LayerNorm(normalized_shape=dims)\n",
+    "\n",
+    "        n_mlp = dims * 4\n",
+    "        self.mlp = nn.Sequential(Linear(in_features=dims, out_features=n_mlp), nn.GELU(), Linear(in_features=n_mlp, out_features=dims))\n",
+    "        self.mlp_ln = LayerNorm(normalized_shape=dims)\n",
+    "\n",
+    "    def forward(self, x, mask=None, kv_cache=None):\n",
+    "        x = self._attn_forward(x=x, mask=mask, kv_cache=kv_cache)\n",
+    "        x = self._mlp_forward(x=x)\n",
+    "        return x\n",
+    "\n",
+    "    def _attn_forward(self, x, mask=None, kv_cache=None):\n",
+    "        residual = x\n",
+    "        x = self.attn_ln(x)\n",
+    "\n",
+    "        if isinstance(self.attn, HybridAttention):\n",
+    "            attn_output = self.attn(x)  \n",
+    "\n",
+    "            x = residual + attn_output\n",
+    "        else:\n",
+    "            attn_output, _ = self.attn(x, mask=mask, kv_cache=kv_cache)  \n",
+    "            x = residual + attn_output\n",
+    "        return x\n",
+    "\n",
+    "    def _mlp_forward(self, x):\n",
+    "        residual = x\n",
+    "        x = self.mlp_ln(x)\n",
+    "        return residual + self.mlp(x)\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\n",
+    "class AudioEncoder(nn.Module):\n",
+    "    def __init__(self, base, mels, dims, head, n_layer, n_ctx, max_dist,\n",
+    "                 win_size, max_span, hybrid, checkpoint, cross, sharpen):\n",
+    "        super().__init__()\n",
+    "        self.conv1 = Conv1d(in_channels=mels, out_channels=dims, kernel_size=3, padding=1)\n",
+    "        self.conv2 = Conv1d(in_channels=dims, out_channels=dims, kernel_size=3, stride=2, padding=1)\n",
+    "        self.pos_embed = SinusoidalEmbedding(n_ctx=n_ctx, dims=dims, checkpoint=checkpoint)\n",
+    "        self.checkpoint = checkpoint\n",
+    "\n",
+    "        self.givens_rotary = CombinedRotaryEmbedding(base=base, dims=dims, head=head)\n",
+    "\n",
+    "        self.blocks = nn.ModuleList(modules=[ResidualAttention(base=base, dims=dims, head=head, max_dist=max_dist, win_size=win_size, max_span=max_span, hybrid=hybrid, checkpoint=checkpoint, cross=cross, sharpen=sharpen) for _ in range(n_layer)])\n",
+    "\n",
+    "        self.ln_post = LayerNorm(normalized_shape=dims)\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        if self.checkpoint:\n",
+    "            x = checkpoint(self._conv_forward, x)\n",
+    "        else:\n",
+    "            x = self._conv_forward(x)\n",
+    "\n",
+    "        for block in self.blocks:\n",
+    "            if self.checkpoint:\n",
+    "                x = checkpoint(block, x)\n",
+    "            else:\n",
+    "                x = block(x)\n",
+    "        return self.ln_post(x)\n",
+    "\n",
+    "    def _conv_forward(self, x):\n",
+    "        x = F.gelu(self.conv1(x))\n",
+    "        x = F.gelu(self.conv2(x))\n",
+    "        x = x.permute(0, 2, 1)\n",
+    "    \n",
+    "        p = self.pos_embed(torch.arange(end=x.size(dim=1), device=x.device)).unsqueeze(0)\n",
+    "        x = (x + p).to(x.dtype)\n",
+    "        x = self.givens_rotary(x)\n",
+    "        return x\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\n",
+    "\n",
+    "class TextDecoder(nn.Module):\n",
+    "    def __init__(self, base, vocab, dims, head, n_layer, n_ctx, max_dist,\n",
+    "                 win_size, max_span, hybrid, checkpoint, cross, sharpen):\n",
+    "        super().__init__()\n",
+    "        \n",
+    "        self.tok_embed = nn.Embedding(num_embeddings=vocab, embedding_dim=dims)\n",
+    "        self.pos_embed = SinusoidalEmbedding(n_ctx=n_ctx, dims=dims, checkpoint=checkpoint)\n",
+    "        self.checkpoint = checkpoint\n",
+    "\n",
+    "        self.givens_rotary = CombinedRotaryEmbedding(base=base, dims=dims, head=head)\n",
+    "\n",
+    "        self.blocks = nn.ModuleList(modules=[ResidualAttention(base=base, dims=dims, head=head, max_dist=max_dist, win_size=win_size, max_span=max_span, hybrid=hybrid, checkpoint=checkpoint, cross=cross, sharpen=sharpen) for _ in range(n_layer)])\n",
+    "\n",
+    "        self.ln_post = LayerNorm(normalized_shape=dims)\n",
+    "        self.ln = LayerNorm(normalized_shape=dims)\n",
+    "\n",
+    "        mask = torch.empty(n_ctx, n_ctx).fill_(value=-np.inf).triu_(diagonal=1)\n",
+    "        self.register_buffer(name=\"mask\", tensor=mask, persistent=False)\n",
+    "        self.mask=mask\n",
+    "\n",
+    "    def forward(self, x, xa, kv_cache=None):\n",
+    "        if self.checkpoint:\n",
+    "            x = checkpoint(self._embedding_forward, x, xa, kv_cache)\n",
+    "        else:\n",
+    "            x = self._embedding_forward(x=x, xa=xa, kv_cache=kv_cache)\n",
+    "\n",
+    "        for block in self.blocks:\n",
+    "            if self.checkpoint:\n",
+    "                x = checkpoint(block, x, self.mask, kv_cache)\n",
+    "            else:\n",
+    "                x = block(x, self.mask, kv_cache)\n",
+    "\n",
+    "        x = self.ln(x)\n",
+    "        x = (x @ torch.transpose(input=self.tok_embed.weight.to(dtype=x.dtype), dim0=0, dim1=1)).float()\n",
+    "        return x\n",
+    "    \n",
+    "    def _embedding_forward(self, x, xa, kv_cache):\n",
+    "        offset = next(iter(kv_cache.values())).shape[1] if kv_cache else 0\n",
+    "        positions = torch.arange(x.shape[1], device=x.device) + offset\n",
+    "        pos_emb = self.pos_embed(positions).unsqueeze(0)\n",
+    "        x = self.tok_embed(x) + pos_emb\n",
+    "        x = self.givens_rotary(x)\n",
+    "        return x"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class EchoConfig(PretrainedConfig):\n",
+    "    model_type = \"Echo\"\n",
+    "    def __init__(\n",
+    "        self,\n",
+    "        checkpoint=False,\n",
+    "        cross=False,\n",
+    "        hybrid=False,\n",
+    "        sharpen=False,\n",
+    "        a_ctx=1500,\n",
+    "        a_head=16,\n",
+    "        a_layer=8,\n",
+    "        a_dims=1024,\n",
+    "        mels=128,\n",
+    "        t_ctx=448,\n",
+    "        t_head=8,\n",
+    "        t_layer=8,\n",
+    "        t_dims=1024,\n",
+    "        win_size=64,\n",
+    "        max_span=64,\n",
+    "        max_dist=64,\n",
+    "        base=10000,\n",
+    "        pad_token_id=50257,\n",
+    "        unk_token_id=50257,\n",
+    "        vocab=51865,\n",
+    "        eos_token_id=50257,\n",
+    "        bos_token_id=50257,\n",
+    "        decoder_start_token_id=50258,\n",
+    "        **kwargs,\n",
+    "    ):\n",
+    "        \n",
+    "        super().__init__(**kwargs) \n",
+    "        self.base = base\n",
+    "        self.bos_token_id = bos_token_id\n",
+    "        self.checkpoint = checkpoint\n",
+    "        self.cross = cross\n",
+    "        self.decoder_start_token_id = decoder_start_token_id\n",
+    "        self.eos_token_id = eos_token_id\n",
+    "        self.hybrid = hybrid\n",
+    "        self.max_dist = max_dist\n",
+    "        self.max_span = max_span\n",
+    "        self.a_ctx = a_ctx\n",
+    "        self.a_head = a_head\n",
+    "        self.a_layer = a_layer\n",
+    "        self.a_dims = a_dims\n",
+    "        self.mels = mels\n",
+    "        self.t_ctx = t_ctx\n",
+    "        self.t_head = t_head\n",
+    "        self.t_layer = t_layer\n",
+    "        self.t_dims = t_dims\n",
+    "        self.pad_token_id = pad_token_id\n",
+    "        self.unk_token_id = unk_token_id\n",
+    "        self.vocab = vocab\n",
+    "        self.win_size = win_size\n",
+    "        self.sharpen=sharpen\n",
+    "\n",
+    "class Echo(nn.Module):\n",
+    "    def __init__(self, config: EchoConfig):\n",
+    "        super().__init__()\n",
+    "        self.config = config\n",
+    "            \n",
+    "        self.encoder = AudioEncoder(\n",
+    "            base=self.config.base,\n",
+    "            mels=self.config.mels,\n",
+    "            dims=self.config.a_dims, \n",
+    "            head=self.config.a_head,\n",
+    "            n_layer=self.config.a_layer,\n",
+    "            n_ctx=self.config.a_ctx,\n",
+    "            max_dist=self.config.max_dist,\n",
+    "            win_size=self.config.win_size,  \n",
+    "            max_span=self.config.max_span,\n",
+    "            hybrid=self.config.hybrid,\n",
+    "            checkpoint=self.config.checkpoint,\n",
+    "            cross=self.config.cross,\n",
+    "            sharpen=self.config.sharpen,\n",
+    "        )\n",
+    "\n",
+    "        self.decoder = TextDecoder(\n",
+    "            base=self.config.base,\n",
+    "            vocab=self.config.vocab,\n",
+    "            dims=self.config.t_dims, \n",
+    "            head=self.config.t_head,\n",
+    "            n_layer=self.config.t_layer,\n",
+    "            n_ctx=self.config.t_ctx,\n",
+    "            max_dist=self.config.max_dist,\n",
+    "            win_size=self.config.win_size,  \n",
+    "            max_span=self.config.max_span,\n",
+    "            hybrid=self.config.hybrid,\n",
+    "            checkpoint=self.config.checkpoint,\n",
+    "            cross=self.config.cross,\n",
+    "            sharpen=self.config.sharpen,\n",
+    "        )\n",
+    "\n",
+    "\n",
+    "        all_heads = torch.zeros(self.config.t_layer, self.config.t_head, dtype=torch.bool) \n",
+    "        all_heads[self.config.t_layer // 2:] = True\n",
+    "        self.register_buffer(name=\"alignment_heads\", tensor=all_heads.to_sparse(), persistent=False)\n",
+    "\n",
+    "        self.base = self.config.base\n",
+    "        self.win_size = self.config.win_size\n",
+    "        self.adjust_counter = 0\n",
+    "        self.best_loss = float('inf')\n",
+    "        self.kv_cache = {}\n",
+    "\n",
+    "\n",
+    "    @property\n",
+    "    def device(self):\n",
+    "        return next(self.parameters()).device\n",
+    "\n",
+    "    def embed_audio(self, mel: torch.Tensor):\n",
+    "        return self.encoder(mel)\n",
+    "\n",
+    "    def logits(self, tokens: torch.Tensor, audio_features: torch.Tensor):\n",
+    "        return self.decoder(tokens, audio_features)\n",
+    "\n",
+    "    def update_window(self, new_window):\n",
+    "        self.win_size = new_window\n",
+    "        for module in self.modules(): \n",
+    "            if isinstance(module, HybridAttention):\n",
+    "                module.update_window(self.win_size)\n",
+    "\n",
+    "    def adjust_window(self, loss, factor=1.00005):\n",
+    "        if self.adjust_counter % 10 == 0:\n",
+    "            if loss < self.best_loss:\n",
+    "                new_window = self.win_size * factor\n",
+    "            else:\n",
+    "                new_window = self.win_size / factor\n",
+    "            self.update_window(new_window=new_window)\n",
+    "            self.best_loss = loss\n",
+    "            self.adjust_counter += 1\n",
+    "            return new_window\n",
+    "        return self.win_size\n",
+    "\n",
+    "    def adjust_base(self, loss, factor=1.0025) -> float | int:\n",
+    "                if self.adjust_counter % 25 == 0:\n",
+    "                    if loss < self.best_loss:\n",
+    "                        new_base=self.base*factor\n",
+    "                    else:\n",
+    "                        new_base=self.base/factor\n",
+    "                    self.update_base(new_base=new_base)\n",
+    "                    self.base=new_base\n",
+    "                    self.best_loss=loss\n",
+    "                self.adjust_counter += 1\n",
+    "                return self.base\n",
+    "            \n",
+    "    def update_base(self, new_base):\n",
+    "        self.new_base=new_base\n",
+    "        for name, module in self.encoder.named_modules():\n",
+    "            if isinstance(module, (CombinedRotaryEmbedding)):\n",
+    "                module.update_base(new_base=self.new_base)\n",
+    "\n",
+    "    @staticmethod\n",
+    "    def shift_tokens_right(input_ids, pad_token_id, decoder_start_token_id):\n",
+    "        shifted_input_ids = input_ids.new_zeros(input_ids.shape)\n",
+    "        shifted_input_ids[:, 1:] = input_ids[:, :-1].clone() \n",
+    "        shifted_input_ids[:, 0] = decoder_start_token_id\n",
+    "        shifted_input_ids.masked_fill_(shifted_input_ids == -100, pad_token_id)\n",
+    "        return shifted_input_ids\n",
+    "\n",
+    "    def forward(self, input_features, labels=None, dec_input_ids=None) -> dict[str, Any | None]:\n",
+    "        if labels is not None:\n",
+    "            if dec_input_ids is None:\n",
+    "                dec_input_ids = self.shift_tokens_right(\n",
+    "                    input_ids=labels, pad_token_id=self.config.pad_token_id, decoder_start_token_id=self.config.decoder_start_token_id\n",
+    "                )\n",
+    "\n",
+    "        encoded_features = self.encoder(input_features).to(self.device)  \n",
+    "        logits = self.decoder(dec_input_ids, encoded_features)\n",
+    "\n",
+    "        loss = None\n",
+    "        if labels is not None:\n",
+    "            loss_fct = nn.CrossEntropyLoss(ignore_index=-100)\n",
+    "            labels = labels.to(logits.device).long()\n",
+    "            loss = loss_fct(logits.view(-1, self.config.vocab), labels.view(-1))\n",
+    "\n",
+    "            self.adjust_window(loss.item())\n",
+    "            # self.adjust_base(loss=loss.item())\n",
+    "        return {\"loss\": loss, \"logits\": logits}\n",
+    "\n",
+    "    def reset_parameters(self):\n",
+    "        for name, module in self.encoder.named_modules():\n",
+    "            if isinstance(module, CombinedRotaryEmbedding):\n",
+    "                module.reset_parameters()\n",
+    "        \n",
+    "    def _initialize_weights(self, module):\n",
+    "            nn.init.normal_(tensor=self.decoder.tok_embed.weight, mean=0.0, std=0.02)\n",
+    "            nn.init.constant_(tensor=self.decoder.ln.weight, val=1)\n",
+    "            nn.init.constant_(tensor=self.decoder.ln.bias, val=0)\n",
+    "            nn.init.xavier_normal_(tensor=self.encoder.conv1.weight)\n",
+    "            nn.init.zeros_(tensor=self.encoder.conv1.bias)\n",
+    "            nn.init.kaiming_normal_(tensor=self.encoder.conv2.weight, mode='fan_out', nonlinearity='relu')\n",
+    "            nn.init.zeros_(tensor=self.encoder.conv2.bias)\n",
+    "            nn.init.constant_(tensor=self.encoder.ln_post.weight, val=1)\n",
+    "            nn.init.constant_(tensor=self.encoder.ln_post.bias, val=0)\n",
+    "\n",
+    "            for block in self.decoder.blocks:\n",
+    "                for layer in block.children():\n",
+    "                    if isinstance(layer, nn.Linear):\n",
+    "                        nn.init.xavier_normal_(tensor=layer.weight)\n",
+    "                        nn.init.zeros_(tensor=layer.bias)\n",
+    "                    if isinstance(layer, LayerNorm):\n",
+    "                        nn.init.constant_(tensor=layer.weight, val=1)\n",
+    "            \n",
+    "            for block in self.encoder.blocks:\n",
+    "                for layer in block.children():\n",
+    "                    if isinstance(layer, nn.Linear):\n",
+    "                        nn.init.xavier_normal_(tensor=layer.weight)\n",
+    "                        nn.init.zeros_(tensor=layer.bias)\n",
+    "                    if isinstance(layer, LayerNorm):\n",
+    "                        nn.init.constant_(tensor=layer.weight, val=1)\n",
+    "\n",
+    "            for module in self.encoder.named_modules():\n",
+    "                if isinstance(module, CombinedRotaryEmbedding):\n",
+    "                    nn.init.constant_(tensor=module.thetas, val=1)\n",
+    "                    nn.init.constant_(tensor=module.r_matrix, val=1)\n",
+    "                    nn.init.constant_(tensor=module.r_pairs, val=1)\n",
+    "                    nn.init.constant_(tensor=module.inv_freq, val=1)\n",
+    "\n",
+    "    def apply_initialization(self, module):\n",
+    "        self._initialize_weights(module=module)\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\n",
+    "from datetime import datetime\n",
+    "log_dir = os.path.join('./output/Echo/', datetime.now().strftime(format='%m-%d_%H'))\n",
+    "os.makedirs(name=log_dir, exist_ok=True)\n",
+    "\n",
+    "config = EchoConfig(\n",
+    "    checkpoint=False,\n",
+    "    cross=False,\n",
+    "    hybrid=False,\n",
+    "    sharpen=False,\n",
+    "    audio_ctx=1500,\n",
+    "    audio_head=4,\n",
+    "    audio_layer=4,\n",
+    "    audio_dims=512,\n",
+    "    mels=128,\n",
+    "    text_ctx=448,\n",
+    "    text_head=4,\n",
+    "    text_layer=4,\n",
+    "    text_dims=512,\n",
+    "    win_size=16,\n",
+    "    max_span=16,\n",
+    "    max_dist=16,\n",
+    "    base=50000,\n",
+    "    pad_token_id=50257,\n",
+    "    unk_token_id=50257,\n",
+    "    vocab=51865,\n",
+    "    eos_token_id=50257,\n",
+    "    bos_token_id=50257,\n",
+    "    decoder_start_token_id=50258,\n",
+    ")\n",
+    "\n",
+    "model = Echo(config=config).to(device=device)\n",
+    "model.apply_initialization(module=model)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "feature_extractor = WhisperFeatureExtractor.from_pretrained(\n",
+    "    pretrained_model_name_or_path=\"openai/whisper-small\", \n",
+    "    feature_size=128, sample_rate=160000, do_normalize=True)\n",
+    "\n",
+    "tokenizer = WhisperTokenizerFast.from_pretrained(\n",
+    "    pretrained_model_name_or_path=\"openai/whisper-small\", \n",
+    "    language=\"en\", task=\"transcribe\")\n",
+    "\n",
+    "processor = WhisperProcessor.from_pretrained(\n",
+    "    pretrained_model_name_or_path=\"openai/whisper-small\", \n",
+    "    feature_size=128, sample_rate=160000, do_normalize=True, \n",
+    "    language=\"en\", task=\"transcribe\")\n",
+    "\n",
+    "class GradientClippingCallback(TrainerCallback):\n",
+    "    def on_step_end(self, args, dims, control, **kwargs):\n",
+    "        torch.nn.utils.clip_grad_norm_(parameters=kwargs[\"model\"].parameters(), max_norm=0.98)\n",
+    "\n",
+    "@dataclass\n",
+    "class DataCollatorSpeechSeq2SeqWithPadding:\n",
+    "    processor: Any\n",
+    "    decoder_start_token_id: int\n",
+    "\n",
+    "    def __call__(self, features: List[Dict[str, Union[List[int], torch.Tensor]]]) -> Dict[str, torch.Tensor]:\n",
+    "        input_features = [{\"input_features\": feature[\"input_features\"]} for feature in features]\n",
+    "        batch = self.processor.feature_extractor.pad(input_features, return_tensors=\"pt\")\n",
+    "        label_features = [{\"input_ids\": feature[\"labels\"]} for feature in features]\n",
+    "        labels_batch = self.processor.tokenizer.pad(label_features, return_tensors=\"pt\")\n",
+    "        labels = labels_batch[\"input_ids\"].masked_fill(labels_batch.attention_mask.ne(1), -100)\n",
+    "        if (labels[:, 0] == self.decoder_start_token_id).all().cpu().item():\n",
+    "            labels = labels[:, 1:]\n",
+    "        batch[\"labels\"] = labels\n",
+    "        return batch\n",
+    "\n",
+    "def get_length_of_dataset(dataset):\n",
+    "    length = 0\n",
+    "    for item in dataset:\n",
+    "        length += len(item[\"audio\"][\"array\"]) / item[\"audio\"][\"sampling_rate\"]\n",
+    "    return length / 3600  \n",
+    "\n",
+    "def prepare_dataset(batch):\n",
+    "    audio = batch[\"audio\"]\n",
+    "    batch[\"input_features\"] = feature_extractor(audio[\"array\"], sampling_rate=audio[\"sampling_rate\"]).input_features[0]\n",
+    "    batch[\"labels\"] = tokenizer(batch[\"sentence\"]).input_ids\n",
+    "    return batch\n",
+    "\n",
+    "data_collator = DataCollatorSpeechSeq2SeqWithPadding(processor=processor, decoder_start_token_id=config.decoder_start_token_id)\n",
+    "\n",
+    "datasets = IterableDatasetDict()\n",
+    "\n",
+    "datasets[\"train\"] = load_dataset(\n",
+    "    path=\"mozilla-foundation/common_voice_17_0\", token=\"\",\n",
+    "    name=\"en\", split=\"train\", streaming=True, trust_remote_code=True).take(10000)\n",
+    "\n",
+    "datasets[\"test\"] = load_dataset(\n",
+    "    path=\"mozilla-foundation/common_voice_17_0\", token=\"\", \n",
+    "    name=\"en\", split=\"test\", streaming=True, trust_remote_code=True).take(100)\n",
+    "\n",
+    "dataset = datasets.cast_column(column=\"audio\", feature=Audio(sampling_rate=16000))\n",
+    "\n",
+    "dataset = dataset.map(function=prepare_dataset, \n",
+    "                      remove_columns=list(next(iter(dataset.values())).features)).with_format(type=\"torch\")\n",
+    "\n",
+    "class MetricsCallback(TrainerCallback):\n",
+    "    def __init__(self, tb_writer, tokenizer, metric, optimizer, scheduler, log_every_n_steps=1):\n",
+    "        super().__init__()\n",
+    "        self.tb_writer = tb_writer\n",
+    "        self.tokenizer = tokenizer\n",
+    "        self.metric = metric\n",
+    "        self.optimizer = optimizer\n",
+    "        self.scheduler = scheduler\n",
+    "        self.log_every_n_steps = log_every_n_steps\n",
+    "        self.predictions = None\n",
+    "        self.label_ids = None\n",
+    "\n",
+    "    def compute_wer(self, pred_str, label_str):\n",
+    "        wer = 100 * self.metric.compute(predictions=pred_str, references=label_str)\n",
+    "        return wer\n",
+    "\n",
+    "    def on_evaluate(self, args, state, control, model, metrics=None, **kwargs):\n",
+    "        if metrics is not None:\n",
+    "            self.eval_loss = metrics.get('eval_loss')\n",
+    "\n",
+    "            current_learning_rate = self.optimizer.param_groups[0]['lr']\n",
+    "            if state.global_step % self.log_every_n_steps == 0:\n",
+    "                self.tb_writer.add_scalar('learning_rate', current_learning_rate, state.global_step)\n",
+    "                print(f\"Learning Rate: {current_learning_rate:.8f}\")\n",
+    "\n",
+    "                self.tb_writer.add_scalar('eval_loss', self.eval_loss, state.global_step)\n",
+    "\n",
+    "            for key, value in metrics.items():\n",
+    "                if key.startswith(\"eval_\"):\n",
+    "                    self.tb_writer.add_scalar(key, value, state.global_step)\n",
+    "\n",
+    "        if self.predictions is not None and self.label_ids is not None:\n",
+    "            pred_str = self.tokenizer.batch_decode(self.predictions, skip_special_tokens=True)\n",
+    "            label_str = self.tokenizer.batch_decode(self.label_ids, skip_special_tokens=True)\n",
+    "\n",
+    "            if state.global_step % self.log_every_n_steps == 0:\n",
+    "                total_samples = len(pred_str)\n",
+    "                random_indices = random.sample(range(total_samples), 1)\n",
+    "\n",
+    "                for sample_index in random_indices:\n",
+    "                    self.tb_writer.add_text(f\"Prediction_{sample_index}\", pred_str[sample_index], state.global_step)\n",
+    "                    self.tb_writer.add_text(f\"Label_{sample_index}\", label_str[sample_index], state.global_step)\n",
+    "                    print(f\"Evaluation: - Step {state.global_step} - Loss: {self.eval_loss:.2f}\")\n",
+    "                    print(f\"Prediction: {pred_str[sample_index]}\")\n",
+    "                    print(f\"Label: {label_str[sample_index]}\")\n",
+    "                    print(\"-\" * 10)\n",
+    "\n",
+    "        self.predictions = None\n",
+    "        self.label_ids = None\n",
+    "\n",
+    "def create_compute_metrics(callback_instance):\n",
+    "    def compute_metrics(eval_pred):\n",
+    "        pred_logits = eval_pred.predictions\n",
+    "        label_ids = eval_pred.label_ids\n",
+    "\n",
+    "        if isinstance(pred_logits, tuple):\n",
+    "            pred_ids = pred_logits[0]\n",
+    "        else:\n",
+    "            pred_ids = pred_logits\n",
+    "        if pred_ids.ndim == 3:\n",
+    "            pred_ids = np.argmax(pred_ids, axis=-1)\n",
+    "\n",
+    "        label_ids[label_ids == -100] = callback_instance.tokenizer.pad_token_id\n",
+    "        callback_instance.predictions = pred_ids\n",
+    "        callback_instance.label_ids = label_ids\n",
+    "        pred_str = callback_instance.tokenizer.batch_decode(pred_ids, skip_special_tokens=True)\n",
+    "        label_str = callback_instance.tokenizer.batch_decode(label_ids, skip_special_tokens=True)\n",
+    "        wer = 100 * callback_instance.metric.compute(predictions=pred_str, references=label_str)\n",
+    "        pred_flat = pred_ids.flatten()\n",
+    "        labels_flat = label_ids.flatten()\n",
+    "        mask = labels_flat != callback_instance.tokenizer.pad_token_id\n",
+    "\n",
+    "        accuracy = accuracy_score(y_true=labels_flat[mask], y_pred=pred_flat[mask])\n",
+    "        precision = precision_score(y_true=labels_flat[mask], y_pred=pred_flat[mask], average='weighted', zero_division=0)\n",
+    "        recall = recall_score(y_true=labels_flat[mask], y_pred=pred_flat[mask], average='weighted', zero_division=0)\n",
+    "        f1 = f1_score(y_true=labels_flat[mask], y_pred=pred_flat[mask], average='weighted', zero_division=0)\n",
+    "        return {\"wer\": wer, \"accuracy\": accuracy, \"precision\": precision, \"recall\": recall, \"f1\": f1}\n",
+    "    return compute_metrics\n",
+    "\n",
+    "metric = evaluate.load(path=\"wer\")\n",
+    "tb_writer = SummaryWriter(log_dir=log_dir)\n",
+    "\n",
+    "training_args = Seq2SeqTrainingArguments(\n",
+    "    output_dir=log_dir,\n",
+    "    per_device_train_batch_size=1,\n",
+    "    per_device_eval_batch_size=1,\n",
+    "    gradient_accumulation_steps=1,\n",
+    "    eval_accumulation_steps=1,\n",
+    "    tf32=True,\n",
+    "    bf16=True,\n",
+    "    eval_strategy=\"steps\",\n",
+    "    save_strategy=\"steps\",\n",
+    "    max_steps=10000,\n",
+    "    save_steps=10000,\n",
+    "    eval_steps=100,\n",
+    "    warmup_steps=100,\n",
+    "    logging_steps=10,\n",
+    "    logging_dir=log_dir + \"/logs_hf\",\n",
+    "    report_to=[\"tensorboard\"],\n",
+    "    load_best_model_at_end=False,\n",
+    "    metric_for_best_model=\"loss\",\n",
+    "    greater_is_better=False,\n",
+    "    push_to_hub=False,\n",
+    "    disable_tqdm=False,\n",
+    "    save_total_limit=1,\n",
+    "    remove_unused_columns=False,\n",
+    "    label_names=[\"labels\"],\n",
+    "    eval_on_start=True,\n",
+    ")\n",
+    "\n",
+    "class MaxFactor(Optimizer):\n",
+    "    def __init__(self, params, lr=0.01, beta2_decay=-0.8, eps=(None, 1e-3), d=1.0, \n",
+    "                 weight_decay=0.0, gamma=0.99, eps_rms=1e-8, maximize=False):\n",
+    "        \n",
+    "        defaults = dict(lr=lr, beta2_decay=beta2_decay, eps=eps, d=d, weight_decay=weight_decay, \n",
+    "                        gamma=gamma, eps_rms=eps_rms, maximize=maximize)\n",
+    "\n",
+    "        super().__init__(params, defaults)\n",
+    "\n",
+    "    @torch.no_grad()\n",
+    "    def step(self, closure=None):\n",
+    "        loss = None\n",
+    "        if closure is not None:\n",
+    "            with torch.enable_grad():\n",
+    "                loss = closure()\n",
+    "\n",
+    "        for group in self.param_groups:\n",
+    "            params_with_grad, grads, row_vars, col_vars, v, state_steps = [], [], [], [], [], []\n",
+    "            eps1, eps2 = group[\"eps\"]\n",
+    "            for p in group[\"params\"]:\n",
+    "                if p.grad is None:\n",
+    "                    continue\n",
+    "                grad = p.grad\n",
+    "                if grad.dtype in {torch.float16, torch.bfloat16}:\n",
+    "                    grad = grad.float()\n",
+    "\n",
+    "                state = self.state[p]\n",
+    "                if len(state) == 0:\n",
+    "                    state[\"step\"] = torch.tensor(0.0, dtype=torch.float32)\n",
+    "                    if p.grad.dim() > 1:\n",
+    "                        row_shape, col_shape = list(p.grad.shape), list(p.grad.shape)\n",
+    "                        row_shape[-1], col_shape[-2] = 1, 1\n",
+    "                        state[\"row_var\"], state[\"col_var\"] = p.grad.new_zeros(row_shape), p.grad.new_zeros(col_shape)\n",
+    "                    state[\"v\"] = torch.zeros_like(p, memory_format=torch.preserve_format)\n",
+    "\n",
+    "                row_vars.append(state.get(\"row_var\", None))\n",
+    "                col_vars.append(state.get(\"col_var\", None))\n",
+    "                v.append(state[\"v\"])\n",
+    "                state_steps.append(state[\"step\"])\n",
+    "                params_with_grad.append(p)\n",
+    "                grads.append(grad)\n",
+    "\n",
+    "            for i, param in enumerate(params_with_grad):\n",
+    "                grad = grads[i]\n",
+    "\n",
+    "                if group[\"maximize\"]:\n",
+    "                    grad = -grad\n",
+    "                step_t, row_var, col_var, vi = state_steps[i], row_vars[i], col_vars[i], v[i]\n",
+    "\n",
+    "                if eps1 is None:\n",
+    "                    eps1 = torch.finfo(param.dtype).eps\n",
+    "                    \n",
+    "                step_t += 1\n",
+    "                step_float = step_t.item()\n",
+    "                one_minus_beta2_t = step_float ** group[\"beta2_decay\"]\n",
+    "                rho_t = min(group[\"lr\"], 1 / (step_float ** 0.5))\n",
+    "                alpha = max(eps2, param.norm(2).item() / (param.numel() ** 0.5)) * rho_t\n",
+    "\n",
+    "                if group[\"weight_decay\"]!= 0:\n",
+    "                    param.mul_(1 - group[\"lr\"] * group[\"weight_decay\"])\n",
+    "\n",
+    "                if grad.dim() > 1:\n",
+    "                    row_mean = torch.norm(grad, dim=-1, keepdim=True).square_().div_(grad.size(-1))\n",
+    "                    row_var.lerp_(row_mean, one_minus_beta2_t)\n",
+    "                    col_mean = torch.norm(grad, dim=-2, keepdim=True).square_().div_(grad.size(-2))\n",
+    "                    col_var.lerp_(col_mean, one_minus_beta2_t)\n",
+    "                    var_estimate = row_var @ col_var\n",
+    "                    max_row_var = row_var.max(dim=-2, keepdim=True)[0]  \n",
+    "                    var_estimate.div_(max_row_var.clamp_(min=eps1))\n",
+    "\n",
+    "                else:\n",
+    "                    vi.mul_(group[\"gamma\"]).add_(1 - group[\"gamma\"], grad ** 2)\n",
+    "                    var_estimate = vi\n",
+    "                \n",
+    "                update = var_estimate.clamp_(min=eps1 * eps1).rsqrt_().mul_(grad)\n",
+    "                update = update.div_(torch.norm(update, float('inf')).clamp_(min=eps1))\n",
+    "                denom = max(1.0, update.norm(2).item() / ((update.numel() ** 0.5) * group[\"d\"]))\n",
+    "                param.add_(-alpha / denom * update.sign() * update.abs().max(dim=-1, keepdim=True)[0])\n",
+    "\n",
+    "        return loss\n",
+    "    \n",
+    "optimizer = MaxFactor(\n",
+    "    model.parameters(), \n",
+    "    lr=0.025,  \n",
+    "    beta2_decay=-0.8,\n",
+    "    eps=(None, 1e-4),\n",
+    "    d=1.0,\n",
+    "    weight_decay=0.0025,\n",
+    "    gamma=0.99, \n",
+    "    eps_rms=1e-8,\n",
+    "    maximize=False,\n",
+    "    )\n",
+    "\n",
+    "scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(\n",
+    "    optimizer=optimizer,\n",
+    "    T_max=training_args.max_steps,\n",
+    "    eta_min=0.0,\n",
+    "    last_epoch=-1  \n",
+    ")\n",
+    "\n",
+    "metrics_callback = MetricsCallback(tb_writer=tb_writer, tokenizer=tokenizer, metric=metric, optimizer=optimizer, scheduler=scheduler, log_every_n_steps=10)\n",
+    "compute_metrics = create_compute_metrics(callback_instance=metrics_callback)\n",
+    "\n",
+    "trainer = Seq2SeqTrainer(\n",
+    "    args=training_args,\n",
+    "    model=model,\n",
+    "    train_dataset=dataset[\"train\"],\n",
+    "    eval_dataset=dataset[\"test\"],\n",
+    "    data_collator=data_collator,\n",
+    "    compute_metrics=compute_metrics,\n",
+    "    processing_class=feature_extractor,\n",
+    "    callbacks=[metrics_callback],\n",
+    "    optimizers=(optimizer, scheduler)\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\n",
+    "trainer.train(resume_from_checkpoint=False)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from tensorboard import program\n",
+    "log_dir = \"D:/new/tensorboard3\" \n",
+    "tb = program.TensorBoard()\n",
+    "tb.configure(argv=[None, '--logdir', log_dir])\n",
+    "url = tb.launch()\n",
+    "print(f\"TensorBoard started at {url}\")"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.12.8"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}