model.py

import base64, gzip, math, os, functools, warnings, numpy as np, torch, transformers, aiohttp, torch.nn.functional as F, evaluate, json, random
from torch import Tensor, amp, optim, nn
from torch.utils.checkpoint import checkpoint
from torch.utils.tensorboard.writer import SummaryWriter
from threading import Thread
from typing import Dict, Optional, Tuple, Union, List, Any
from dataclasses import dataclass
from transformers import (Seq2SeqTrainer, Seq2SeqTrainingArguments, PretrainedConfig, TrainerCallback, WhisperProcessor, WhisperFeatureExtractor, WhisperTokenizerFast)
from torch.optim import Optimizer
import evaluate
from evaluate import module
from sklearn.metrics import accuracy_score, precision_score, f1_score, recall_score
from datasets import load_dataset, IterableDatasetDict, Audio, load_from_disk
from torch.nn.functional import scaled_dot_product_attention
transformers.utils.logging.set_verbosity_error()
warnings.filterwarnings(action="ignore")
warnings.warn = lambda *args, **kwargs: None
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
dtype = torch.float32


class Linear(nn.Linear):
    def forward(self, x: Tensor) -> Tensor:# type: ignore
        return F.linear(x, self.weight.to(x.dtype),
                         None if self.bias is None else self.bias.to(x.dtype))

class Conv1d(nn.Conv1d):
    def _conv_forward(self, x: Tensor, weight: Tensor, bias: Optional[Tensor] = None) -> Tensor:# type: ignore
        return super()._conv_forward(x, weight.to(x.dtype),
                                     None if bias is None else bias.to(x.dtype))

class LayerNorm(nn.LayerNorm):
    def forward(self, x: Tensor) -> Tensor:  # type: ignore
        return super().forward(x.float()).type(x.dtype) 



class CombinedRotaryEmbedding(nn.Module):
    def __init__(self, base, dims, head, theta_learnable=True, rot_learnable=True,
                 matrix_learnable=False, freq_learnable=True):
        super(CombinedRotaryEmbedding, self).__init__()

        self.base = base
        self.dims = dims
        self.head = head

        self.h_dim = self.dims // self.head
        self.rot = (self.dims // self.head) // 2

        self.thetas = nn.Parameter(torch.zeros(self.rot))
        self.r_pairs = nn.Parameter(data=torch.rand(self.rot, 2) * self.h_dim)

        self.theta_scale = nn.Parameter(torch.ones(1), requires_grad=theta_learnable)
        self.rot_scale = nn.Parameter(torch.ones(1), requires_grad=rot_learnable)

        self.r_matrix = nn.Parameter(torch.eye(n=self.h_dim), requires_grad=matrix_learnable)

        freq_data = 1.0 / (self.base ** (torch.arange(start=0, end=self.h_dim, step=2).float() / self.h_dim))
        self.inv_freq = nn.Parameter(freq_data, requires_grad=freq_learnable)

        self.orthogonal_reg_weight = 0.01

    def blended_rotation_matrix(self, dims, i, j, theta):
        G = torch.eye(dims).to(theta.device)
        G[i, i] = torch.cos(theta)
        G[i, j] = -torch.sin(theta)
        G[j, i] = torch.sin(theta)
        G[j, j] = torch.cos(theta)

        v = torch.zeros(dims).to(theta.device)
        v[i] = torch.cos(theta)
        v[j] = torch.sin(theta)
        H = torch.eye(dims).to(theta.device) - 2 * torch.outer(v, v) / torch.dot(v, v)

        R = torch.eye(dims).to(theta.device)
        R[i, i] = torch.cos(theta)
        R[i, j] = -torch.sin(theta)
        R[j, i] = torch.sin(theta)
        R[j, j] = torch.cos(theta)

        return (G + H + R) / 3

    def apply_blended_rotation(self, x):
        adjusted_rot = int(torch.round(self.rot_scale * self.rot))
        for k in range(adjusted_rot):
            i, j = self.r_pairs[k].long()
            theta = self.thetas[k] * self.theta_scale
            B = self.blended_rotation_matrix(dims=self.h_dim, i=i, j=j, theta=theta)
            x = torch.matmul(input=x, other=B)
        return x

    def update_base(self, new_base):
        if new_base is not None and new_base != self.base:
            self.base = new_base
            inv_freq = 1.0 / (self.base ** (torch.arange(start=0, end=self.h_dim, step=2).float() / self.h_dim))
            self.inv_freq.data.copy_(inv_freq)
            self.update_pairs()

    def reset_parameters(self):
        nn.init.orthogonal_(self.r_matrix)
        nn.init.zeros_(self.thetas)
        nn.init.zeros_(self.r_pairs)
        nn.init.ones_(self.theta_scale)
        nn.init.ones_(self.rot_scale)

    def orthogonal_regularization_term(self):
        loss = torch.tensor(0.0, device=self.r_matrix.device)
        if self.r_matrix.requires_grad:
            product = torch.matmul(self.r_matrix, self.r_matrix.t())
            identity = torch.eye(self.r_matrix.size(0)).to(self.r_matrix.device)
            loss = ((product - identity) ** 2).sum()
        return self.orthogonal_reg_weight * loss

    def update_pairs(self):
        pairs = []
        while len(pairs) < self.rot:
            i, j = torch.randint(0, self.h_dim - 1, (2,))
            if i != j and (i, j) not in pairs and (j, i) not in pairs:
                pairs.append((i, j))
        self.r_pairs.data.copy_(torch.tensor(pairs, dtype=torch.float32))

    def forward(self, x, global_step=None):
        if x.dim() not in [3, 4]:
            raise ValueError(f"Expected input tensor to be 3D or 4D, but got {x.dim()}D")

        batch_size, seq_len, *rest = x.size()

        if x.dim() == 3:
            dims = rest[0]
            if dims != self.head * self.h_dim:
                raise ValueError(f"Expected dims ({dims}) to be compatible with head ({self.head}) * h_dim ({self.h_dim}={self.head * self.h_dim})")
        else:
            head, h_dim = rest
            if head != self.head or h_dim != self.h_dim:
                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}")

        x = x.view(batch_size, seq_len, self.head, self.h_dim)
        x = x.reshape(-1, self.h_dim)

        x = self.apply_blended_rotation(x)

        x = torch.matmul(input=x, other=self.r_matrix)

        x = x.view(batch_size, seq_len, self.head, self.h_dim)

        sinusoid_inp = torch.einsum('i, j -> i j', torch.arange(end=seq_len, device=x.device), self.inv_freq.to(device=x.device))
        sin = sinusoid_inp.sin()[None, :, None, :]
        cos = sinusoid_inp.cos()[None, :, None, :]

        x1, x2 = x[..., ::2], x[..., 1::2]
        x = torch.cat(tensors=[x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1)
        x = x.view(batch_size, seq_len, self.dims)

        return x

class SinusoidalEmbedding(nn.Module):
    def __init__(self, n_ctx, dims, checkpoint):
        super().__init__()
        self.n_ctx = n_ctx
        self.dims = dims
        self.checkpoint = checkpoint

        position = torch.arange(0, n_ctx, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, dims, 2).float() * -(math.log(10000.0) / dims))
        features = torch.zeros(n_ctx, dims)
        features[:, 0::2] = torch.sin(position * div_term)
        features[:, 1::2] = torch.cos(position * div_term)
        self.register_buffer('my_big_toe', features)
        self.pos_embeds = nn.Parameter(self.my_big_toe.clone())

    def forward(self, positions):
        if self.checkpoint:
            position_embeddings = checkpoint(lambda x: self.pos_embeds[x], positions)
        else:
            position_embeddings = self.pos_embeds[positions]
        return F.normalize(position_embeddings, p=2, dim=-1) 

class CombinedPositionalEmbedding(nn.Module):
    def __init__(self, base, dims, head, n_ctx, theta_learnable=True, rot_learnable=True, 
                 matrix_learnable=False, freq_learnable=True, checkpoint=False):
        super().__init__()
        self.rotary_embedding = CombinedRotaryEmbedding(base, dims, head, theta_learnable, 
                                                        rot_learnable, matrix_learnable, freq_learnable)
        self.sinusoidal_embedding = SinusoidalEmbedding(n_ctx, dims, checkpoint)

    def forward(self, x, positions, global_step=None):
        rotary_embed = self.rotary_embedding(x, global_step)
        sinusoidal_embed = self.sinusoidal_embedding(positions)
        
        combined_embedding = rotary_embed + sinusoidal_embed
        return combined_embedding

class MultiheadAttention(nn.Module):
    use_sdpa = True

    def __init__(self, base, dims, head, max_dist):
        super().__init__()
        assert dims % head == 0, "dims must be divisible by head"
        self.head = head
        self.h_dim = dims // head
        assert self.h_dim % 2 == 0, "Head dimension must be even for rotary embeddings"

        self.query = nn.Linear(dims, dims)
        self.key = nn.Linear(dims, dims, bias=False)
        self.value = nn.Linear(dims, dims)
        self.out = nn.Linear(dims, dims)

        # self.givens_rotary = CombinedRotaryEmbedding(base=base, dims=dims, head=head)

    def forward(self, x, xa = None, mask = None, kv_cache = None):

        q = self.query(x)

        if kv_cache is None or xa is None or self.key not in kv_cache:
            k = self.key(x if xa is None else xa)
            v = self.value(x if xa is None else xa)

        else:
            k = kv_cache[self.key]
            v = kv_cache[self.value]

        # q = self.givens_rotary(q)
        # k = self.givens_rotary(k)

        wv, qk = self.qkv_attention(q=q, k=k, v=v, mask=mask)

        out = self.out(wv)
        return out, qk
    
    def qkv_attention(self, q: Tensor, k: Tensor, v: Tensor, mask: Optional[Tensor] = None):
        
        n_batch, n_ctx, dims = q.shape
        scale = (dims // self.head) ** -0.25
        q = q.view(*q.shape[:2], self.head, -1).permute(0, 2, 1, 3)
        k = k.view(*k.shape[:2], self.head, -1).permute(0, 2, 1, 3)
        v = v.view(*v.shape[:2], self.head, -1).permute(0, 2, 1, 3)

        if MultiheadAttention.use_sdpa:
            a = scaled_dot_product_attention(query=q, key=k, value=v, is_causal=mask is not None and n_ctx > 1)
            out = a.permute(0, 2, 1, 3).flatten(start_dim=2)
            qk = None
        else:
            qk = (q * scale) @ (k * scale).transpose(-1, -2)
            if mask is not None:
                qk = qk + mask[:n_ctx, :n_ctx]
            qk = qk.float()

            w = F.softmax(qk, dim=-1).to(dtype=q.dtype)
            out = (w @ v).permute(0, 2, 1, 3).flatten(start_dim=2)
            qk = qk.detach()

        return out, qk
    
class AdaptiveSpanAttention(nn.Module):
    def __init__(self, base, dims, head, max_dist, sharpen, win_size, max_span, temp_scale=0.01):
        super().__init__()
        self.max_dist = max_dist
        self.win_size = win_size
        self.max_span = max_span
        self.temp_scale = temp_scale
        self.multihead_attn = MultiheadAttention(base=base, dims=dims, head=head, max_dist=max_dist)
        self.span_scale = nn.Parameter(torch.tensor(1.0))
        self.sharpen = sharpen

    def forward(self, query, key, value, span_scale):
        span_len = int(self.max_span * span_scale.mean().item())
        span_len = min(span_len, query.shape[1], key.shape[1], value.shape[1])
        eff_span = min(span_len, self.max_dist)

        q_span = query[:, :eff_span, :]
        k_span = key[:, :eff_span, :]
        v_span = value[:, :eff_span, :]

        batch_size, _, dims = query.shape
        scale = (dims // self.multihead_attn.head) ** -0.25

        q = q_span.view(q_span.shape[0], q_span.shape[1], self.multihead_attn.head, -1).permute(0, 2, 1, 3)
        k = k_span.view(k_span.shape[0], k_span.shape[1], self.multihead_attn.head, -1).permute(0, 2, 1, 3)
        v = v_span.view(v_span.shape[0], v_span.shape[1], self.multihead_attn.head, -1).permute(0, 2, 1, 3)

        if self.sharpen:
            temperature = 1.0 + self.temp_scale * (1.0 - span_scale.mean().item())
        else:
            temperature = 0.5 + self.temp_scale * span_scale.mean().item()

        attn_scores = torch.matmul(q, k.transpose(-2, -1))
        attn_weights = torch.softmax((attn_scores / temperature) * scale, dim=-1)
        attn_out = torch.matmul(attn_weights, v)
        attn_out = attn_out.permute(0, 2, 1, 3).flatten(start_dim=2)
        attn_out = attn_out.contiguous().view(batch_size, eff_span, dims)

        return attn_out, attn_weights

class SpanPredictor(nn.Module):
    def __init__(self, dims):
        super().__init__()
        self.linear = nn.Linear(in_features=dims, out_features=1)

    def forward(self, global_out):
        scale = torch.sigmoid(self.linear(global_out))
        return scale

class HybridAttention(nn.Module):
    def __init__(self, base, dims, head, max_dist, sharpen, win_size=32, max_span=32, slid_win=32):
        super().__init__()
        self.max_dist = max_dist
        self.win_size = win_size
        self.max_span = max_span
        self.slid_win = slid_win

        self.span_pred = SpanPredictor(dims=dims)
        self.dist_local = max_dist
        self.dist_global = max_dist

        self.attn_local = AdaptiveSpanAttention(base=base, dims=dims, head=head, max_dist=max_dist, sharpen=sharpen, win_size=win_size, max_span=max_span)
        self.attn_global = MultiheadAttention(base=base, dims=dims, head=head, max_dist=self.dist_global)
        self.ln_local = LayerNorm(normalized_shape=dims)
        self.ln_global = LayerNorm(normalized_shape=dims)
        self.projection = Linear(in_features=2 * dims, out_features=dims)

    def forward(self, x, new_dist=None, new_base=None, xa=None, mask=None, kv_cache=None):
        local = self.ln_local(x)
        globe = self.ln_global(x)

        globe_out, _ = self.attn_global(globe, globe, globe)

        span_scale = self.span_pred(globe_out.mean(dim=1))

        win_size = max(1, int(self.slid_win * span_scale.mean().item()))
        span_len = max(1, int(self.max_span * span_scale.mean().item()))

        effective_max_dist = min(self.max_dist, local.size(1))
        local_max_dist = min(self.dist_local, span_len, win_size)
        globe_max_dist = effective_max_dist

        self.attn_local.max_dist = local_max_dist
        self.attn_global.max_dist = globe_max_dist

        local_out = self.slide_win(x=local, win_size=win_size, span_len=span_len, span_scale=span_scale)

        combined = torch.cat(tensors=[local_out, globe_out], dim=-1)
        x = self.projection(combined)

        return x

    def slide_win(self, x, win_size, span_len, span_scale):
        batch_size, seq_len, dims = x.size()
        out = torch.zeros_like(x, device=x.device)

        for i in range(0, seq_len, win_size):
            end = min(i + win_size, seq_len)
            query = x[:, i:end, :]

            start = max(0, i - span_len + win_size)
            key = x[:, start:i + span_len, :]
            value = x[:, start:i + span_len, :]
            attn_out, _ = self.attn_local(query, key, value, span_scale)
            out[:, i:end, :] = attn_out

        return out

class ResidualAttention(nn.Module):
    def __init__(self, base, dims, head, max_dist, win_size, max_span, hybrid, checkpoint, cross, sharpen):
        super().__init__()

        if hybrid:
            self.attn = HybridAttention(base=base, dims=dims, head=head, max_dist=max_dist, sharpen=sharpen)
            self.attn_ln = LayerNorm(normalized_shape=dims)
        else:
            self.attn = MultiheadAttention(base=base, dims=dims, head=head, max_dist=max_dist)
            self.attn_ln = LayerNorm(normalized_shape=dims)

        n_mlp = dims * 4
        self.mlp = nn.Sequential(Linear(in_features=dims, out_features=n_mlp), nn.GELU(), Linear(in_features=n_mlp, out_features=dims))
        self.mlp_ln = LayerNorm(normalized_shape=dims)

    def forward(self, x, mask=None, kv_cache=None):
        x = self._attn_forward(x=x, mask=mask, kv_cache=kv_cache)
        x = self._mlp_forward(x=x)
        return x

    def _attn_forward(self, x, mask=None, kv_cache=None):
        residual = x
        x = self.attn_ln(x)

        if isinstance(self.attn, HybridAttention):
            attn_output = self.attn(x)  

            x = residual + attn_output
        else:
            attn_output, _ = self.attn(x, mask=mask, kv_cache=kv_cache)  
            x = residual + attn_output
        return x

    def _mlp_forward(self, x):
        residual = x
        x = self.mlp_ln(x)
        return residual + self.mlp(x)

class AudioEncoder(nn.Module):
    def __init__(self, base, mels, dims, head, n_layer, n_ctx, max_dist,
                 win_size, max_span, hybrid, checkpoint, cross, sharpen):
        super().__init__()
        self.conv1 = Conv1d(in_channels=mels, out_channels=dims, kernel_size=3, padding=1)
        self.conv2 = Conv1d(in_channels=dims, out_channels=dims, kernel_size=3, stride=2, padding=1)
        self.pos_embed = SinusoidalEmbedding(n_ctx=n_ctx, dims=dims, checkpoint=checkpoint)
        self.checkpoint = checkpoint

        self.givens_rotary = CombinedRotaryEmbedding(base=base, dims=dims, head=head)
        # self.combine = CombinedPositionalEmbedding(base=base, dims=dims, head=head)
        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)])
        self.ln_post = LayerNorm(normalized_shape=dims)

    def forward(self, x):
        if self.checkpoint:
            x = checkpoint(self._conv_forward, x)
        else:
            x = self._conv_forward(x)

        for block in self.blocks:
            if self.checkpoint:
                x = checkpoint(block, x)
            else:
                x = block(x)
        return self.ln_post(x)

    def _conv_forward(self, x):
        x = F.gelu(self.conv1(x))
        x = F.gelu(self.conv2(x))
        x = x.permute(0, 2, 1)
    
        p = self.pos_embed(torch.arange(end=x.size(dim=1), device=x.device)).unsqueeze(0)
        x = (x + p).to(x.dtype)
        x = self.givens_rotary(x)
        # x = self.combine(x)
        return x

class TextDecoder(nn.Module):
    def __init__(self, base, vocab, dims, head, n_layer, n_ctx, max_dist,
                 win_size, max_span, hybrid, checkpoint, cross, sharpen):
        super().__init__()
        
        self.tok_embed = nn.Embedding(num_embeddings=vocab, embedding_dim=dims)
        self.pos_embed = SinusoidalEmbedding(n_ctx=n_ctx, dims=dims, checkpoint=checkpoint)
        self.checkpoint = checkpoint

        self.givens_rotary = CombinedRotaryEmbedding(base=base, dims=dims, head=head)

        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)])

        self.ln_post = LayerNorm(normalized_shape=dims)
        self.ln = LayerNorm(normalized_shape=dims)

        mask = torch.empty(n_ctx, n_ctx).fill_(value=-np.inf).triu_(diagonal=1)
        self.register_buffer(name="mask", tensor=mask, persistent=False)
        self.mask=mask

    def forward(self, x, xa, kv_cache=None):
        if self.checkpoint:
            x = checkpoint(self._embedding_forward, x, xa, kv_cache)
        else:
            x = self._embedding_forward(x=x, xa=xa, kv_cache=kv_cache)

        for block in self.blocks:
            if self.checkpoint:
                x = checkpoint(block, x, self.mask, kv_cache)
            else:
                x = block(x, self.mask, kv_cache)

        x = self.ln(x)
        x = (x @ torch.transpose(input=self.tok_embed.weight.to(dtype=x.dtype), dim0=0, dim1=1)).float()
        return x
    
    def _embedding_forward(self, x, xa, kv_cache):
        offset = next(iter(kv_cache.values())).shape[1] if kv_cache else 0
        positions = torch.arange(x.shape[1], device=x.device) + offset
        pos_emb = self.pos_embed(positions).unsqueeze(0)
        x = self.tok_embed(x) + pos_emb
        x = self.givens_rotary(x)
        return x

class EchoConfig(PretrainedConfig):
    model_type = "Echo"
    def __init__(
        self,
        checkpoint=False,
        cross=False,
        hybrid=False,
        sharpen=False,
        a_ctx=1500,
        a_head=16,
        a_layer=8,
        a_dims=1024,
        mels=128,
        t_ctx=448,
        t_head=8,
        t_layer=8,
        t_dims=1024,
        win_size=64,
        max_span=64,
        max_dist=64,
        base=10000,
        pad_token_id=50257,
        unk_token_id=50257,
        vocab=51865,
        eos_token_id=50257,
        bos_token_id=50257,
        decoder_start_token_id=50258,
        **kwargs,
    ):
        
        super().__init__(**kwargs) 
        self.base = base
        self.bos_token_id = bos_token_id
        self.checkpoint = checkpoint
        self.cross = cross
        self.decoder_start_token_id = decoder_start_token_id
        self.eos_token_id = eos_token_id
        self.hybrid = hybrid
        self.max_dist = max_dist
        self.max_span = max_span
        self.a_ctx = a_ctx
        self.a_head = a_head
        self.a_layer = a_layer
        self.a_dims = a_dims
        self.mels = mels
        self.t_ctx = t_ctx
        self.t_head = t_head
        self.t_layer = t_layer
        self.t_dims = t_dims
        self.pad_token_id = pad_token_id
        self.unk_token_id = unk_token_id
        self.vocab = vocab
        self.win_size = win_size
        self.sharpen=sharpen

class Echo(nn.Module):
    def __init__(self, config: EchoConfig):
        super().__init__()
        self.config = config
            
        self.encoder = AudioEncoder(
            base=self.config.base,
            mels=self.config.mels,
            dims=self.config.a_dims, 
            head=self.config.a_head,
            n_layer=self.config.a_layer,
            n_ctx=self.config.a_ctx,
            max_dist=self.config.max_dist,
            win_size=self.config.win_size,  
            max_span=self.config.max_span,
            hybrid=self.config.hybrid,
            checkpoint=self.config.checkpoint,
            cross=self.config.cross,
            sharpen=self.config.sharpen,
        )

        self.decoder = TextDecoder(
            base=self.config.base,
            vocab=self.config.vocab,
            dims=self.config.t_dims, 
            head=self.config.t_head,
            n_layer=self.config.t_layer,
            n_ctx=self.config.t_ctx,
            max_dist=self.config.max_dist,
            win_size=self.config.win_size,  
            max_span=self.config.max_span,
            hybrid=self.config.hybrid,
            checkpoint=self.config.checkpoint,
            cross=self.config.cross,
            sharpen=self.config.sharpen,
        )

        all_heads = torch.zeros(self.config.t_layer, self.config.t_head, dtype=torch.bool) 
        all_heads[self.config.t_layer // 2:] = True
        self.register_buffer(name="alignment_heads", tensor=all_heads.to_sparse(), persistent=False)

        self.base = self.config.base
        self.win_size = self.config.win_size
        self.adjust_counter = 0
        self.best_loss = float('inf')
        self.kv_cache = {}

    @property
    def device(self):
        return next(self.parameters()).device

    def embed_audio(self, mel: torch.Tensor):
        return self.encoder(mel)

    def logits(self, tokens: torch.Tensor, audio_features: torch.Tensor):
        return self.decoder(tokens, audio_features)

    def update_window(self, new_window):
        self.win_size = new_window
        for module in self.modules(): 
            if isinstance(module, HybridAttention):
                module.update_window(self.win_size)

    def adjust_window(self, loss, factor=1.00005):
        if self.adjust_counter % 10 == 0:
            if loss < self.best_loss:
                new_window = self.win_size * factor
            else:
                new_window = self.win_size / factor
            self.update_window(new_window=new_window)
            self.best_loss = loss
            self.adjust_counter += 1
            return new_window
        return self.win_size

    def adjust_base(self, loss, factor=1.0025) -> float | int:
                if self.adjust_counter % 25 == 0:
                    if loss < self.best_loss:
                        new_base=self.base*factor
                    else:
                        new_base=self.base/factor
                    self.update_base(new_base=new_base)
                    self.base=new_base
                    self.best_loss=loss
                self.adjust_counter += 1
                return self.base
            
    def update_base(self, new_base):
        self.new_base=new_base
        for name, module in self.encoder.named_modules():
            if isinstance(module, (CombinedRotaryEmbedding)):
                module.update_base(new_base=self.new_base)

    @staticmethod
    def shift_tokens_right(input_ids, pad_token_id, decoder_start_token_id):
        shifted_input_ids = input_ids.new_zeros(input_ids.shape)
        shifted_input_ids[:, 1:] = input_ids[:, :-1].clone() 
        shifted_input_ids[:, 0] = decoder_start_token_id
        shifted_input_ids.masked_fill_(shifted_input_ids == -100, pad_token_id)
        return shifted_input_ids

    def forward(self, input_features, labels=None, dec_input_ids=None) -> dict[str, Any | None]:
        if labels is not None:
            if dec_input_ids is None:
                dec_input_ids = self.shift_tokens_right(
                    input_ids=labels, pad_token_id=self.config.pad_token_id, decoder_start_token_id=self.config.decoder_start_token_id
                )

        encoded_features = self.encoder(input_features).to(self.device)  
        logits = self.decoder(dec_input_ids, encoded_features)

        loss = None
        if labels is not None:
            loss_fct = nn.CrossEntropyLoss(ignore_index=-100)
            labels = labels.to(logits.device).long()
            loss = loss_fct(logits.view(-1, self.config.vocab), labels.view(-1))

            self.adjust_window(loss.item())
            # self.adjust_base(loss=loss.item())
        return {"loss": loss, "logits": logits}

    def reset_parameters(self):
        for name, module in self.encoder.named_modules():
            if isinstance(module, CombinedRotaryEmbedding):
                module.reset_parameters()
        
    def _initialize_weights(self, module):
            nn.init.normal_(tensor=self.decoder.tok_embed.weight, mean=0.0, std=0.02)
            nn.init.constant_(tensor=self.decoder.ln.weight, val=1)
            nn.init.constant_(tensor=self.decoder.ln.bias, val=0)
            nn.init.xavier_normal_(tensor=self.encoder.conv1.weight)
            nn.init.zeros_(tensor=self.encoder.conv1.bias)
            nn.init.kaiming_normal_(tensor=self.encoder.conv2.weight, mode='fan_out', nonlinearity='relu')
            nn.init.zeros_(tensor=self.encoder.conv2.bias)
            nn.init.constant_(tensor=self.encoder.ln_post.weight, val=1)
            nn.init.constant_(tensor=self.encoder.ln_post.bias, val=0)

            for block in self.decoder.blocks:
                for layer in block.children():
                    if isinstance(layer, nn.Linear):
                        nn.init.xavier_normal_(tensor=layer.weight)
                        nn.init.zeros_(tensor=layer.bias)
                    if isinstance(layer, LayerNorm):
                        nn.init.constant_(tensor=layer.weight, val=1)
            
            for block in self.encoder.blocks:
                for layer in block.children():
                    if isinstance(layer, nn.Linear):
                        nn.init.xavier_normal_(tensor=layer.weight)
                        nn.init.zeros_(tensor=layer.bias)
                    if isinstance(layer, LayerNorm):
                        nn.init.constant_(tensor=layer.weight, val=1)

            for module in self.encoder.named_modules():
                if isinstance(module, CombinedRotaryEmbedding):
                    nn.init.constant_(tensor=module.thetas, val=1)
                    nn.init.constant_(tensor=module.r_matrix, val=1)
                    nn.init.constant_(tensor=module.r_pairs, val=1)
                    nn.init.constant_(tensor=module.inv_freq, val=1)

    def apply_initialization(self, module):
        self._initialize_weights(module=module)

from datetime import datetime
log_dir = os.path.join('./output/Echo/', datetime.now().strftime(format='%m-%d_%H'))
os.makedirs(name=log_dir, exist_ok=True)

config = EchoConfig(
    checkpoint=False,
    cross=False,
    hybrid=False,
    sharpen=False,
    audio_ctx=1500,
    audio_head=4,
    audio_layer=4,
    audio_dims=512,
    mels=128,
    text_ctx=448,
    text_head=4,
    text_layer=4,
    text_dims=512,
    win_size=16,
    max_span=16,
    max_dist=16,
    base=50000,
    pad_token_id=50257,
    unk_token_id=50257,
    vocab=51865,
    eos_token_id=50257,
    bos_token_id=50257,
    decoder_start_token_id=50258,

)

model = Echo(config=config).to(device=device)
model.apply_initialization(module=model)


feature_extractor = WhisperFeatureExtractor.from_pretrained(
    pretrained_model_name_or_path="openai/whisper-small", 
    feature_size=128, sample_rate=160000, do_normalize=True)

tokenizer = WhisperTokenizerFast.from_pretrained(
    pretrained_model_name_or_path="openai/whisper-small", 
    language="en", task="transcribe")

processor = WhisperProcessor.from_pretrained(
    pretrained_model_name_or_path="openai/whisper-small", 
    feature_size=128, sample_rate=160000, do_normalize=True, 
    language="en", task="transcribe")

class GradientClippingCallback(TrainerCallback):
    def on_step_end(self, args, dims, control, **kwargs):
        torch.nn.utils.clip_grad_norm_(parameters=kwargs["model"].parameters(), max_norm=0.98)

@dataclass
class DataCollatorSpeechSeq2SeqWithPadding:
    processor: Any
    decoder_start_token_id: int

    def __call__(self, features: List[Dict[str, Union[List[int], torch.Tensor]]]) -> Dict[str, torch.Tensor]:
        input_features = [{"input_features": feature["input_features"]} for feature in features]
        batch = self.processor.feature_extractor.pad(input_features, return_tensors="pt")
        label_features = [{"input_ids": feature["labels"]} for feature in features]
        labels_batch = self.processor.tokenizer.pad(label_features, return_tensors="pt")
        labels = labels_batch["input_ids"].masked_fill(labels_batch.attention_mask.ne(1), -100)
        if (labels[:, 0] == self.decoder_start_token_id).all().cpu().item():
            labels = labels[:, 1:]
        batch["labels"] = labels
        return batch

def get_length_of_dataset(dataset):
    length = 0
    for item in dataset:
        length += len(item["audio"]["array"]) / item["audio"]["sampling_rate"]
    return length / 3600  

def prepare_dataset(batch):
    audio = batch["audio"]
    batch["input_features"] = feature_extractor(audio["array"], sampling_rate=audio["sampling_rate"]).input_features[0]
    batch["labels"] = tokenizer(batch["sentence"]).input_ids
    return batch

data_collator = DataCollatorSpeechSeq2SeqWithPadding(processor=processor, decoder_start_token_id=config.decoder_start_token_id)

datasets = IterableDatasetDict()

datasets["train"] = load_dataset(
    path="mozilla-foundation/common_voice_17_0", token="",
    name="en", split="train", streaming=True, trust_remote_code=True).take(10000)

datasets["test"] = load_dataset(
    path="mozilla-foundation/common_voice_17_0", token="", 
    name="en", split="test", streaming=True, trust_remote_code=True).take(100)

dataset = datasets.cast_column(column="audio", feature=Audio(sampling_rate=16000))

dataset = dataset.map(function=prepare_dataset, 
                      remove_columns=list(next(iter(dataset.values())).features)).with_format(type="torch")

class MetricsCallback(TrainerCallback):
    def __init__(self, tb_writer, tokenizer, metric, optimizer, scheduler, log_every_n_steps=1):
        super().__init__()
        self.tb_writer = tb_writer
        self.tokenizer = tokenizer
        self.metric = metric
        self.optimizer = optimizer
        self.scheduler = scheduler
        self.log_every_n_steps = log_every_n_steps
        self.predictions = None
        self.label_ids = None

    def compute_wer(self, pred_str, label_str):
        wer = 100 * self.metric.compute(predictions=pred_str, references=label_str)
        return wer

    def on_evaluate(self, args, state, control, model, metrics=None, **kwargs):
        if metrics is not None:
            self.eval_loss = metrics.get('eval_loss')

            current_learning_rate = self.optimizer.param_groups[0]['lr']
            if state.global_step % self.log_every_n_steps == 0:
                self.tb_writer.add_scalar('learning_rate', current_learning_rate, state.global_step)
                print(f"Learning Rate: {current_learning_rate:.8f}")

                self.tb_writer.add_scalar('eval_loss', self.eval_loss, state.global_step)

            for key, value in metrics.items():
                if key.startswith("eval_"):
                    self.tb_writer.add_scalar(key, value, state.global_step)

        if self.predictions is not None and self.label_ids is not None:
            pred_str = self.tokenizer.batch_decode(self.predictions, skip_special_tokens=True)
            label_str = self.tokenizer.batch_decode(self.label_ids, skip_special_tokens=True)

            if state.global_step % self.log_every_n_steps == 0:
                total_samples = len(pred_str)
                random_indices = random.sample(range(total_samples), 1)

                for sample_index in random_indices:
                    self.tb_writer.add_text(f"Prediction_{sample_index}", pred_str[sample_index], state.global_step)
                    self.tb_writer.add_text(f"Label_{sample_index}", label_str[sample_index], state.global_step)
                    print(f"Evaluation: - Step {state.global_step} - Loss: {self.eval_loss:.2f}")
                    print(f"Prediction: {pred_str[sample_index]}")
                    print(f"Label: {label_str[sample_index]}")
                    print("-" * 10)

        self.predictions = None
        self.label_ids = None

def create_compute_metrics(callback_instance):
    def compute_metrics(eval_pred):
        pred_logits = eval_pred.predictions
        label_ids = eval_pred.label_ids

        if isinstance(pred_logits, tuple):
            pred_ids = pred_logits[0]
        else:
            pred_ids = pred_logits
        if pred_ids.ndim == 3:
            pred_ids = np.argmax(pred_ids, axis=-1)

        label_ids[label_ids == -100] = callback_instance.tokenizer.pad_token_id
        callback_instance.predictions = pred_ids
        callback_instance.label_ids = label_ids
        pred_str = callback_instance.tokenizer.batch_decode(pred_ids, skip_special_tokens=True)
        label_str = callback_instance.tokenizer.batch_decode(label_ids, skip_special_tokens=True)
        wer = 100 * callback_instance.metric.compute(predictions=pred_str, references=label_str)
        pred_flat = pred_ids.flatten()
        labels_flat = label_ids.flatten()
        mask = labels_flat != callback_instance.tokenizer.pad_token_id

        accuracy = accuracy_score(y_true=labels_flat[mask], y_pred=pred_flat[mask])
        precision = precision_score(y_true=labels_flat[mask], y_pred=pred_flat[mask], average='weighted', zero_division=0)
        recall = recall_score(y_true=labels_flat[mask], y_pred=pred_flat[mask], average='weighted', zero_division=0)
        f1 = f1_score(y_true=labels_flat[mask], y_pred=pred_flat[mask], average='weighted', zero_division=0)
        return {"wer": wer, "accuracy": accuracy, "precision": precision, "recall": recall, "f1": f1}
    return compute_metrics

metric = evaluate.load(path="wer")
tb_writer = SummaryWriter(log_dir=log_dir)

training_args = Seq2SeqTrainingArguments(
    output_dir=log_dir,
    per_device_train_batch_size=1,
    per_device_eval_batch_size=1,
    gradient_accumulation_steps=1,
    eval_accumulation_steps=1,
    tf32=True,
    bf16=True,
    eval_strategy="steps",
    save_strategy="steps",
    max_steps=10000,
    save_steps=10000,
    eval_steps=100,
    warmup_steps=100,
    logging_steps=10,
    logging_dir=log_dir + "/logs_hf",
    report_to=["tensorboard"],
    load_best_model_at_end=False,
    metric_for_best_model="loss",
    greater_is_better=False,
    push_to_hub=False,
    disable_tqdm=False,
    save_total_limit=1,
    remove_unused_columns=False,
    label_names=["labels"],
    eval_on_start=True,
)

class MaxFactor(Optimizer):
    def __init__(self, params, lr=0.01, beta2_decay=-0.8, eps=(None, 1e-3), d=1.0, 
                 weight_decay=0.0, gamma=0.99, eps_rms=1e-8, maximize=False):
        
        defaults = dict(lr=lr, beta2_decay=beta2_decay, eps=eps, d=d, weight_decay=weight_decay, 
                        gamma=gamma, eps_rms=eps_rms, maximize=maximize)

        super().__init__(params, defaults)

    @torch.no_grad()
    def step(self, closure=None):
        loss = None
        if closure is not None:
            with torch.enable_grad():
                loss = closure()

        for group in self.param_groups:
            params_with_grad, grads, row_vars, col_vars, v, state_steps = [], [], [], [], [], []
            eps1, eps2 = group["eps"]
            for p in group["params"]:
                if p.grad is None:
                    continue
                grad = p.grad
                if grad.dtype in {torch.float16, torch.bfloat16}:
                    grad = grad.float()

                state = self.state[p]
                if len(state) == 0:
                    state["step"] = torch.tensor(0.0, dtype=torch.float32)
                    if p.grad.dim() > 1:
                        row_shape, col_shape = list(p.grad.shape), list(p.grad.shape)
                        row_shape[-1], col_shape[-2] = 1, 1
                        state["row_var"], state["col_var"] = p.grad.new_zeros(row_shape), p.grad.new_zeros(col_shape)
                    state["v"] = torch.zeros_like(p, memory_format=torch.preserve_format)

                row_vars.append(state.get("row_var", None))
                col_vars.append(state.get("col_var", None))
                v.append(state["v"])
                state_steps.append(state["step"])
                params_with_grad.append(p)
                grads.append(grad)

            for i, param in enumerate(params_with_grad):
                grad = grads[i]

                if group["maximize"]:
                    grad = -grad
                step_t, row_var, col_var, vi = state_steps[i], row_vars[i], col_vars[i], v[i]

                if eps1 is None:
                    eps1 = torch.finfo(param.dtype).eps
                    
                step_t += 1
                step_float = step_t.item()
                one_minus_beta2_t = step_float ** group["beta2_decay"]
                rho_t = min(group["lr"], 1 / (step_float ** 0.5))
                alpha = max(eps2, param.norm(2).item() / (param.numel() ** 0.5)) * rho_t

                if group["weight_decay"]!= 0:
                    param.mul_(1 - group["lr"] * group["weight_decay"])

                if grad.dim() > 1:
                    row_mean = torch.norm(grad, dim=-1, keepdim=True).square_().div_(grad.size(-1))
                    row_var.lerp_(row_mean, one_minus_beta2_t)
                    col_mean = torch.norm(grad, dim=-2, keepdim=True).square_().div_(grad.size(-2))
                    col_var.lerp_(col_mean, one_minus_beta2_t)
                    var_estimate = row_var @ col_var
                    max_row_var = row_var.max(dim=-2, keepdim=True)[0]  
                    var_estimate.div_(max_row_var.clamp_(min=eps1))

                else:
                    vi.mul_(group["gamma"]).add_(1 - group["gamma"], grad ** 2)
                    var_estimate = vi
                
                update = var_estimate.clamp_(min=eps1 * eps1).rsqrt_().mul_(grad)
                update = update.div_(torch.norm(update, float('inf')).clamp_(min=eps1))
                denom = max(1.0, update.norm(2).item() / ((update.numel() ** 0.5) * group["d"]))
                param.add_(-alpha / denom * update.sign() * update.abs().max(dim=-1, keepdim=True)[0])

        return loss
    
optimizer = MaxFactor(
    model.parameters(), 
    lr=0.025,  
    beta2_decay=-0.8,
    eps=(None, 1e-4),
    d=1.0,
    weight_decay=0.0025,
    gamma=0.99, 
    eps_rms=1e-8,
    maximize=False,
    )

scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
    optimizer=optimizer,
    T_max=training_args.max_steps,
    eta_min=0.0,
    last_epoch=-1  
)

metrics_callback = MetricsCallback(tb_writer=tb_writer, tokenizer=tokenizer, metric=metric, optimizer=optimizer, scheduler=scheduler, log_every_n_steps=10)
compute_metrics = create_compute_metrics(callback_instance=metrics_callback)

trainer = Seq2SeqTrainer(
    args=training_args,
    model=model,
    train_dataset=dataset["train"],
    eval_dataset=dataset["test"],
    data_collator=data_collator,
    compute_metrics=compute_metrics,
    processing_class=feature_extractor,
    callbacks=[metrics_callback],
    optimizers=(optimizer, scheduler)
)

trainer.train(resume_from_checkpoint=False)

from tensorboard import program
log_dir = "D:/new/tensorboard3" 
tb = program.TensorBoard()
tb.configure(argv=[None, '--logdir', log_dir])
url = tb.launch()
print(f"TensorBoard started at {url}")