echoutils.py

import torch
import os
import pyworld as pw
import numpy as np
import torchaudio
import torch.nn.functional as F
from datasets import load_dataset, Audio
from dataclasses import dataclass
from typing import Any, List, Dict
import math
import matplotlib.pyplot as plt
import torch.nn as nn
import torch.nn.init as init
from torch import Tensor
from typing import Optional, Union, Tuple
from torch.nn.functional import scaled_dot_product_attention
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
dtype = torch.float32

class LayerNorm(nn.Module):
    def __init__(self, emb_dim):
        super().__init__()
        self.eps = 1e-5
        self.scale = nn.Parameter(torch.ones(emb_dim))
        self.shift = nn.Parameter(torch.zeros(emb_dim))

    def forward(self, x):
        mean = x.mean(dim=-1, keepdim=True)
        var = x.var(dim=-1, keepdim=True, unbiased=False)
        norm_x = (x - mean) / torch.sqrt(var + self.eps)
        return self.scale * norm_x + self.shift

class GELU(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x):
        return 0.5 * x * (1 + torch.tanh(
            torch.sqrt(torch.tensor(2.0 / torch.pi)) * 
            (x + 0.044715 * torch.pow(x, 3))
        ))

def sinusoids(ctx, dims, max_tscale=10000):
    assert dims % 2 == 0
    pos = torch.log(torch.tensor(float(max_tscale))) / (dims // 2 - 1)
    tscales = torch.exp(-pos * torch.arange(dims // 2, device=device, dtype=torch.float32))
    scaled = torch.arange(ctx, device=device, dtype=torch.float32).unsqueeze(1) * tscales.unsqueeze(0)
    position = torch.cat([torch.sin(scaled), torch.cos(scaled)], dim=1) 
    positional_embedding = nn.Parameter(position, requires_grad=True)
    return positional_embedding

def get_activation(act: str) -> nn.Module:
    act_map = {
        "gelu": nn.GELU(), 
        "relu": nn.ReLU(), 
        "sigmoid": nn.Sigmoid(), 
        "tanh": nn.Tanh(), 
        "swish": nn.SiLU(), 
        "tanhshrink": nn.Tanhshrink(), 
        "softplus": nn.Softplus(), 
        "softshrink": nn.Softshrink(), 
        "leaky_relu": nn.LeakyReLU(), 
        "elu": nn.ELU()
    }
    return act_map.get(act, nn.GELU())

def cos_sim(q: Tensor, k: Tensor, v: Tensor, mask) -> Tensor:
    q_norm = torch.nn.functional.normalize(q, dim=-1, eps=1e-12)
    k_norm = torch.nn.functional.normalize(k, dim=-1, eps=1e-12)
    qk_cosine = torch.matmul(q_norm, k_norm.transpose(-1, -2))
    qk_cosine = qk_cosine + mask
    weights = F.softmax(qk_cosine, dim=-1)
    out = torch.matmul(weights, v)
    return out

def taylor_softmax_2nd_order(x):
    exp_approx = 1 + x + (x**2) / 2
    return exp_approx / torch.sum(exp_approx, dim=-1, keepdim=True)

def taylor_softmax_approximation(x, order=2):
    if order == 0:
        return torch.ones_like(x) / x.size(-1) 
    elif order == 1:
        numerator = 1 + x
    elif order == 2:
        numerator = 1 + x + 0.5 * x**2
    else:
        raise NotImplementedError("Higher orders are not implemented yet.")
    denominator = torch.sum(numerator, dim=-1, keepdim=True)
    return numerator / denominator

def rbf_scores(q, k, rbf_sigma=1.0, rbf_ratio=0.0):
    dot_scores = torch.matmul(q, k.transpose(-1, -2))
    if rbf_ratio <= 0.0:
        return dot_scores
    q_norm = q.pow(2).sum(dim=-1, keepdim=True)
    k_norm = k.pow(2).sum(dim=-1, keepdim=True)
    qk = torch.matmul(q, k.transpose(-1, -2))
    dist_sq = q_norm + k_norm.transpose(-1, -2) - 2 * qk
    rbf_scores = torch.exp(-dist_sq / (2 * rbf_sigma**2))
    return (1 - rbf_ratio) * dot_scores + rbf_ratio * rbf_scores

def sliding_window_mask(q_len, k_len, window, device):
    # mask[i, j] = 1 if j in [i-window+1, i], else 0
    idxs = torch.arange(q_len, device=device).unsqueeze(1)
    jdxs = torch.arange(k_len, device=device).unsqueeze(0)
    mask = (jdxs >= (idxs - window + 1)) & (jdxs <= idxs)
    return mask.float()  # shape: (q_len, k_len)

def mask_win(text_ctx, aud_ctx):
    mask = torch.tril(torch.ones(text_ctx, text_ctx, device=device, dtype=dtype), diagonal=0)
    audio_mask = torch.tril(torch.ones(text_ctx, aud_ctx - text_ctx, device=device, dtype=dtype))
    full_mask = torch.cat([mask, audio_mask], dim=-1)
    return full_mask

def maskc(ctx, device):
    return torch.tril(torch.ones(ctx, ctx, device=device, dtype=dtype), diagonal=0)

def create_attention_mask(batch_size, ctx, is_causal=True, padding_mask=None, device=None):
    if is_causal:
        mask = torch.triu(torch.ones((ctx, ctx), device=device), diagonal=0)
        mask = mask.expand(batch_size, 1, ctx, ctx)
    else:
        mask = torch.zeros((batch_size, 1, ctx, ctx), device=device)
    if padding_mask is not None:
        padding_mask = padding_mask.unsqueeze(1).unsqueeze(2).bool()
        mask = (mask.bool() | (~padding_mask)).float()
    return mask

def calculate_attention(q, k, v, mask=None, temp=1.0):
    scaled_q = q
    if temp != 1.0 and temp > 0:
        scaled_q = q * (1.0 / temp)**.5
    out = scaled_dot_product_attention(scaled_q, k, v, is_causal=mask is not None and q.shape[1] > 1)        
    return out

def calculate_attentionb(q_norm, k_norm, v_iter, mask=None, temp=1.0):
    d_k = q_norm.size(-1)
    scores = torch.matmul(q_norm, k_norm.transpose(-2, -1)) / (torch.sqrt(torch.tensor(d_k, dtype=torch.float32)) / temp)
    if mask is not None:
        scores = scores.masked_fill(mask == 0, float('-inf'))
    attention_weights = F.softmax(scores, dim=-1)
    output = torch.matmul(attention_weights, v_iter)
    return output

class LocalOut(nn.Module):
    def __init__(self, dims: int, head: int):
        super().__init__()
        self.head_dim = dims // head
        self.dims = dims
        self.q_module = nn.Linear(self.head_dim, self.head_dim)
        self.k_module = nn.Linear(self.head_dim, self.head_dim)
        self.v_module = nn.Linear(self.head_dim, self.head_dim)
        self.o_proj = nn.Linear(self.head_dim, self.head_dim)

    def _reshape_to_output(self, attn_output: Tensor) -> Tensor:
        batch, _, ctx, _ = attn_output.shape
        return attn_output.transpose(1, 2).contiguous().view(batch, ctx, self.dims)      

def qkv_init(dims, head):
    head_dim = dims // head
    q = nn.Linear(dims, dims)
    k = nn.Linear(dims, dims)
    v = nn.Linear(dims, dims)
    o = nn.Linear(dims, dims)
    lna = nn.LayerNorm(dims)
    lnb = nn.LayerNorm(dims)
    lnc = nn.LayerNorm(head_dim)
    lnd = nn.LayerNorm(head_dim)
    return q, k, v, o, lna, lnb, lnc, lnd

def shape(dims, head, q, k, v):
    batch_size = q.shape[0]
    seq_len_q = q.shape[1]
    seq_len_kv = k.shape[1]
    head_dim = dims // head

    q = q.view(batch_size, seq_len_q, head, head_dim).transpose(1, 2)
    k = k.view(batch_size, seq_len_kv, head, head_dim).transpose(1, 2)
    v = v.view(batch_size, seq_len_kv, head, head_dim).transpose(1, 2)
    return q, k, v

def create_qkv(dims, head, q, k, v, x, xa):
    head_dim = dims // head
    scale = head_dim ** -0.25
    q = q(x) * scale
    k = k(xa) * scale
    v = v(xa)
    batch, ctx, dims = x.shape
    def _shape(tensor):
        return tensor.view(batch, ctx, head, head_dim).transpose(1, 2).contiguous()
    return _shape(q), _shape(k), _shape(v)

# def calculate_attention(q, k, v, mask=None, temp=1.0):
#     scaled_q = q
#     if temp != 1.0 and temp > 0:
#         scaled_q = q * (1.0 / temp)**.5
#     out = scaled_dot_product_attention(scaled_q, k, v, is_causal=mask is not None and q.shape[1] > 1)        
#     return out

# class LocalOut(nn.Module):
#     def __init__(self, dims: int, head: int):
#         super().__init__()
#         head_dim = dims // head
#         self.head_dim = head_dim
#         self.query_module = nn.Linear(head_dim, head_dim)
#         self.key_module = nn.Linear(head_dim, head_dim)
#         self.value_module = nn.Linear(head_dim, head_dim)
#         self.out_proj = nn.Linear(head_dim, head_dim)
    
#     def _reshape_to_output(self, x):
#         return x

# class attentiona(nn.Module):
#     def __init__(self, dims: int, head: int, max_iter: int = 3, threshold: float = 0.01, factor: float = 0.1, dropout: float = 0.1, temp = 1.0):
#         super(attentiona, self).__init__()
#         self.q,  self.k,  self.v,  self.o, self.lna, self.lnb = qkv_init(dims, head)
#         self.dims = dims
#         self.head = head
#         self.head_dim = dims // head
#         self.max_iter = max_iter
#         self.threshold = nn.Parameter(torch.tensor(threshold))
#         self.temp = nn.Parameter(torch.tensor(temp), requires_grad=True)        
#         self.factor = nn.Parameter(torch.tensor(factor))
#         self.lnc = nn.LayerNorm(self.head_dim, bias=False)
#         self.lnd = nn.LayerNorm(self.head_dim, bias=False)     
#         self.attn_local = LocalOut(self.head_dim)   

#     def _focus(self, x: Tensor, xa: Optional[Tensor] = None, mask: Optional[Tensor] = None):
#         z = default(xa, x)
#         q, k, v = create_qkv(self.dims, self.head, self.q, self.k, self.v, self.lna(x), self.lna(z))    

#         iteration = 0
#         temp = self.temp.item()
#         prev_out = torch.zeros_like(q)
#         attn_out = torch.zeros_like(q)
#         threshold = self.threshold.item()
#         factor = self.factor.item()
#         qcur = q

#         while iteration < self.max_iter:
#             eff_span = min(qcur.shape[1], k.shape[1])
#             if xa is not None:
#                 eff_span = min(eff_span, xa.shape[1])
#             if eff_span == 0: 
#                 break

#             qiter = qcur[:, :, :eff_span, :]
#             kiter = k[:, :, :eff_span, :]
#             viter = v[:, :, :eff_span, :]
#             q = self.attn_local.query_module(qiter)
#             k = self.attn_local.key_module(kiter)
#             v = self.attn_local.value_module(viter)

#             iter_mask = None
#             if mask is not None:
#                 if mask.dim() == 4: 
#                     iter_mask = mask[:, :, :eff_span, :eff_span]
#                 elif mask.dim() == 2: 
#                     iter_mask = mask[:eff_span, :eff_span]

#             attn_iter = calculate_attention(
#                 self.lnc(q), self.lnd(k), v,
#                 mask=iter_mask, temp=temp)

#             iter_out = torch.zeros_like(qcur)
#             iter_out[:, :, :eff_span, :] = attn_iter
#             diff = torch.abs(iter_out - prev_out).mean()
#             dthresh = threshold + factor * diff
#             if diff < dthresh and iteration > 0:
#                 attn_out = iter_out
#                 break

#             prev_out = iter_out.clone()
#             qcur = qcur + iter_out
#             attn_out = iter_out
#             iteration += 1
#             temp += 0.005

#         output = attn_out.permute(0, 2, 1, 3).flatten(start_dim=2)
#         return self.o(output), None

    # def _slide_win_local(self, x: Tensor, win_size: int, span_len: int, mask: Optional[Tensor] = None) -> Tensor:

    #     batch, ctx, dims = x.shape
    #     output = torch.zeros_like(x)
    #     num_win = (ctx + win_size - 1) // win_size

    #     for i in range(num_win):
    #         qstart = i * win_size
    #         qend = min(qstart + win_size, ctx)
    #         win_qlen = qend - qstart
    #         if win_qlen == 0: 
    #             continue

    #         kstart = max(0, qend - span_len)
    #         kend = qend
    #         qwin = x[:, qstart:qend, :]
    #         kwin = x[:, kstart:kend, :]

    #         win_mask = None
    #         if mask is not None:
    #             if mask.dim() == 4:
    #                 win_mask = mask[:, :, qstart:qend, kstart:kend]
    #             elif mask.dim() == 2:
    #                 win_mask = mask[qstart:qend, kstart:kend]

    #         attn_out, _ = self._focus(x=qwin, xa=kwin, mask=win_mask)
    #         output[:, qstart:qend, :] = attn_out
    #     return output

    # def forward(self, x: Tensor, xa: Optional[Tensor] = None, mask: Optional[Tensor] = None, 
    #             use_sliding_win: bool = False, win_size: int = 512, span_len: int = 1024) -> Tensor:
    #     if use_sliding_win:
    #         return self._slide_win_local(x, win_size, span_len, mask)
    #     else:
    #         output, _ = self._focus(x, xa, mask)
    #         return output
   
class KVCache(nn.Module):
    def __init__(self, max_batch_size, max_seq_length, n_heads, head_dim, dtype=torch.bfloat16):
        super().__init__()
        cache_shape = (max_batch_size, n_heads, max_seq_length, head_dim)
        self.register_buffer('k_cache', torch.zeros(cache_shape, dtype=dtype))
        self.register_buffer('v_cache', torch.zeros(cache_shape, dtype=dtype))

    def update(self, input_pos, k_val, v_val):
        # input_pos: [S], k_val: [B, H, S, D]
        assert input_pos.shape[0] == k_val.shape[2]

        k_out = self.k_cache
        v_out = self.v_cache
        k_out[:, :, input_pos] = k_val  # pyright: ignore[reportIndexIssue]
        v_out[:, :, input_pos] = v_val # pyright: ignore[reportIndexIssue]

        return k_out, v_out

def mel_scale_scalar(freq: float) -> float:
    return 1127.0 * math.log(1.0 + freq / 700.0)

def mel_scale(freq: Tensor) -> Tensor:
    return 1127.0 * (1.0 + freq / 700.0).log()

def trace_x(func):
    def wrapper(*args, **kwargs):
        print(f"Calling {func.__name__}")
        result = func(*args, **kwargs)
        if isinstance(result, torch.Tensor):
            print(f"  {func.__name__} returned shape: {result.shape}")
        return result
    return wrapper

def track_x(new_x, operation=""): 
    """ track_x(x, "x") """
    x_id = [id(new_x)]
    if new_x is None:
        return new_x
    current_id = id(new_x)
    if current_id != x_id[0]:
        print(f"x FLOW: {x_id[0]} → {current_id} in {operation}")
        x_id[0] = current_id
    else:
        print(f"x REUSE: {current_id} in {operation}")
    return new_x

def track_xa(new_xa, operation=""): 
    """ track_xa(xa, "xa - decoder") """
    xa_id = [id(new_xa)] if new_xa is not None else [None]
    if new_xa is None:
        return new_xa
    current_id = id(new_xa)
    if current_id != xa_id[0]:
        print(f"xa FLOW: {xa_id[0]} → {current_id} in {operation}")
        xa_id[0] = current_id  # pyright: ignore[reportArgumentType, reportCallIssue]
    else:
        print(f"xa REUSE: {current_id} in {operation}")
    return new_xa

def get_activation(act: str) -> nn.Module:
    """Get activation function by name."""
    act_map = {
        "gelu": nn.GELU(), 
        "relu": nn.ReLU(), 
        "sigmoid": nn.Sigmoid(), 
        "tanh": nn.Tanh(), 
        "swish": nn.SiLU(), 
        "tanhshrink": nn.Tanhshrink(), 
        "softplus": nn.Softplus(), 
        "softshrink": nn.Softshrink(), 
        "leaky_relu": nn.LeakyReLU(), 
        "elu": nn.ELU()
    }
    return act_map.get(act, nn.GELU())

def get_generation_config(param):
    return GenerationConfig(    # type: ignore
        max_length=param.text_ctx,
        pad_token_id=getattr(param, "pad_token_id", 0),
        bos_token_id=getattr(param, "bos_token_id", 1),
        eos_token_id=getattr(param, "eos_token_id", 2),
        do_sample=False,
        num_beams=1,
        early_stopping=False,
        length_penalty=1.0,
        no_repeat_ngram_size=0,
        repetition_penalty=1.0,
        temperature=1.0,
        decoder_start_token_id=1,
        is_multilingual=False,
        use_cache=False,
        return_timestamps=False)

class feature_encoder(nn.Module):
    def __init__(self, mels, input_dims, dims, head, layer, act, features, feature=None, use_rope=False, spec_shape=None, debug=[], attend_feature=False, target_length=None):
        """
        Feature encoder for audio processing.
        """
        super().__init__()

        self.dims = dims
        self.head = head
        self.head_dim = dims // head  
        self.dropout = 0.01 
        self.use_rope = use_rope
        self.attend_feature = attend_feature
        self.target_length = target_length
        self.feature = feature

        self.debug = debug
        act_fn = get_activation(act)

        if self.attend_feature:
            # self.q, self.k, self.v, self.o, self.scale = qkv_init(dims, head)
            self.mlp = nn.Sequential(nn.Linear(dims, dims), nn.ReLU(), nn.Linear(dims, dims))
        else:
            self.q, self.k, self.v, self.o, self.scale = None, None, None, None, None
            self.mlp = None

        self.spectrogram = nn.Sequential(
            Conv1d(mels, dims, kernel_size=3), act_fn,
            Conv1d(dims, dims, kernel_size=3), act_fn,
            Conv1d(dims, dims, kernel_size=3, groups=dims), act_fn)

        self.waveform = nn.Sequential(
            Conv1d(1, dims//4, kernel_size=15, stride=4, padding=7), act_fn,
            Conv1d(dims//4, dims//2, kernel_size=7, stride=2, padding=3), act_fn,
            Conv1d(dims//2, dims, kernel_size=5, stride=2, padding=2), act_fn)

        self.pitch = nn.Sequential(
            Conv1d(1, dims, kernel_size=7, stride=1, padding=3), act_fn,
            Conv1d(dims, dims, kernel_size=5, stride=1, padding=2), act_fn,
            Conv1d(dims, dims, kernel_size=3, stride=1, padding=1, groups=dims), act_fn)

        if use_rope:
            # if spec_shape is not None:
            self.positional = lambda length, dims, max_tscale: sinusoids(length, dims, max_tscale)
            self.rope = rotary(dims=dims, head=head, radii=False, debug=[], use_pbias=False, axial=False, spec_shape=spec_shape) # type: ignore
        else:
            self.rope = None # type: ignore
            self.positional = lambda length, dims, max_tscale: sinusoids(length, dims, max_tscale)
        self.norm = RMSNorm(dims)

    def rope(self, x, xa=None, mask=None, feats=None, feature=None, layer=None):
        if isinstance(x, int):
            ctx = x 
        elif isinstance(x, torch.Tensor):
            ctx = x.shape[1] if x.dim() > 1 else x.shape[0]
            batch, ctx, dims = x.shape[0], ctx, x.shape[-1]

            x = x.view(batch, ctx, self.head, self.head_dim).permute(0, 2, 1, 3)
        freqs = self.rope(ctx, feats=feats, feature=feature, layer=layer)
        x = self.rope.apply_rotary(x, freqs)  # pyright: ignore[reportOptionalSubscript, reportAttributeAccessIssue]
        x = x.permute(0, 2, 1, 3).contiguous().view(batch, ctx, dims)
        return x

    def mel_scalar(self, freq: float) -> float:
        return 1127.0 * math.log(1.0 + freq / 700.0)

    def forward(self, x, xa=None, mask=None, feats=None, feature=None, layer=None, max_tscale=36000):
        target_length = x.shape[1] if self.target_length is None else self.target_length

        if feature == "pitch":
            xp = x.clone()
            enc_dict = feats if feats is not None else {}
            enc_dict = dict(enc_dict)  
            enc_dict["f0"] = xp
            # xp = self.mel_scalar(xp.mean())
            # print(f"Using pitch scalar: {xp}")
            # max_tscale = xp*300
            # print(f"Using max_tscale: {max_tscale}")
            feats = enc_dict
            if x.dim() == 2:
                x = x.unsqueeze(0)
            x = self.pitch(x).permute(0, 2, 1)
  
        if feature == "phase":
            if x.dim() == 2:
                x = x.unsqueeze(0)
            x = self.pitch(x).permute(0, 2, 1)

        if feature == "waveform":
            if x.dim() == 2:
                x = x.unsqueeze(0)
            x = self.waveform(x).permute(0, 2, 1)
            if target_length and x.shape[1] != self.target_length:
                x = F.adaptive_avg_pool1d(x.transpose(1, 2), target_length).transpose(1, 2)
        
        if feature == "harmonics":
            if x.dim() == 2:
                x = x.unsqueeze(0)
            x = self.spectrogram(x).permute(0, 2, 1)

        if feature == "aperiodic":
            if x.dim() == 2:
                x = x.unsqueeze(0)
            x = self.spectrogram(x).permute(0, 2, 1)            

        if feature == "spectrogram":
            if x.dim() == 2:
                x = x.unsqueeze(0)
            x = self.spectrogram(x).permute(0, 2, 1)

        if self.use_rope:
            x = x + self.positional(x.shape[1], x.shape[-1], max_tscale).to(device, dtype)
            x = self.rope(x=x, xa=None, mask=None, feats=feats, feature=feature, layer=layer)
        else:
            max_tscale = x.shape[1] * 1000 if max_tscale is None else max_tscale
            x = x + self.positional(x.shape[1], x.shape[-1], max_tscale).to(device, dtype)
        x = nn.functional.dropout(x, p=self.dropout, training=self.training)
        x = self.norm(x)

        # if self.attend_feature:
        #     xa = feats[feature]  # pyright: ignore[reportOptionalSubscript]
        #     if xa is not None:
        #         q, k, v = create_qkv(self.q, self.k, self.v, x=xa, xa=x, head=self.head)
        #         out, _ = calculate_attention(q, k, v, mask=None, temperature=1.0, is_causal=True)
        #         x = x + out

        x = nn.functional.dropout(x, p=self.dropout, training=self.training)
        x = self.norm(x)
        return x

class OneShot(nn.Module):
    def __init__(self, dims: int, head: int, scale: float = 0.3, features: Optional[List[str]] = None):
        super().__init__()
        if features is None:    
            features = ["spectrogram", "waveform", "pitch", "aperiodic", "harmonics"]
        self.head = head
        self.head_dim = dims // head
        self.scale = 1.0 // len(features) if features else scale

        self.q = Linear(dims, dims)
        self.k = Linear(dims, dims)

    def forward(self, x: Tensor, xa: Tensor, feature=None) -> Tensor | None:
        B, L, D = x.shape
        K = xa.size(1)
        q = self.q(x).view(B, L, self.head, self.head_dim).transpose(1,2)
        k = self.k(xa).view(B, K, self.head, self.head_dim).transpose(1,2)
        bias = (q @ k.transpose(-1, -2)) * self.scale / math.sqrt(self.head_dim)
        return bias

class curiosity(nn.Module):
    def __init__(self, d, h, bias=True):
        super().__init__()
        self.h  = h
        self.dh = d // h
        self.qkv = nn.Linear(d, d * 3, bias=bias)
        self.qkv_aux = nn.Linear(d, d * 3, bias=bias)
        self.o  = nn.Linear(d, d, bias=bias)
        self.g  = nn.Parameter(torch.zeros(h))

    def split(self, x):
        b, t, _ = x.shape
        return x.view(b, t, self.h, self.dh).transpose(1, 2)

    def merge(self, x):
        b, h, t, dh = x.shape
        return x.transpose(1, 2).contiguous().view(b, t, h * dh)

    def forward(self, x, xa, mask=None):
        q, k, v   = self.qkv(x).chunk(3, -1)
        qa, ka, va = self.qkv_aux(xa).chunk(3, -1)
        q, k, v   = map(self.split, (q, k, v))
        qa, ka, va = map(self.split, (qa, ka, va))
        dots      = (q @ k.transpose(-2, -1)) / self.dh**0.5
        dots_aux  = (q @ ka.transpose(-2, -1)) / self.dh**0.5
        if mask is not None: dots = dots.masked_fill(mask, -9e15)
        p   = dots.softmax(-1)
        pa  = dots_aux.softmax(-1)
        h_main = p  @ v
        h_aux  = pa @ va
        g = torch.sigmoid(self.g).view(1, -1, 1, 1)
        out = self.merge(h_main * (1 - g) + h_aux * g)
        return self.o(out)

class PositionalEncoding(nn.Module):
    def __init__(self, dims, ctx):
        super(PositionalEncoding, self).__init__()
        self.dims = dims
        self.ctx = ctx
        self.pe = self.get_positional_encoding(max_ctx=ctx)

    def get_positional_encoding(self, max_ctx):
        pe = torch.zeros(max_ctx, self.dims)
        position = torch.arange(0, max_ctx, dtype=torch.float32).unsqueeze(1)
        div_term = torch.exp(
            torch.arange(0, self.dims, 2, dtype=torch.float32)
            * (-math.log(10000.0) / self.dims)
        )
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(0)
        return pe.to(device)

    def forward(self, x):
        ctx = x.size(1)
        pe = self.pe[:, :ctx, :]
        x = x * math.sqrt(self.dims)
        x = x + pe
        return x

def valid(default_value, *items):
    for item in items:
        if item is not None:
            return item
    return default_value

def dict_to(d, device, dtype=dtype):
    return {k: v.to(device, dtype) if isinstance(v, torch.Tensor) else v 
            for k, v in d.items()}
    
def exists(v):
    return v is not None

def default(v, b):
    return v if exists(v) else b

class Conv1d(nn.Conv1d):
    def _conv_forward(
        self, x: Tensor, weight: Tensor, bias) -> Tensor:
        return super()._conv_forward(x, weight.to(x.device, x.dtype), None if bias is None else bias.to(x.device, x.dtype))

class Conv2d(nn.Conv2d):
    def _conv_forward(
        self, x: Tensor, weight: Tensor, bias) -> Tensor:
        return super()._conv_forward(x, weight.to(x.device, x.dtype), None if bias is None else bias.to(x.device, x.dtype))

class Linear(nn.Module):
    def __init__(self, in_features: int, out_features: int, bias: bool = True) -> None:
        super(Linear, self).__init__()
        self.linear = nn.Linear(in_features, out_features, bias=bias)
        init.xavier_uniform_(self.linear.weight)
        if bias:
            init.zeros_(self.linear.bias)
    def forward(self, x: Tensor) -> Tensor:
        return self.linear(x)
    
class RMSNorm(nn.Module):
    def __init__(self, dims: Union[int, Tensor, List, Tuple], 
                 eps = 1e-8, elementwise_affine = True):
        super(RMSNorm, self).__init__()
        if isinstance(dims, int):
            self.normalized_shape = (dims,)
        else:
            self.normalized_shape = tuple(dims)
        self.eps = eps
        self.elementwise_affine = elementwise_affine
        if self.elementwise_affine:
            self.weight = nn.Parameter(torch.empty(self.normalized_shape))  # type: ignore
            init.ones_(self.weight)  
        else:
            self.register_parameter("weight", None)
    def forward(self, x):
        return F.rms_norm(x, self.normalized_shape, self.weight, self.eps)  # type: ignore
    
def LayerNorm(x: Tensor, normalized_shape: Union[int, Tensor, List, Tuple],
               weight: Optional[Tensor] = None, bias: Optional[Tensor] = None,
               eps: float = 1e-5) -> Tensor:
    return F.layer_norm(x, normalized_shape, weight, bias, eps)  # type: ignore

def get_device():
    return torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

def get_dtype():
    return torch.float32 if torch.cuda.is_available() else torch.float64

def tox():
    return {"device": get_device(), "dtype": get_dtype()}

def sinusoids(ctx, dims, max_tscale=10000):
    assert dims % 2 == 0
    pos = torch.log(torch.tensor(float(max_tscale))) / (dims // 2 - 1)
    tscales = torch.exp(-pos * torch.arange(dims // 2, device=device, dtype=torch.float32))
    scaled = torch.arange(ctx, device=device, dtype=torch.float32).unsqueeze(1) * tscales.unsqueeze(0)
    position = torch.cat([torch.sin(scaled), torch.cos(scaled)], dim=1) 
    positional_embedding = nn.Parameter(position, requires_grad=True)
    return positional_embedding

class SelfCriticalRL(nn.Module):
    def __init__(self, model, tokenizer, reward_fn):
        super().__init__()
        self.model = model
        self.tokenizer = tokenizer
        self.reward_fn = reward_fn

    def forward(self, input_ids, features, labels=None, max_len=128, feature_name="spectrogram"):

        with torch.no_grad():
            greedy_ids = self.model.generate(input_ids=input_ids, **{feature_name: features}, max_length=max_len)
        greedy_text = [self.tokenizer.decode(ids) for ids in greedy_ids]
        sampled_ids = self.model.generate(input_ids=input_ids, **{feature_name: features}, max_length=max_len, do_sample=True, top_k=5)
        sampled_text = [self.tokenizer.decode(ids) for ids in sampled_ids]
        
        rewards = []
        baseline = []
        for s, g, ref in zip(sampled_text, greedy_text, labels): # type: ignore
            ref_text = self.tokenizer.decode(ref)
            rewards.append(self.reward_fn(s, ref_text))
            baseline.append(self.reward_fn(g, ref_text))
        rewards = torch.tensor(rewards, device=device, dtype=torch.float)
        baseline = torch.tensor(baseline, device=device, dtype=torch.float)
        advantage = rewards - baseline
        logits = self.model(input_ids=sampled_ids, **{feature_name: features})["logits"]  # logits: [batch, sampled_seq_len, vocab_size]
        log_probs = F.log_softmax(logits, dim=-1)
        log_probs_seq = torch.gather(log_probs, 2, sampled_ids.unsqueeze(-1)).squeeze(-1)
        log_probs_sum = log_probs_seq.sum(dim=1)
        loss = -(advantage * log_probs_sum).mean()
        return loss

class SelfTrainingModule(nn.Module):
    def __init__(self, model, tokenizer, quality_fn=None, threshold=0.8):
        super().__init__()
        self.model = model
        self.tokenizer = tokenizer
        self.quality_fn = quality_fn
        self.threshold = threshold

    def generate_pseudo_labels(self, unlabeled_batch, features, max_len=128, feature_name="spectrogram"):
        with torch.no_grad():
            pred_ids = self.model.generate(input_ids=unlabeled_batch, **{feature_name: features}, max_length=max_len)

        if self.quality_fn is not None:
            quality_scores = self.quality_fn(pred_ids, self.model, features)
            mask = quality_scores > self.threshold
            pred_ids = pred_ids[mask]
        return pred_ids

    def forward(self, unlabeled_batch, features, max_len=128, feature_name="spectrogram"):
        pseudo_labels = self.generate_pseudo_labels(unlabeled_batch, features, max_len, feature_name=feature_name)
        logits = self.model(input_ids=unlabeled_batch, **{feature_name: features}, labels=pseudo_labels)["logits"]
        loss = nn.functional.cross_entropy(
            logits.view(-1, logits.shape[-1]), pseudo_labels.view(-1), ignore_index=0)
        return loss

def confidence_indicator(pred_ids, model, features):
    with torch.no_grad():
        logits = model(input_ids=pred_ids, **features)["logits"]
    probs = torch.softmax(logits, dim=-1)
    max_probs, _ = probs.max(dim=-1)
    return max_probs.mean(dim=1)

def wer_reward(hyp, ref):
    hyp_words = hyp.split()
    ref_words = ref.split()
    d = [[0] * (len(ref_words)+1) for _ in range(len(hyp_words)+1)]
    for i in range(len(hyp_words)+1):
        d[i][0] = i
    for j in range(len(ref_words)+1):
        d[0][j] = j
    for i in range(1, len(hyp_words)+1):
        for j in range(1, len(ref_words)+1):
            if hyp_words[i-1] == ref_words[j-1]:
                d[i][j] = d[i-1][j-1]
            else:
                d[i][j] = 1 + min(d[i-1][j], d[i][j-1], d[i-1][j-1])
    wer = d[-1][-1] / max(1, len(ref_words))
    return -wer  # negative WER as reward

def clean_ids(ids, pad_token_id=0, bos_token_id=1, eos_token_id=2):
    if isinstance(ids, torch.Tensor):
        ids = ids.tolist()
    return [int(id) for id in ids if id != -100 and id != pad_token_id and id != bos_token_id and id != eos_token_id]

def clean_batch(batch_ids, pad_token_id=0, bos_token_id=1, eos_token_id=2):
    return [clean_ids(seq, pad_token_id, bos_token_id, eos_token_id) for seq in batch_ids]

def setup_tokenizer(dir: str):
    from tokenizers import Tokenizer
    tokenizer = Tokenizer.from_file(f"{dir}")
    orig_encode = tokenizer.encode
    orig_decode = tokenizer.decode

    def enc(text, add_special_tokens=True):
        ids = orig_encode(text).ids
        if not add_special_tokens:
            sp_ids = [tokenizer.token_to_id(t) for t in ["<PAD>", "<BOS>", "<EOS>"]]
            ids = [id for id in ids if id not in sp_ids]
        return ids

    def bdec(ids_list, pad_token_id=0, bos_token_id=1, eos_token_id=2, skip_special_tokens=True):
        results = []
        if isinstance(ids_list, torch.Tensor):
            ids_list = ids_list.tolist()
        elif isinstance(ids_list, np.ndarray):
            ids_list = ids_list.tolist()
        for ids in ids_list:
            ids = [int(id) for id in ids if id not in (pad_token_id, bos_token_id, eos_token_id, -100)]
            results.append(orig_decode(ids))
        return results

    def dec(ids, pad_token_id=0, bos_token_id=1, eos_token_id=2):
        ids = [int(id) for id in ids if id not in (pad_token_id, bos_token_id, eos_token_id, -100)]
        return orig_decode(ids)

    def save_pretrained(save_dir):
        os.makedirs(save_dir, exist_ok=True)
        tokenizer.save(f"{save_dir}/tokenizer.json")

    tokenizer.encode = enc
    tokenizer.batch_decode = bdec
    tokenizer.decode = dec
    tokenizer.save_pretrained = save_pretrained
    tokenizer.pad_token_id = 0
    tokenizer.bos_token_id = 1
    tokenizer.eos_token_id = 2
    return tokenizer

def tokenize_pitch(pitch_features, target_length):
    pitch_len = pitch_features.shape[-1]
    token_len = target_length
    if pitch_len > token_len:
        pitch_tokens = F.adaptive_avg_pool1d(pitch_features, token_len)
    else:
        pitch_tokens = F.interpolate(pitch_features, token_len)
    return pitch_tokens

def load_wave(wave_data, sample_rate=16000):

    if isinstance(wave_data, str):
        waveform, sample_rate = torchaudio.load(uri=wave_data, normalize=False)
    elif isinstance(wave_data, dict):
        waveform = torch.tensor(data=wave_data["array"]).float()
        sample_rate = wave_data["sampling_rate"]  # noqa: F841
    else:
        raise TypeError("Invalid wave_data format.")
    return waveform

def world_to_mel(sp, ap, sample_rate=16000, n_mels=128):
    import librosa
    mel_basis = librosa.filters.mel(sr=sample_rate, n_fft=1024, n_mels=n_mels)
    mel_basis = torch.from_numpy(mel_basis).float()
    sp_mel = torch.matmul(sp, mel_basis.T)  # (frames, 128)
    ap_mel = torch.matmul(ap, mel_basis.T)  # (frames, 128)
    return sp_mel, ap_mel

def extract_features(batch, tokenizer, waveform=False, spec=False, pitch_tokens=False, pitch=False, harmonics=False, sample_rate=16000, hop_length=256, mode="mean", debug=False, phase_mod=False, crepe=False, aperiodics=False, dummy=False):

    # import torch
    # import torchaudio
    # import torchaudio.functional as F
    # import torchaudio.transforms as T

    torch_windows = {
        'hann': torch.hann_window,
        'hamming': torch.hamming_window,
        'blackman': torch.blackman_window,
        'bartlett': torch.bartlett_window,
        'ones': torch.ones,
        None: torch.ones,
    }
    if dummy:
        return {
            "spectrogram": torch.zeros((1, 128, 100)),
            "f0": torch.zeros((1, 100)),
            "pitch_tokens": torch.zeros((1, 100)),
            "pitch": torch.zeros((1, 100)),
            "harmonics": torch.zeros((1, 128, 100)),
            "aperiodics": torch.zeros((1, 128, 100)),
            "crepe_time": None,
            "crepe_frequency": None,
            "crepe_confidence": None,
            "crepe_activation": None,
        }

    audio = batch["audio"]
    sample_rate = audio["sampling_rate"]
    # audio_length = len(audio["array"]) / audio["sampling_rate"]
    labels = tokenizer.encode(batch["transcription"])
    # sentence_length = len(batch["transcription"]) 
    wav = load_wave(wave_data=audio, sample_rate=sample_rate)
    
    def crepe_predict(wav, sample_rate, viterbi=False):
        import torchcrepe
        wav = wav.numpy().astype(np.float32)
        time, frequency, confidence, activation = torchcrepe.predict(
            wav, sample_rate=sample_rate, viterbi=viterbi)
        crepe_time = torch.from_numpy(time)
        crepe_frequency = torch.from_numpy(frequency)
        crepe_confidence = torch.from_numpy(confidence)
        crepe_activation = torch.from_numpy(activation)
        return crepe_time, crepe_frequency, crepe_confidence, crepe_activation

    if crepe:
        crepe_time, crepe_frequency, crepe_confidence, crepe_activation = crepe_predict(wav, sample_rate, viterbi=True)

    else:
        crepe_time = None
        crepe_frequency = None
        crepe_confidence = None
        crepe_activation = None

    def spectrogram(wav, sample_rate, n_fft=1024, hop_length=256, window_fn=torch.hann_window):
        if isinstance(window_fn, str):
            window_fn = torch_windows[window_fn]
        if window_fn is None:
            window_fn = torch.ones(n_fft)
        if isinstance(window_fn, torch.Tensor):
            window_fn = window_fn.to(device)
        return torchaudio.functional.spectrogram(
            wav, n_fft=n_fft, hop_length=hop_length, win_length=n_fft,
            window=window_fn, center=True, pad_mode="reflect", power=1.0)

    def mel_spectrogram(wav, sample_rate):

        spectrogram_config = {
            "hop_length": 256,
            "f_min": 150,
            "f_max": 2000,
            "n_mels": 128,
            "n_fft": 1024,
            "sample_rate": 16000,
            "pad_mode": "constant",
            "center": True, 
            "power": 1.0,
            "window_fn": torch.hann_window,
            "mel_scale": "htk",
            "norm": None,
            "normalized": False,
        }
        transform = torchaudio.transforms.MelSpectrogram(**spectrogram_config)
        mel_spectrogram = transform(wav)
        log_mel = torch.clamp(mel_spectrogram, min=1e-10).log10()
        log_mel = torch.maximum(log_mel, log_mel.max() - 8.0)
        spectrogram_tensor = (log_mel + 4.0) / 4.0
        return spectrogram_tensor

    if spec: 
        spectrogram_tensor = mel_spectrogram(wav, sample_rate)

    def mfcc(wav, sample_rate, n_mels=128, n_fft=1024, hop_length=256, window_fn=torch.hann_window):
        transform = torchaudio.transforms.MFCC(
            sample_rate=sample_rate,
            n_mfcc=n_mels,
            melkwargs={
                "n_fft": n_fft,
                "hop_length": hop_length,
                "window_fn": window_fn,
                "n_mels": n_mels,
                "center": True,
                "pad_mode": "reflect",
                "norm": None,
                "mel_scale": "htk",
            }
        )
        mfcc_tensor = transform(wav)
        return mfcc_tensor

    # def compute_pitch(wav, sample_rate, hop_length=256):
    #     # pitch = F.detect_pitch_frequency(wav, sample_rate)
    #     # f0 = pitch
    #     import pyworld as pw
    #     wav_np = wav.numpy().astype(np.float64)
    #     f0, t = pw.dio(wav_np, sample_rate, frame_period=hop_length / sample_rate * 1000)
    #     f0 = pw.stonemask(wav_np, f0, t, sample_rate)
    #     return f0, t

    def harmonics_and_aperiodics(wav, f0, t, sample_rate):
        import pyworld as pw
        wav_np = wav.numpy().astype(np.float64)
        sp = pw.cheaptrick(wav_np, f0, t, sample_rate, fft_size=256)
        ap = pw.d4c(wav_np, f0, t, sample_rate, fft_size=256)
        harmonic_tensor = torch.from_numpy(sp)
        aperiodic_tensor = torch.from_numpy(ap)
        harmonic_tensor = harmonic_tensor[:, :128].contiguous().T
        aperiodic_tensor = aperiodic_tensor[:, :128].contiguous().T
        harmonic_tensor = torch.where(harmonic_tensor == 0.0, torch.zeros_like(harmonic_tensor), harmonic_tensor / 1.0)
        aperiodic_tensor = torch.where(aperiodic_tensor == 0.0, torch.zeros_like(aperiodic_tensor), aperiodic_tensor / 1.0)
        return harmonic_tensor, aperiodic_tensor

    if pitch or pitch_tokens or harmonics or aperiodics:
        wavnp = wav.numpy().astype(np.float64)
        f0_np, t = pw.dio(wavnp, sample_rate, frame_period=hop_length / sample_rate * 1000)
        f0_np = pw.stonemask(wavnp, f0_np, t, sample_rate)

    if pitch_tokens:
        wav = torch.from_numpy(wavnp)
        t2 = torch.from_numpy(t)
        audio_duration = len(wav) / sample_rate
        T = len(labels)
        tok_dur_sec = audio_duration / T
        token_starts = torch.arange(T) * tok_dur_sec
        token_ends = token_starts + tok_dur_sec
        start_idx = torch.searchsorted(t2, token_starts, side="left")
        end_idx = torch.searchsorted(t2, token_ends, side="right")
        pitch_tok = torch.zeros(T, dtype=torch.float32)
        for i in range(T):
            lo, hi = start_idx[i], max(start_idx[i]+1, end_idx[i]) # type: ignore
            segment = f0_np[lo:hi]
            if mode == "mean":
                pitch_tok[i] = segment.mean()
            elif mode == "median":
                pitch_tok[i] = torch.median(segment)
            else:
                pitch_tok[i] = segment[-1]
        pitch_tok[pitch_tok < 100.0] = 0.0
        bos_pitch = pitch_tok[0] if len(pitch_tok) > 0 else 0.0
        pitch_tokens_tensor = torch.cat([torch.tensor([bos_pitch]), pitch_tok])
        pitch_tokens_tensor = torch.where(pitch_tokens_tensor == 0.0, torch.zeros_like(pitch_tokens_tensor), (pitch_tokens_tensor - 71.0) / (500.0 - 71.0))
    else:
        pitch_tokens_tensor = None

    if phase_mod:
        tframe = torch.mean(t2[1:] - t2[:-1])
        phi0 = 0.0
        omega = 2 * torch.pi * f0_tensor # type: ignore
        dphi = omega * tframe
        phi = torch.cumsum(dphi, dim=0) + phi0
        phase = torch.remainder(phi, 2 * torch.pi)
    else:
        phase = None

    if pitch:
        p_tensor = torchaudio.functional.detect_pitch_frequency(wav, sample_rate)
        # p_tensor = torch.from_numpy(f0_np)
        # p_tensor = p_tensor.unsqueeze(0) 
    else:
        p_tensor = None

    if harmonics or aperiodics:
        spnp = pw.cheaptrick(wavnp, f0_np, t, sample_rate, fft_size=256)
        apnp = pw.d4c(wavnp, f0_np, t, sample_rate, fft_size=256)
        harmonic_tensor = torch.from_numpy(spnp)
        aperiodic_tensor = torch.from_numpy(apnp)
        harmonic_tensor = harmonic_tensor[:, :128].contiguous().T
        aperiodic_tensor = aperiodic_tensor[:, :128].contiguous().T
        harmonic_tensor = torch.where(harmonic_tensor == 0.0, torch.zeros_like(harmonic_tensor), harmonic_tensor / 1.0)
        aperiodic_tensor = torch.where(aperiodic_tensor == 0.0, torch.zeros_like(aperiodic_tensor), aperiodic_tensor / 1.0)
    else:
        harmonic_tensor = None
        aperiodic_tensor = None

    if waveform:
        wave_tensor = wav
    else:
        wave_tensor = None

    if dummy:   
        if spectrogram_tensor is not None:
            dummy_tensor = torch.ones_like(spectrogram_tensor)
        elif p_tensor is not None:
            dummy_tensor = torch.ones_like(p_tensor) 
        elif pitch_tokens_tensor is not None:
            dummy_tensor = torch.ones_like(pitch_tokens_tensor)
        else:
            batch_size = 128
            seq_len = 1024
            dummy_tensor = torch.ones(batch_size, seq_len)
            dummy_tensor = dummy_tensor.to(device)
    else:
        dummy_tensor = None
    if debug:
        print(f"['pitch_tokens']: {pitch_tokens_tensor.shape if pitch_tokens else None}")
        print(f"['harmonic']: {harmonic_tensor.shape if harmonics else None}")
        print(f"['aperiodic']: {aperiodic_tensor.shape if aperiodics else None}")
        print(f"['spectrogram']: {spectrogram_tensor.shape if spec else None}")
        print(f"['waveform']: {wave_tensor.shape if waveform else None}")
        print(f"['labels']: {len(labels) if labels else None}")
        print(f"['phase']: {phase.shape if phase else None}")
        print(f"['pitch']: {p_tensor.shape if pitch else None}")
        print(f"['crepe_time']: {crepe_time.shape if crepe else None}")  
        print(f"['crepe_frequency']: {crepe_frequency.shape if crepe else None}")
        print(f"['crepe_confidence']: {crepe_confidence.shape if crepe else None}")
        print(f"['crepe_activation']: {crepe_activation.shape if crepe else None}")
        print(f"['dummy']: {dummy_tensor.shape if dummy else None}")

    return {
        "waveform": wave_tensor if waveform else None,
        "spectrogram": spectrogram_tensor if spec else None,
        "pitch_tokens": pitch_tokens_tensor if pitch_tokens else None,
        "pitch": p_tensor if pitch else None,
        "harmonic": harmonic_tensor if harmonics else None,
        "aperiodic": aperiodic_tensor if aperiodics else None,  
        "labels": labels,
        "phase": phase if phase_mod else None,
        "crepe_time": crepe_time if crepe else None,
        "crepe_frequency": crepe_frequency if crepe else None,
        "crepe_confidence": crepe_confidence if crepe else None,
        "crepe_activation": crepe_activation if crepe else None,
        "dummy": dummy_tensor if dummy else None,
    }

# class PEncoder(nn.Module): # pitch encoder
#     def __init__(self, dims: int, head: int, layer: int, kernel_size: int, act: str, 
#                  max_seq_len: int, input_dims: int = 1, use_rope=False): 
#         super().__init__()
        
#         self.head = head
#         self.head_dim = dims // head
#         self.dims = dims
#         self.use_rope=use_rope
#         self.dropout_rate = 0.01 
#         act_fn = get_activation(act)

#         self.positional_encoding = nn.Parameter(torch.randn(1, max_seq_len, dims)) 

#         if use_rope:
#                 self.rope = rotary(dims=dims, head=head, radii=False, debug=[], use_pbias=False, axial=False, spec_shape=spec_shape)# type: ignore
#         else:
#             self.rope = None
#             self.positional = lambda length, dims, max_tscale: sinusoids(length, dims, max_tscale)

#         self.attend_pitch = False 
#         if self.attend_pitch:
#             self.mlp = nn.Sequential(
#                 nn.Linear(dims, dims),
#                 nn.ReLU(),
#                 nn.Linear(dims, dims),
#             )
#         else:
#             self.mlp = None

#         self.pitch_encoder = nn.Sequential(
#             nn.Conv1d(input_dims, dims, kernel_size=kernel_size, stride=1, padding=kernel_size // 2), act_fn,
#             nn.Conv1d(dims, dims, kernel_size=kernel_size - 2, stride=1, padding=(kernel_size - 2) // 2), act_fn,
#             nn.Conv1d(dims, dims, kernel_size=kernel_size - 4, stride=1, padding=(kernel_size - 4) // 2, groups=dims), act_fn
#         )
    
#         def rope_to_feature(self, x, xa=None, mask=None, feats=None, feature="pitch", layer="PEncoder"):
#             batch, ctx, dims = x.shape
#             x = x.view(batch, ctx, head, self.head_dim).permute(0, 2, 1, 3)
#             freqs = self.rope(ctx, feats=feats, feature=feature, layer=layer) # type: ignore
#             x = self.rope.apply_rotary(x, freqs)# type: ignore
#             x = x.permute(0, 2, 1, 3).contiguous().view(batch, ctx, dims)
#             return x

#         self.norm = nn.LayerNorm(dims) 

#     def forward(self, x: Tensor, xa: Optional[Tensor] = None, mask: Optional[Tensor] = None, 
#                 feats: Optional[Any] = None, feature: str = "pitch", layer: str = "PEncoder", 
#                 audio_duration: Optional[float] = None, sample_rate: Optional[int] = None, 
#                 labels_len: Optional[int] = None, f0_np: Optional[np.ndarray] = None, 
#                 t2_np: Optional[np.ndarray] = None, mode: str = "mean") -> Tensor:
        
#         if x.dim() == 2 and feature == "pitch":
#             x_processed = x.unsqueeze(1) # Input to pitch_encoder: (batch_size, 1, num_pitch_tokens)
#             x_processed = self.pitch_encoder(x_processed)             # Output: (batch_size, dims, num_pitch_tokens)
#             x = x_processed.permute(0, 2, 1) # Reassign to x for consistency

#         if self.use_rope:
#             pass # Placeholder for RoPE application
    
#         seq_len = x.shape[1]
#         # x = x + self.positional_encoding[:, :seq_len, :] 
#         x = x + sinusoids(x.shape[1], x.shape[-1], 36000).to(device, dtype)

#         if self.mlp is not None:
#             x = self.mlp(x)

#         x = nn.functional.dropout(x, p=self.dropout_rate, training=self.training)
#         x = self.norm(x)    

#         return x

def plot_waveform(waveform, sr, title="Waveform", ax=None):
    waveform = waveform.numpy()

    num_channels, num_frames = waveform.shape
    time_axis = torch.arange(0, num_frames) / sr

    if ax is None:
        _, ax = plt.subplots(num_channels, 1)
    ax.plot(time_axis, waveform[0], linewidth=1)
    ax.grid(True)
    ax.set_xlim([0, time_axis[-1]])
    ax.set_title(title)

def plot_spectrogram(specgram, title=None, ylabel="freq_bin", ax=None):
    import librosa
    if ax is None:
        _, ax = plt.subplots(1, 1)
    if title is not None:
        ax.set_title(title)
    ax.set_ylabel(ylabel)
    ax.imshow(librosa.power_to_db(specgram), origin="lower", aspect="auto", interpolation="nearest")

def plot_fbank(fbank, title=None):
    fig, axs = plt.subplots(1, 1)
    axs.set_title(title or "Filter bank")
    axs.imshow(fbank, aspect="auto")
    axs.set_ylabel("frequency bin")
    axs.set_xlabel("mel bin")

def plot_pitch(waveform, sr, pitch):
    figure, axis = plt.subplots(1, 1)
    axis.set_title("Pitch Feature")
    axis.grid(True)

    end_time = waveform.shape[1] / sr
    time_axis = torch.linspace(0, end_time, waveform.shape[1])
    axis.plot(time_axis, waveform[0], linewidth=1, color="gray", alpha=0.3)

    axis2 = axis.twinx()
    time_axis = torch.linspace(0, end_time, pitch.shape[1])
    axis2.plot(time_axis, pitch[0], linewidth=2, label="Pitch", color="green")

    axis2.legend(loc=0)

def prepare_datasets(tokenizer, token, sanity_check=False, sample_rate=16000, streaming=False,
        load_saved=False, save_dataset=False, cache_dir=None, extract_args=None, max_ctx=2048):

    if extract_args is None:
        extract_args = {
        "waveform": False,
        "spec": False,
        "f0": False,
        "pitch_tokens": False,
        "pitch": False,
        "harmonic": False,
        "aperiodic": False,
        "sample_rate": 16000,
        "hop_length": 256,
        "mode": "mean",
        "debug": False,
        "phase_mod": False,
        "crepe": False,
        "dummy": False,
        }

    if load_saved:
        if cache_dir is None:
            cache_dir = "./processed_datasets"
        else:
            cache_dir = cache_dir

        os.makedirs(cache_dir, exist_ok=True)
        cache_file_train = os.path.join(cache_dir, "train.arrow")
        cache_file_test = os.path.join(cache_dir, "test.arrow")

        if os.path.exists(cache_file_train) and os.path.exists(cache_file_test):
            from datasets import Dataset
            train_dataset = Dataset.load_from_disk(cache_file_train)
            test_dataset = Dataset.load_from_disk(cache_file_test)
            return train_dataset, test_dataset   

    if sanity_check:
        test = load_dataset(
            "google/fleurs", "en_us", token=token, split="test", trust_remote_code=True, streaming=streaming).cast_column("audio", Audio(sampling_rate=sample_rate)).take(1)
        dataset = test.map(lambda x: extract_features(x, tokenizer, **extract_args), remove_columns=test.column_names)
        train_dataset = dataset
        test_dataset = dataset
        return train_dataset, test_dataset
    else:
        def filter_func(x):
            return (0 < len(x["transcription"]) < max_ctx and
                    len(x["audio"]["array"]) > 0 and
                    len(x["audio"]["array"]) < max_ctx * 160)

        # raw_train  = load_dataset("mozilla-foundation/common_voice_17_0", "en", token=token, split="train", trust_remote_code=True, streaming=True).rename_column("sentence", "transcription")
        # raw_test = load_dataset("mozilla-foundation/common_voice_17_0", "en", token=token, split="test", trust_remote_code=True, streaming=True).rename_column("sentence", "transcription").take(1000)

        raw_train = load_dataset(
            "google/fleurs", "en_us", token=token, split="train", trust_remote_code=True, streaming=streaming).take(1000)
        raw_test = load_dataset(
            "google/fleurs", "en_us", token=token, split="test", trust_remote_code=True, streaming=streaming).take(100)

        raw_train = raw_train.filter(filter_func).cast_column("audio", Audio(sampling_rate=sample_rate))
        raw_test = raw_test.filter(filter_func).cast_column("audio", Audio(sampling_rate=sample_rate))
        train_dataset = raw_train.map(lambda x: extract_features(x, tokenizer, **extract_args), remove_columns=raw_train.column_names)
        test_dataset = raw_test.map(lambda x: extract_features(x, tokenizer, **extract_args), remove_columns=raw_test.column_names)
        train_dataset.save_to_disk(cache_file_train) if save_dataset is True else None
        test_dataset.save_to_disk(cache_file_test) if save_dataset is True else None
        return train_dataset, test_dataset

class tgate(nn.Module):
    def __init__(self, dims, num_types=4):
        super().__init__()
        self.gates = nn.ModuleList([nn.Sequential(Linear(dims, 1), nn.Sigmoid()) for _ in range(num_types)])
        self.classifier = nn.Sequential(Linear(dims, num_types), nn.Softmax(dim=-1))
    def forward(self, x):
        types = self.classifier(x)
        gates = torch.stack([gate(x) for gate in self.gates], dim=-1)
        cgate = torch.sum(gates * types.unsqueeze(2), dim=-1)
        return cgate

def get_feature_encoder(feature: str, mels: int, input_dims: int, dims: int, head: int, layer: int, act=None, features=None) -> nn.Module:
    if feature == "spectrogram":
        return FEncoder(mels=mels, input_dims=input_dims, dims=dims, head=head)
    elif feature == "waveform":
        return WEncoder(input_dims, dims, head, layer, act, feature, features)
    elif feature == "pitch":
        return PEncoder(input_dims, dims, head, layer, act, feature, features)
    else:
        raise ValueError(f"Unknown feature type: {feature}")

class FEncoder(nn.Module):
    def __init__(self, mels, input_dims, dims, head, layer, act, feature, features, use_rope=False, spec_shape=None, debug=[]):
        super().__init__()
        
        self.head = head
        self.head_dim = dims // head  
        self.dropout = 0.01 
        self.use_rope = use_rope
        self.dims = dims
        self.debug = debug
        self.feature = feature
        self.mels = mels
        self.input_dims = input_dims
        act_fn = get_activation(act)

        self.encoder = nn.Sequential(
            Conv1d(mels, dims, kernel_size=3, stride=1, padding=1), act_fn,
            Conv1d(dims, dims, kernel_size=3, stride=1, padding=1), act_fn,
            Conv1d(dims, dims, kernel_size=3, stride=1, padding=1, groups=dims), act_fn)

        if use_rope:
            if spec_shape is not None:
                self.rope = rotary(dims=dims, head=head, radii=False, debug=[], use_pbias=False, axial=False, spec_shape=spec_shape) # type: ignore
        else:
            self.rope = None
            self.positional = lambda length, dims, max_tscale: sinusoids(length, dims, max_tscale)
        self.norm = RMSNorm(dims)

    def apply_rope_to_features(self, x, xa=None, mask=None, feats=None, feature="audio", layer="FEncoder"):
        batch, ctx, dims = x.shape
        x = x.view(batch, ctx, self.head, self.head_dim).permute(0, 2, 1, 3)
        freqs = self.rope(ctx, feats=feats, feature=feature, layer=layer)# type: ignore
        x = self.rope.apply_rotary(x, freqs)# type: ignore
        x = x.permute(0, 2, 1, 3).contiguous().view(batch, ctx, dims)

        return x

    def forward(self, x, xa=None, mask=None, feats=None, feature="audio", layer="FEncoder"):
        x = self.encoder(x).permute(0, 2, 1)
        if self.use_rope:
            x = self.apply_rope_to_features(x, xa=xa, mask=mask, feats=feats, feature=feature, layer=layer)
        else:
            x = x + self.positional(x.shape[1], x.shape[-1], 10000).to(device, dtype)
        x = nn.functional.dropout(x, p=self.dropout, training=self.training)
        print(f"feature encoder: {x.shape} {feature}") if "fencoder" in self.debug else None
        x = self.norm(x)
        return x

class WEncoder(nn.Module): # waveform encoder
    def __init__(self, input_dims, dims, head, layer, kernel_size, act, use_rope=False, debug=[], spec_shape=None):
        super().__init__()
        
        self.head = head
        self.head_dim = dims // head
        self.dropout = 0.01
        self.use_rope = use_rope
        self.dims = dims
        self.debug = debug
        act_fn = get_activation(act)
        self.target_length = None
        self.encoder = nn.Sequential(
            Conv1d(input_dims, dims//4, kernel_size=15, stride=4, padding=7), act_fn,
            Conv1d(dims//4, dims//2, kernel_size=7, stride=2, padding=3), act_fn,
            Conv1d(dims//2, dims, kernel_size=5, stride=2, padding=2), act_fn)
            
        if use_rope:
            if spec_shape is not None:
                self.rope = rotary(dims=dims, head=head, radii=False, debug=[], use_pbias=False, axial=False, spec_shape=spec_shape)# type: ignore
        else:
            self.rope = None
            self.positional = lambda length, dims, max_tscale: sinusoids(length, dims, max_tscale)
        self.norm = RMSNorm(dims)

    def apply_rope_to_features(self, x, xa=None, mask=None, feats=None, feature="waveform", layer="WEncoder"):
        batch, ctx, dims = x.shape
        x = x.view(batch, ctx, self.head, self.head_dim).permute(0, 2, 1, 3)
        freqs = self.rope(ctx, feats=feats, feature=feature, layer=layer)# type: ignore
        x = self.rope.apply_rotary(x, freqs)# type: ignore
        x = x.permute(0, 2, 1, 3).contiguous().view(batch, ctx, dims)
        return x
        
    def forward(self, x, xa=None, mask=None, feats= None, feature="waveform", layer = "WEncoder"):
        x = self.encoder(x).permute(0, 2, 1)  # (batch, time, dims)
        if self.target_length and x.shape[1] != self.target_length:
            x = F.adaptive_avg_pool1d(x.transpose(1, 2), self.target_length).transpose(1, 2)
        if self.use_rope:
            x = self.apply_rope_to_features(x, xa=xa, mask=mask, feats=feats, feature=feature, layer=layer)
        else:
            x = x + self.positional(x.shape[1], x.shape[-1], 10000).to(device, dtype)
        x = nn.functional.dropout(x, p=self.dropout, training=self.training)
        print(f"waveform encoder: {x.shape} {feature}") if "fencoder" in self.debug else None
        return self.norm(x)

class PEncoder(nn.Module): # pitch encoder
    def __init__(self, input_dims, dims, head, layer, act, use_rope=False, debug=[], one_shot=False, spec_shape=None):
        super().__init__()
        
        self.head = head
        self.head_dim = dims // head
        self.dims = dims
        self.dropout = 0.01
        self.use_rope = use_rope
        self.debug = debug
        act_fn = get_activation(act)

        self.attend_pitch = False

        if self.attend_pitch:
            self.q, self.k, self.v, self.o, self.scale = qkv_init(dims, head)
            self.mlp = nn.Sequential(
                nn.Linear(dims, dims),
                nn.ReLU(),
                nn.Linear(dims, dims),
            )
        else:
            self.q, self.k, self.v, self.o, self.scale = None, None, None, None, None
            self.mlp = None

        self.pitch_encoder = nn.Sequential(
            Conv1d(input_dims, dims, kernel_size=7, stride=1, padding=3), act_fn,
            Conv1d(dims, dims, kernel_size=5, stride=1, padding=2), act_fn,
            Conv1d(dims, dims, kernel_size=3, stride=1, padding=1, groups=dims), act_fn)

        if use_rope:
                self.rope = rotary(dims=dims, head=head, radii=False, debug=[], use_pbias=False, axial=False, spec_shape=spec_shape)# type: ignore
        else:
            self.rope = None
            self.positional = lambda length, dims, max_tscale: sinusoids(length, dims, max_tscale)
        self.norm = RMSNorm(dims)
        
    def rope_to_feature(self, x, xa=None, mask=None, feats=None, feature="pitch", layer="PEncoder"):
        batch, ctx, dims = x.shape
        x = x.view(batch, ctx, self.head, self.head_dim).permute(0, 2, 1, 3)
        freqs = self.rope(ctx, feats=feats, feature=feature, layer=layer) # type: ignore
        x = self.rope.apply_rotary(x, freqs)# type: ignore
        x = x.permute(0, 2, 1, 3).contiguous().view(batch, ctx, dims)
        return x
        
    def forward(self, x, xa=None, mask=None, feats= None, feature="pitch", layer="PEncoder"):
        # f0=x
        # freqs = self.rope(f0.shape[1], feats=feats, feature=feature, layer=layer)
        if x.dim() == 2:
            x = x.unsqueeze(0)
        # if feature == "pitch":
        x = self.pitch_encoder(x).permute(0, 2, 1)

        if self.use_rope:
            x = self.rope_to_feature(x, xa=xa, mask=mask, feats=feats, feature=feature, layer=layer)
    
        x = x + self.positional(x.shape[1], x.shape[-1], 10000).to(device, dtype)
        if self.mlp is not None:
            x = self.mlp(x)

        if self.attend_pitch:
            if xa is not None:
                q, k, v = create_qkv(self.q, self.k, self.v, x=xa, xa=x, head=self.head)
                out, _ = calculate_attention(q, k, v, mask=None, temperature=1.0, is_causal=True)
                x = x + out

        # x = nn.functional.dropout(x, p=self.dropout, training=self.training)
        x = self.norm(x)    
        print(f"Pitch encoder: {x.shape} {feature}") if "fencoder" in self.debug else None
        return x

@dataclass
class DataCollator:
    tokenizer: Any

    def __call__(self, features: List[Dict[str, torch.Tensor]]) -> Dict[str, torch.Tensor]:
        all_keys = set()
        for f in features:
            all_keys.update(f.keys())
        batch = {}
        pad_token_id = getattr(self.tokenizer, 'pad_token_id', 0)
        bos_token_id = getattr(self.tokenizer, 'bos_token_id', 1)
        eos_token_id = getattr(self.tokenizer, 'eos_token_id', 2)

        for key in all_keys:
            if key == "labels":
                labels_list = [f["labels"] for f in features]
                max_len = max(len(l) for l in labels_list)  # noqa: E741
                all_ids, all_labels = [], []

                for label in labels_list:
                    label_list = label.tolist() if isinstance(label, torch.Tensor) else label
                    decoder_input = [bos_token_id] + label_list
                    label_eos = label_list + [eos_token_id]
                    input_len = max_len + 1 - len(decoder_input)
                    label_len = max_len + 1 - len(label_eos)
                    padded_input = decoder_input + [pad_token_id] * input_len
                    padded_labels = label_eos + [pad_token_id] * label_len
                    all_ids.append(padded_input)
                    all_labels.append(padded_labels)
                batch["input_ids"] = torch.tensor(all_ids, dtype=torch.long)
                batch["labels"] = torch.tensor(all_labels, dtype=torch.long)

            elif key in ["spectrogram", "waveform", "pitch", "harmonic", "aperiodic", "pitch_tokens", "f0", "phase", "crepe_time", "crepe_frequency", "crepe_confidence", "crepe_activation", "dummy"]:

                items = [f[key] for f in features if key in f]
                items = [item for item in items if item is not None]
                if not items:  
                    continue
                items = [torch.tensor(item) if not isinstance(item, torch.Tensor) else item for item in items]
                max_len = max(item.shape[-1] for item in items)
                padded = []
                for item in items:
                    pad_width = max_len - item.shape[-1]
                    if pad_width > 0:
                        pad_item = F.pad(item, (0, pad_width), mode='constant', value=pad_token_id)
                    else:
                        pad_item = item
                    padded.append(pad_item)
                batch[key] = torch.stack(padded)
                # if key == "spectrogram":
                #     batch["spectrogram"] = batch[key]
        return batch

# import tiktoken
# import torch
# from torch.utils.data import Dataset, DataLoader

# class tokenize(Dataset):
#     def __init__(self, txt, tokenizer, max_length, stride):

#         self.input_ids = []
#         self.labels = []

#         token_ids = tokenizer.encode(txt, allowed_special={"<eos>"})

#         for i in range(0, len(token_ids) - max_length, stride):

#             input_chunk = token_ids[i:i + max_length]
#             target_chunk = token_ids[i + 1: i + max_length + 1]
#             self.input_ids.append(torch.tensor(input_chunk))
#             self.labels.append(torch.tensor(target_chunk))

#     def __len__(self):
#         return len(self.input_ids)

#     def __getitem__(self, idx):
#         return self.input_ids[idx], self.labels[idx]

# def create_dataloader_v1(txt, batch_size, max_length, stride, shuffle=True, drop_last=True, num_workers=0):
#     tokenizer = tiktoken.get_encoding("gpt2")
#     dataset = tokenize(txt, tokenizer, max_length, stride)
#     dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=shuffle, drop_last=drop_last, num_workers=num_workers)
#     return dataloader

# def custom_collate_fn(batch, tokenizer_pad_token_id):
#     max_len_in_batch = max(len(seq) for seq in batch)

#     padded_input_ids = []
#     attention_masks = []
#     for seq in batch:
#         padded_seq = F.pad(seq, (0, max_len_in_batch - len(seq)), value=tokenizer_pad_token_id)
#         attention_mask = torch.ones(max_len_in_batch)
#         attention_mask[len(seq):] = 0 
#         padded_input_ids.append(padded_seq)
#         attention_masks.append(attention_mask)

#     input_ids_tensor = torch.stack(padded_input_ids)
#     attention_mask_tensor = torch.stack(attention_masks)
#     labels_tensor = input_ids_tensor.clone() 

#     return {
#         'input_ids': input_ids_tensor,
#         'attention_mask': attention_mask_tensor,
#         'labels': labels_tensor
#     }

# with open("the-verdict.txt", "r", encoding="utf-8") as f:
#     raw_text = f.read()

# vocab_size = 50257
# output_dim = 256
# context_length = 1024

# token_embedding_layer = torch.nn.Embedding(vocab_size, output_dim)
# pos_embedding_layer = torch.nn.Embedding(context_length, output_dim)

# batch_size = 8
# max_length = 4
# dataloader = create_dataloader_v1(
#     raw_text,
#     batch_size=batch_size,
#     max_length=max_length,
#     stride=max_length
# )

def levenshtein(reference_words, hypothesis_words):
    m, n = len(reference_words), len(hypothesis_words)
    dist_matrix = [[0 for _ in range(n+1)] for _ in range(m+1)]
    for i in range(m+1):
        dist_matrix[i][0] = i
    for j in range(n+1):
        dist_matrix[0][j] = j
    for i in range(1, m+1):
        for j in range(1, n+1):
            if reference_words[i-1] == hypothesis_words[j-1]:
                dist_matrix[i][j] = dist_matrix[i-1][j-1]
            else:
                substitution = dist_matrix[i-1][j-1] + 1
                insertion = dist_matrix[i][j-1] + 1
                deletion = dist_matrix[i-1][j] + 1
                dist_matrix[i][j] = min(substitution, insertion, deletion)
    return dist_matrix[m][n]

def wer_batch(references, hypotheses):
    total_errors = 0
    total_words = 0
    for ref, hyp in zip(references, hypotheses):
        ref_words = ref.lower().split()
        errors = levenshtein(ref_words, hyp.lower().split()) 
        total_errors += errors
        total_words += len(ref_words)
    return (total_errors / total_words) * 100 if total_words > 0 else 0.0

def compute_metrics(pred, tokenizer=None, model=None, print_pred=False, num_samples=0, logits=None):
    def clean(ids, pad_token_id=0, bos_token_id=1, eos_token_id=2):
        
        if isinstance(ids, torch.Tensor):
            ids = ids.tolist()
        if isinstance(ids[0], (list, torch.Tensor, np.ndarray)):
            return [[int(i) for i in seq if i not in (-100, pad_token_id, bos_token_id, eos_token_id)] for seq in ids]
        else:
            return [int(i) for i in ids if i not in (-100, pad_token_id, bos_token_id, eos_token_id)]

    pred_ids = pred.predictions
    label_ids = pred.label_ids

    if isinstance(pred_ids, tuple):
        pred_ids = pred_ids[0]

    if not isinstance(pred_ids, torch.Tensor):
        pred_ids = torch.tensor(pred_ids)

    label_ids = clean(label_ids)
    pred_ids = clean(pred_ids)
    pred_str = tokenizer.batch_decode(pred_ids)
    label_str = tokenizer.batch_decode(label_ids)

    if print_pred:
        for i in range(min(num_samples, len(pred_ids))):

            print(f"Pred tokens: {pred_ids[i]}")
            print(f"Label tokens: {label_ids[i]}")
            print(f"Pred: '{pred_str[i]}'")
            print(f"Label: '{label_str[i]}'")
            print("-" * 40)
    
    wer = wer_batch(label_str, pred_str)
    if model is not None:
        trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad) / 1_000_000
        efficiency_score = (100 - wer) / trainable_params if trainable_params > 0 else 0.0
    else:
        trainable_params = 0.0
        efficiency_score = 0.0

    return {
        "wer": float(wer),
        "efficiency_score": float(efficiency_score),
    }

def preprocess_logits_for_metrics(logits, labels):
    pred_ids = torch.argmax(logits, dim=-1)
    return pred_ids, labels

def hilbert_transform(x):
    N = x.shape[-1]
    xf = torch.fft.rfft(x)
    h = torch.zeros(N // 2 + 1, device=x.device, dtype=x.dtype)
    if N % 2 == 0:
        h[0] = h[N//2] = 1
        h[1:N//2] = 2
    else:
        h[0] = 1
        h[1:(N+1)//2] = 2
    return torch.fft.irfft(xf * h, n=N)

def analytic_signal(x):
    return x + 1j * hilbert_transform(x)

def hilbert_transform_2d(x, dim=-1):
    N = x.shape[dim]
    if dim == -1 or dim == len(x.shape) - 1:
        xf = torch.fft.rfft(x)
    else:
        xf = torch.fft.rfft(x, dim=dim)
    h_shape = [1] * len(x.shape)
    h_shape[dim] = N // 2 + 1
    h = torch.zeros(h_shape, device=x.device, dtype=x.dtype)
    if dim == -1 or dim == len(x.shape) - 1:
        if N % 2 == 0:
            h[..., 0] = h[..., -1] = 1
            h[..., 1:-1] = 2
        else:
            h[..., 0] = 1
            h[..., 1:] = 2
    else:
        pass
    return torch.fft.irfft(xf * h, n=N, dim=dim)

def hilbert_transform_true_2d(x):
    xf = torch.fft.rfft2(x)
    h1, h2 = torch.meshgrid(
        torch.fft.rfftfreq(x.shape[-2]) * 2 - 1,
        torch.fft.rfftfreq(x.shape[-1]) * 2 - 1,
        indexing='ij')
    h = -1j / (math.pi * (h1 + 1j*h2))
    h[0, 0] = 0 
    return torch.fft.irfft2(xf * h.to(x.device))

def process_spectrogram_with_hilbert(spec):
    analytic = spec + 1j * hilbert_transform(spec)
    envelope = torch.abs(analytic)
    phase = torch.angle(analytic)
    return envelope, phase

# import torch
# import torch.nn as nn
# import torch.nn.functional as F
# from torch import Tensor
# from typing import Optional, Tuple
# import numpy as np
# from torch.nn.functional import scaled_dot_product_attention
# from torch.cuda.amp import autocast
# from torch.nn import LayerNorm, Linear
# import logging

# logging.basicConfig(level=logging.WARNING)
# log = logging.getLogger(__name__)

# class ProjectionModule(nn.Module):
#     """
#     Projects input embeddings into query, key, or value representations
#     for multi-head attention, handling scaling for Q/K.
#     """
#     def __init__(self, dims: int, head: int, proj_type: str = "query", use_bias: bool = True):
#         """
#         Args:
#             dims: Input and output dimension.
#             head: Number of attention heads.
#             proj_type: Type of projection ("query", "key", "value").
#             use_bias: Whether to use bias in the linear layer.
#         """
#         super().__init__()
#         assert dims % head == 0, f"dims ({dims}) must be divisible by head ({head})"
#         assert proj_type in ["query", "key", "value"], \
#                f"proj_type must be 'query', 'key', or 'value', got {proj_type}"

#         self.dims = dims
#         self.head = head
#         self.head_dim = dims // head
#         self.proj_type = proj_type
#         self.scale = self.head_dim ** -0.5 if proj_type != "value" else 1.0
#         self.proj = Linear(in_features=dims, out_features=dims, bias=use_bias)
#         self.init_weights()

#     def init_weights(self):
#         """Initialize projection weights."""
#         nn.init.normal_(tensor=self.proj.weight, std=0.02)
#         if self.proj.bias is not None:
#             nn.init.zeros_(tensor=self.proj.bias)

#     def forward(self, x: Tensor) -> Tensor:
#         """
#         Applies projection, scaling (for Q/K), and reshapes for multi-head attention.

#         Args:
#             x: Input tensor of shape (batch, seq_len, dims).

#         Returns:
#             Projected tensor of shape (batch, head, seq_len, head_dim).
#         """
#         batch, seq_len, _ = x.shape
#         proj = self.proj(x)

#         proj = proj.view(batch, seq_len, self.head, self.head_dim).permute(0, 2, 1, 3)

#         if self.proj_type in ["query", "key"]:
#             proj = proj * self.scale
#         return proj

# def calculate_attention(
#     q: Tensor,
#     k: Tensor,
#     v: Tensor,
#     mask: Optional[Tensor] = None,
#     temperature: float = 1.0,
#     use_sdpa: bool = True,
#     is_causal: bool = False,
#     dropout_p: float = 0.0
# ) -> Tuple[Tensor, Optional[Tensor]]:
#     """
#     Calculates scaled dot-product attention.

#     Uses torch.nn.functional.scaled_dot_product_attention if use_sdpa is True
#     and inputs are compatible, otherwise falls back to manual implementation.

#     Args:
#         q: Query tensor (Batch, Heads, SeqLen_Q, HeadDim). Already scaled if needed.
#         k: Key tensor (Batch, Heads, SeqLen_K, HeadDim). Already scaled if needed.
#         v: Value tensor (Batch, Heads, SeqLen_K, HeadDim).
#         mask: Attention mask. Can be boolean (True means ignore) or float (-inf means ignore).
#               Shape should be broadcastable to (Batch, Heads, SeqLen_Q, SeqLen_K).
#         temperature: Softmax temperature scaling. Applied *before* softmax.
#         use_sdpa: Flag to attempt using PyTorch's optimized SDPA implementation.
#         is_causal: If True, applies a causal mask (for decoder self-attention).
#                    Used only if mask is None and use_sdpa is True.
#         dropout_p: Dropout probability for attention weights.

#     Returns:
#         A tuple containing:
#         - attn_output: Attention output tensor (Batch, Heads, SeqLen_Q, HeadDim).
#         - attn_weights: Attention weights tensor (Batch, Heads, SeqLen_Q, SeqLen_K),
#                         or None if SDPA implementation doesn't return them or if fallback used.
#                         *Note: SDPA's default doesn't return weights, requires specific backend support.*
#                         *Manual path always returns weights.*
#     """
#     batch_size, num_heads, q_len, head_dim = q.shape
#     k_len = k.size(2)

#     temp_scale = 1.0 / temperature if temperature > 0 else 1.0

#     attn_output, attn_weights = None, None

#     if use_sdpa:
#         try:
#             if temperature != 1.0:
#                  raise NotImplementedError("SDPA does not directly support temperature scaling. Use manual path or scale Q.")

#             attn_output = scaled_dot_product_attention(
#                 q, k, v,
#                 attn_mask=mask,
#                 dropout_p=dropout_p,
#                 is_causal=is_causal and mask is None
#             )
#             attn_weights = None
#             return attn_output, attn_weights
#         except (RuntimeError, NotImplementedError) as e:
#             log.warning(f"SDPA failed or not used ({e}), falling back to manual attention.")
#     attn_scores = torch.matmul(q, k.transpose(-2, -1)) * temp_scale

#     if mask is not None:
#         if mask.dim() == 2:
#             mask = mask.unsqueeze(0).unsqueeze(0)
#         elif mask.dim() == 3:
#             mask = mask.unsqueeze(1)

#         expected_mask_shape = (batch_size, num_heads, q_len, k_len)
#         if mask.shape != expected_mask_shape:
#              try:
#                  mask = mask.expand(expected_mask_shape)
#              except RuntimeError:
#                  raise ValueError(f"Mask shape {mask.shape} is not compatible with attention scores shape {expected_mask_shape}")

#         if mask.dtype == torch.bool:
#             attn_scores = attn_scores.masked_fill(mask, float("-inf"))
#         else:
#             attn_scores = attn_scores + mask

#     attn_weights = F.softmax(attn_scores, dim=-1)

#     if dropout_p > 0.0:
#         attn_weights = F.dropout(attn_weights, p=dropout_p)

#     attn_output = torch.matmul(attn_weights, v)

#     return attn_output, attn_weights

# class BaseAttention(nn.Module):
#     """Base class for attention mechanisms with common functionality."""
#     use_sdpa = True

#     def __init__(self, dims: int, head: int, max_dist: int = 512, dropout: float = 0.0):
#         """
#         Args:
#             dims: Embedding dimension.
#             head: Number of attention heads.
#             max_dist: Maximum attention distance (used by some subclasses).
#             dropout: Dropout probability for attention weights.
#         """
#         super().__init__()
#         assert dims % head == 0, f"dims ({dims}) must be divisible by head ({head})"
#         self.dims = dims
#         self.head = head
#         self.head_dim = dims // head
#         self.max_dist = max_dist
#         self.dropout = dropout

#     def _shape(self, tensor: torch.Tensor) -> torch.Tensor:
#         """
#         Reshape tensor from (batch, seq_len, dims) to
#         (batch, head, seq_len, head_dim) for multi-head attention.
#         """
#         batch, seq_len, _ = tensor.shape
#         return tensor.view(batch, seq_len, self.head, self.head_dim).transpose(1, 2).contiguous()

#     def _reshape_to_output(self, attn_output: Tensor) -> Tensor:
#         """
#         Reshape attention output from (batch, head, seq_len, head_dim)
#         back to (batch, seq_len, dims).
#         """
#         batch, _, seq_len, _ = attn_output.shape
#         return attn_output.transpose(1, 2).contiguous().view(batch, seq_len, self.dims)

# class AttentionCombiner(BaseAttention):
#     """
#     Computes attention given Q, K, V projections and applies an output projection.
#     Assumes Q, K, V inputs are already projected and appropriately shaped/scaled.
#     """
#     def __init__(self, dims: int, head: int, use_bias: bool = True, dropout: float = 0.0):
#         """
#         Args:
#             dims: Embedding dimension.
#             head: Number of attention heads.
#             use_bias: Whether to use bias in the output projection.
#             dropout: Dropout probability for attention weights.
#         """
#         super().__init__(dims, head, dropout=dropout)
#         self.out = Linear(in_features=dims, out_features=dims, bias=use_bias)
#         self._init_weights()

#     def _init_weights(self):
#         """Initialize output projection weights."""
#         nn.init.normal_(tensor=self.out.weight, std=0.02)
#         if self.out.bias is not None:
#             nn.init.zeros_(tensor=self.out.bias)

#     # @autocast('cuda', enabled=torch.cuda.is_available())
#     def forward(self, q: Tensor, k: Tensor, v: Tensor,
#                 mask: Optional[Tensor] = None, is_causal: bool = False) -> Tensor:
#         """
#         Processes Q, K, V through attention and output projection.

#         Args:
#             q: Query tensor (Batch, Heads, SeqLen_Q, HeadDim). Assumed scaled.
#             k: Key tensor (Batch, Heads, SeqLen_K, HeadDim). Assumed scaled.
#             v: Value tensor (Batch, Heads, SeqLen_K, HeadDim).
#             mask: Attention mask.
#             is_causal: Whether to apply causal masking (if mask is None).

#         Returns:
#             Output tensor (Batch, SeqLen_Q, Dims).
#         """
#         attn_output, _ = calculate_attention(
#             q, k, v, mask=mask,
#             temperature=1.0,
#             use_sdpa=BaseAttention.use_sdpa,
#             is_causal=is_causal,
#             dropout_p = self.dropout
#         )

#         output = self._reshape_to_output(attn_output)
#         return self.out(output)

# class AdaptiveUpdateAttention(BaseAttention):
#     """
#     Attention implementation where Key and Value representations are cached
#     and only updated based on content-dependent predictors. Suitable for
#     encoder layers or cross-attention where K/V context changes less frequently.

#     Note: Current implementation focuses on conditional update based on *current*
#     input, not standard auto-regressive KV caching for generation.
#     """
#     def __init__(self, dims: int, head: int, max_dist: int = 512, update_threshold: float = 0.5, dropout: float = 0.0):
#         """
#         Args:
#             dims: Embedding dimension.
#             head: Number of attention heads.
#             max_dist: Maximum attention distance (inherited, may not be directly used here).
#             update_threshold: Threshold for sigmoid output of predictors to trigger update.
#             dropout: Dropout probability for attention weights.
#         """
#         super().__init__(dims, head, max_dist, dropout=dropout)

#         self.query_module = ProjectionModule(dims, head, "query")
#         self.key_module = ProjectionModule(dims, head, "key")
#         self.value_module = ProjectionModule(dims, head, "value")
#         self.combiner = AttentionCombiner(dims, head, dropout=dropout)

#         self.key_update_predictor = nn.Sequential(
#             Linear(dims, dims // 4), nn.ReLU(), Linear(dims // 4, 1), nn.Sigmoid())
#         self.value_update_predictor = nn.Sequential(
#             Linear(dims, dims // 4), nn.ReLU(), Linear(dims // 4, 1), nn.Sigmoid())

#         self.update_threshold = update_threshold
#         self.stored_key_cache: Optional[Tensor] = None
#         self.stored_value_cache: Optional[Tensor] = None
#         self.reset_cache_on_forward = True

#     def _should_update(self, x: torch.Tensor, predictor: nn.Module) -> torch.Tensor:
#         """Predict whether K or V should be updated based on content."""
#         avg_rep = x.mean(dim=1)
#         update_prob = predictor(avg_rep)
#         return update_prob > self.update_threshold

#     def forward(self, x: Tensor, xa: Optional[Tensor] = None,
#                 mask: Optional[Tensor] = None,
#                 is_causal: bool = False) -> Tensor:
#         """
#         Process inputs with adaptive K/V update mechanism.

#         Args:
#             x: Input tensor for queries (Batch, SeqLen_Q, Dims).
#             xa: Optional input tensor for keys/values (for cross-attention).
#                 If None, uses x for self-attention (Batch, SeqLen_KV, Dims).
#             mask: Attention mask.
#             is_causal: Whether attention should be causal.

#         Returns:
#             Output tensor (Batch, SeqLen_Q, Dims).
#         """
#         if self.reset_cache_on_forward:
#              self.stored_key_cache = None
#              self.stored_value_cache = None

#         batch, ctx_q, _ = x.shape
#         q = self.query_module(x)

#         kv_input = xa if xa is not None else x
#         ctx_kv = kv_input.size(1)

#         update_k_batch = self._should_update(kv_input, self.key_update_predictor)
#         update_v_batch = self._should_update(kv_input, self.value_update_predictor)

#         if self.stored_key_cache is None or self.stored_key_cache.shape[2] != ctx_kv or self.stored_key_cache.shape[0] != batch:
#             k = self.key_module(kv_input)
#             self.stored_key_cache = k
#         elif update_k_batch.any():
#             new_k = self.key_module(kv_input)
#             update_mask_k = update_k_batch.view(-1, 1, 1, 1).expand_as(self.stored_key_cache)
#             k = torch.where(update_mask_k, new_k, self.stored_key_cache)
#             self.stored_key_cache = k
#         else:
#             k = self.stored_key_cache

#         if self.stored_value_cache is None or self.stored_value_cache.shape[2] != ctx_kv or self.stored_value_cache.shape[0] != batch:
#             v = self.value_module(kv_input)
#             self.stored_value_cache = v
#         elif update_v_batch.any():
#             new_v = self.value_module(kv_input)
#             update_mask_v = update_v_batch.view(-1, 1, 1, 1).expand_as(self.stored_value_cache)
#             v = torch.where(update_mask_v, new_v, self.stored_value_cache)
#             self.stored_value_cache = v
#         else:
#             v = self.stored_value_cache

#         output = self.combiner(q, k, v, mask=mask, is_causal=is_causal)
#         return output

# class Refiner:
#     """
#     Q-learning based agent to refine parameters (e.g., attention span).
#     Operates outside the standard backpropagation loop.
#     """
#     def __init__(self, states: int, actions: int, alpha: float = 0.1, gamma: float = 0.9, epsilon: float = 0.1):
#         self.states = states
#         self.actions = actions
#         self.R = {}
#         self.alpha = alpha
#         self.gamma = gamma
#         self.epsilon = epsilon
#         self.default_value = 0.0

#     def get_value(self, state: int, action: int) -> float:
#         """Get Q-value for state-action pair."""
#         return self.R.get((state, action), self.default_value)

#     def set_value(self, state: int, action: int, value: float):
#         """Set Q-value for state-action pair."""
#         self.R[(state, action)] = value

#     def choose_action(self, state: int) -> int:
#         """Choose action using epsilon-greedy strategy."""
#         if np.random.random() < self.epsilon:
#             return np.random.randint(self.actions)
#         else:
#             action_values = [self.get_value(state, a) for a in range(self.actions)]
#             return np.argmax(action_values).item()

#     def update(self, state: int, action: int, reward: float, next_state: int):
#         """Update Q-value using the Q-learning rule."""
#         next_values = [self.get_value(next_state, a) for a in range(self.actions)]
#         best_next_value = max(next_values) if next_values else self.default_value

#         old_value = self.get_value(state, action)
#         td_target = reward + self.gamma * best_next_value
#         td_error = td_target - old_value
#         new_value = old_value + self.alpha * td_error
#         self.set_value(state, action, new_value)

# class Predictor(nn.Module):
#     """Neural predictor for estimating a scale value (e.g., for adaptive span)."""
#     def __init__(self, dims: int):
#         super().__init__()
#         self.linear = Linear(in_features=dims, out_features=1)
#         self._init_weights()

#     def _init_weights(self):
#         """Initialize predictor weights."""
#         nn.init.xavier_normal_(self.linear.weight)
#         if self.linear.bias is not None:
#            nn.init.zeros_(self.linear.bias)

#     def forward(self, x: Tensor) -> Tensor:
#         """
#         Predicts a scale factor (0-1) from input features.

#         Args:
#             x: Input tensor (Batch, SeqLen, Dims) or (Batch, Dims).

#         Returns:
#             Scale tensor (Batch, 1).
#         """
#         if x.dim() > 2:
#             x = x.mean(dim=1)
#         scale = torch.sigmoid(self.linear(x))
#         return scale

# class AdaptiveSpanAttention(BaseAttention):
#     """
#     Attention mechanism where the span is dynamically adjusted based on a
#     learnable parameter or predicted scale. This version focuses on slicing
#     the input sequence to the effective span.

#     Note: This implementation attends only to the *first* `eff_span` tokens.
#     For attending to a *relative* window, different logic (e.g., sliding window
#     or masking) would be needed in `calculate_attention`.
#     """
#     def __init__(self, dims: int, head: int, max_dist: int = 512,
#                  initial_span_scale: float = 1.0, learnable_scale: bool = True,
#                  sharpen: bool = True, temp_scale: float = 0.01, dropout: float = 0.0):
#         """
#         Args:
#             dims, head, max_dist, dropout: Standard BaseAttention params.
#             initial_span_scale: Initial value for the span scale.
#             learnable_scale: If True, span_scale is an nn.Parameter.
#             sharpen, temp_scale: Parameters for dynamic temperature adjustment.
#         """
#         super().__init__(dims, head, max_dist, dropout=dropout)
#         self.sharpen = sharpen
#         self.temp_scale = temp_scale
#         if learnable_scale:
#             self.span_scale = nn.Parameter(torch.tensor(initial_span_scale))
#         else:
#             self.register_buffer("span_scale", torch.tensor(initial_span_scale))

#         self.query_module = ProjectionModule(dims, head, "query")
#         self.key_module = ProjectionModule(dims, head, "key")
#         self.value_module = ProjectionModule(dims, head, "value")
#         self.out_proj = Linear(dims, dims)

#     @autocast('cuda', enabled=torch.cuda.is_available())
#     def forward(self, x: Tensor, xa: Optional[Tensor] = None,
#                 mask: Optional[Tensor] = None,
#                 span_scale_override: Optional[Tensor] = None,
#                 is_causal: bool = False) -> Tuple[Tensor, Optional[Tensor]]:
#         """
#         Computes attention over an adaptively determined span.

#         Args:
#             x: Input tensor for Q (Batch, SeqLen_Q, Dims).
#             xa: Optional input for K/V (Batch, SeqLen_KV, Dims). If None, use x.
#             mask: External attention mask.
#             span_scale_override: Optional tensor (Batch, 1) or scalar to override internal span_scale.
#             is_causal: Whether to apply causal masking.

#         Returns:
#             Tuple of (output tensor (Batch, SeqLen_Q, Dims), attention weights (optional)).
#         """
#         kv_input = xa if xa is not None else x
#         batch, ctx_q, _ = x.shape
#         ctx_kv = kv_input.size(1)

#         current_span_scale = span_scale_override if span_scale_override is not None else self.span_scale
#         if isinstance(current_span_scale, nn.Parameter):
#              span_scale_val = current_span_scale.sigmoid()
#         elif current_span_scale.numel() == 1:
#              span_scale_val = current_span_scale.expand(batch, 1)
#         else:
#              span_scale_val = current_span_scale

#         span_mean = span_scale_val.mean().item()
#         max_span_len = ctx_kv
#         target_span_len = max(1, int(max_span_len * span_mean))

#         eff_span = min(target_span_len, self.max_dist, ctx_q, ctx_kv)

#         if eff_span == 0:
#             return (torch.zeros_like(x), None)

#         q_span = x[:, :eff_span, :]
#         k_span = kv_input[:, :eff_span, :]
#         v_span = kv_input[:, :eff_span, :]

#         q_proj = self.query_module(q_span)
#         k_proj = self.key_module(k_span)
#         v_proj = self.value_module(v_span)

#         temperature = (1.0 + self.temp_scale * (1.0 - span_mean)
#                        if self.sharpen
#                        else 0.5 + self.temp_scale * span_mean)
#         temperature = max(temperature, 1e-3)

#         span_mask = None
#         if mask is not None:
#              if mask.dim() == 4:
#                  span_mask = mask[:, :, :eff_span, :eff_span]
#              elif mask.dim() == 2:
#                  span_mask = mask[:eff_span, :eff_span]
#         attn_output_span, attn_weights = calculate_attention(
#             q_proj, k_proj, v_proj,
#             mask=span_mask,
#             temperature=temperature,
#             use_sdpa=BaseAttention.use_sdpa,
#             is_causal=is_causal,
#             dropout_p=self.dropout
#         )

#         output_span = self._reshape_to_output(attn_output_span)
#         projected_output_span = self.out_proj(output_span)

#         output = torch.zeros_like(x)
#         output[:, :eff_span, :] = projected_output_span

#         return output, attn_weights

# class MyelinatedLayer(BaseAttention):
#     """
#     A complex Transformer layer featuring:
#     - Integrated local/global attention (via IntegratedAttention).
#     - Optional adapters within sub-layers.
#     - Node importance prediction for sparsity.
#     - MLP block.
#     - Working memory component.
#     - Potential layer skipping ("jumping") based on a learned policy.

#     (This version assumes IntegratedAttention is the core attention mechanism).
#     """
#     def __init__(self, dims: int, head: int, num_layers: int = 3,
#                  sparsity_threshold: float = 0.1, max_dist: int = 512,
#                  dropout: float = 0.1, mlp_ratio: int = 4):
#         super().__init__(dims, head, max_dist, dropout)
#         self.num_layers = num_layers
#         self.sparsity_threshold = sparsity_threshold

#         self.attention = IntegratedAttention(dims, head, max_dist=max_dist, dropout=dropout)

#         self.sub_layers = nn.ModuleList()
#         self.node_predictors = nn.ModuleList([
#             nn.Sequential(LayerNorm(dims), Linear(dims, 1), nn.Sigmoid())
#             for _ in range(num_layers)])

#         for i in range(num_layers):
#             self.sub_layers.append(nn.ModuleDict({
#                 'ln': LayerNorm(dims),
#                 'gate': nn.Sequential(Linear(dims, 1), nn.Sigmoid()),
#                 'adapter': Linear(dims, dims) if i % 2 == 0 else None
#             }))

#         self.policy_net = nn.Sequential(Linear(dims, 128), nn.ReLU(), Linear(128, num_layers))
#         self.jump_weights = nn.Parameter(torch.tensor([0.1, 0.05, 0.01]))

#         n_mlp = dims * mlp_ratio
#         self.mlp_gate = nn.Sequential(Linear(dims, 1), nn.Sigmoid())
#         self.mlp = nn.Sequential(Linear(dims, n_mlp), nn.GELU(), Linear(n_mlp, dims), nn.Dropout(dropout))
#         self.mlp_ln = LayerNorm(dims)

#         self.working_memory = nn.Parameter(torch.zeros(1, 1, dims))
#         self.memory_gate = nn.Sequential(Linear(dims, 1), nn.Sigmoid())

#         self.last_memory_gate_values: Optional[Tensor] = None

#     def predict_node_importance(self, x: Tensor, layer_idx: int) -> Tensor:
#         """Predict token importance mask (0.0 or 1.0) for sparsity."""
#         importance = self.node_predictors[layer_idx](x)
#         is_important = (importance > self.sparsity_threshold).float()
#         return is_important

#     def forward(self, x: Tensor, xa: Optional[Tensor] = None,
#                 mask: Optional[Tensor] = None, kv_cache: Optional[Tensor] = None,
#                 is_causal: bool = True) -> Tensor:
#         batch, ctx, _ = x.shape
#         working_memory = self.working_memory.expand(batch, 1, -1).to(x.device)
#         original_x = x

#         pooled_representation = x.mean(dim=1)
#         policy_logits = self.policy_net(pooled_representation)
#         policy = F.softmax(policy_logits, dim=-1)

#         jump_history = []
#         i = 0
#         last_processed_output = x

#         while i < self.num_layers:
#             layer = self.sub_layers[i]

#             node_importance_mask = self.predict_node_importance(x, i)

#             if node_importance_mask.mean() < 0.2 and i > 0:
#                 i += 1
#                 jump_history.append(f"skip_low_imp->{i}")
#                 continue

#             norm_x = layer['ln'](x)

#             current_attn_mask = node_importance_mask.permute(0, 2, 1)
#             if mask is not None:
#                  pass

#             attn_output = self.attention(
#                 norm_x * node_importance_mask,
#                 xa=xa,
#                 mask=mask,
#                 kv_cache=kv_cache,
#                 is_causal=is_causal
#             )

#             if layer['adapter'] is not None:
#                 attn_output = layer['adapter'](attn_output)

#             gate_value = layer['gate'](norm_x)
#             x = x + gate_value * attn_output * node_importance_mask
#             last_processed_output = x

#             memory_gate = self.memory_gate(x.mean(dim=1, keepdim=True))
#             current_mean_x = x.mean(dim=1, keepdim=True)
#             working_memory = memory_gate * working_memory + (1 - memory_gate) * current_mean_x

#             if i < self.num_layers - 1:
#                  jump_prob_dist = policy[:, 1:]
#                  jump_prob = jump_prob_dist.sum(dim=-1)

#                  should_jump_batch = torch.rand_like(jump_prob) < jump_prob

#                  if should_jump_batch.any():
#                      jump_len_probs = policy[should_jump_batch, :self.num_layers-i]
#                      sampled_jump_len = torch.multinomial(jump_len_probs, 1)[:, 0] + 1

#                      jump_length = sampled_jump_len.max().item()
#                      i_next = min(i + jump_length, self.num_layers)

#                      skip_weight_idx = min(jump_length - 1, len(self.jump_weights) - 1)
#                      skip_weight = self.jump_weights[skip_weight_idx]

#                      x = skip_weight * original_x + (1 - skip_weight) * working_memory.expand_as(x) + x * (1-skip_weight)
#                      jump_history.append(f"jump_{jump_length} S:{skip_weight.item():.2f} ->{i_next}")
#                      i = i_next
#                      continue

#             i += 1

#         mlp_input = last_processed_output
#         norm_mlp_input = self.mlp_ln(mlp_input)
#         mlp_output = self.mlp(norm_mlp_input)
#         mlp_gate_value = self.mlp_gate(norm_mlp_input)
#         final_output = mlp_input + mlp_gate_value * mlp_output

#         if 'memory_gate' in locals():
#              self.last_memory_gate_values = memory_gate.detach().clone()

#         return final_output

# class IntegratedAttention(BaseAttention):
#     """
#     Integrates multiple attention strategies:
#     - Local attention (sliding window or adaptive span via AdaptiveSpanAttention).
#     - Global attention (potentially with adaptive updates via AdaptiveUpdateAttention).
#     - Cross-attention capability.
#     - RL-based refinement (`Refiner`) of the local attention span.
#     - Iterative refinement (`_focus`) within local attention.
#     """
#     def __init__(self, dims: int, head: int, max_dist: int = 512,
#                  win_size: int = 256, max_span: int = 384, temp_scale: float = 0.01,
#                  dropout: float = 0.1,
#                  rl_states: int = 10000, rl_actions: int = 10, rl_alpha: float = 0.1,
#                  rl_gamma: float = 0.9, rl_epsilon: float = 0.1):
#         super().__init__(dims, head, max_dist, dropout=dropout)
#         self.max_span = max_span
#         self.sliding_window = win_size
#         self.temp_scale = temp_scale
#         self.sharpen = True

#         self.refiner = Refiner(
#             states=rl_states, actions=rl_actions, alpha=rl_alpha,
#             gamma=rl_gamma, epsilon=rl_epsilon)
#         self.span_pred = Predictor(dims=dims)

#         self.attn_local = AdaptiveSpanAttention(
#             dims=dims, head=head, max_dist=max_dist, sharpen=self.sharpen,
#             temp_scale=temp_scale, learnable_scale=False,
#             dropout=dropout)

#         self.attn_global = AdaptiveUpdateAttention(
#             dims=dims, head=head, max_dist=max_dist, dropout=dropout)

#         self.cross_attn = AttentionCombiner(dims=dims, head=head, dropout=dropout)

#         self.self_projection = Linear(in_features=2 * dims, out_features=dims)
#         self.global_cross_projection = Linear(in_features=dims, out_features=dims)

#         self.ln_local_in = LayerNorm(normalized_shape=dims)
#         self.ln_global_in = LayerNorm(normalized_shape=dims)
#         self.ln_cross_in = LayerNorm(normalized_shape=dims)

#         self.register_buffer("threshold", torch.tensor(1e-4), persistent=False)
#         self.register_buffer("s_factor", torch.tensor(0.1), persistent=False)

#     def forward(self, x: Tensor, xa: Optional[Tensor] = None,
#                 mask: Optional[Tensor] = None, kv_cache: Optional[Tensor] = None,
#                 is_causal: bool = True) -> Tensor:
#         """
#         Main forward pass distributing to cross or self-attention pathways.

#         Args:
#             x: Primary input tensor (Batch, SeqLen_Q, Dims).
#             xa: Context tensor for cross-attention (Batch, SeqLen_KV, Dims).
#             mask: Attention mask (padding or causal).
#             kv_cache: Key/Value cache for generation (specific usage depends on sub-modules).
#             is_causal: Flag for causal masking in self-attention.

#         Returns:
#             Output tensor (Batch, SeqLen_Q, Dims).
#         """
#         batch, ctx_q, _ = x.shape

#         if xa is not None:
#             q_norm = self.ln_cross_in(x)
#             k_cross = self.attn_global.key_module(xa)
#             v_cross = self.attn_global.value_module(xa)
#             q_cross = self.attn_global.query_module(q_norm)

#             cross_out = self.cross_attn(q=q_cross, k=k_cross, v=v_cross, mask=mask, is_causal=False)
#             return self.global_cross_projection(cross_out)

#         local_input = self.ln_local_in(x)
#         global_input = self.ln_global_in(x)

#         globe_out_raw = self.attn_global(
#              global_input,
#              xa=None,
#              mask=mask,
#              is_causal=is_causal
#         )
#         globe_out = self.global_cross_projection(globe_out_raw)

#         base_freq_scale = self.span_pred(globe_out)

#         state = self._extract_rl_state(local_input)
#         action = self.refiner.choose_action(state=state)
#         refinement_scale = self._action_to_scale(action=action)
#         final_span_scale = torch.clamp(base_freq_scale * refinement_scale.expand_as(base_freq_scale), min=0.0, max=1.0)

#         span_mean = final_span_scale.mean().item()
#         with torch.no_grad():
#             current_win_size = max(1, int(self.sliding_window * span_mean))
#             current_span_len = max(1, int(self.max_span * span_mean))
#         local_out_raw = self._slide_win_local(
#             x=local_input,
#             win_size=current_win_size,
#             span_len=current_span_len,
#             span_scale=final_span_scale,
#             mask=mask,
#             is_causal=is_causal
#         )
#         with torch.no_grad():
#              reward = self._calculate_rl_reward(output=local_out_raw)
#              next_state = self._extract_rl_state(local_out_raw)
#              self.refiner.update(state=state, action=action, reward=reward, next_state=next_state)

#         combined = torch.cat([local_out_raw, globe_out], dim=-1)
#         output = self.self_projection(combined)

#         return output

#     def _calculate_rl_reward(self, output: Tensor) -> float:
#         """Calculate quality metric (reward) for reinforcement learning."""
#         with torch.no_grad():
#             output_probs = torch.softmax(output, dim=-1)
#             safe_probs = torch.clamp(output_probs, min=1e-10)
#             entropy = -(safe_probs * torch.log(safe_probs)).sum(-1).mean()
#             coverage = (output.abs() > 0.01).float().mean()
#             reward = float(coverage - 0.1 * entropy)
#         return reward

#     def _extract_rl_state(self, x: Tensor) -> int:
#         """Extract discrete state features for RL agent from tensor."""
#         with torch.no_grad():
#             pooled = x.mean(dim=1)
#             mean_state = pooled[0].mean()
#             var_state = pooled[0].var(unbiased=False)
#             state_features = torch.stack([mean_state, var_state]).cpu().numpy()
#             state_id = self._discretize_state(state_features)
#         return state_id

#     def _discretize_state(self, state: np.ndarray) -> int:
#         """Convert continuous state numpy array to a discrete state ID."""
#         bins = np.linspace(-1, 1, num=10)
#         state_discrete = np.digitize(state, bins)
#         state_hash = sum(val * (10**i) for i, val in enumerate(state_discrete))
#         state_id = int(state_hash % self.refiner.states)
#         return state_id

#     def _action_to_scale(self, action: int) -> Tensor:
#         """Convert discrete RL action index to a continuous scale factor [0, 1]."""
#         span_value = action / (self.refiner.actions - 1)
#         scale_tensor = torch.tensor([span_value], device=self.span_pred.linear.weight.device, dtype=torch.float)
#         return scale_tensor

#     def _focus(self, query: Tensor, key: Tensor, value: Tensor,
#                span_scale: Tensor, mask: Optional[Tensor] = None,
#                is_causal: bool = False) -> Tuple[Tensor, Optional[Tensor]]:
#         """
#         Iterative attention refinement. Applies attention multiple times,
#         adding the output back to the query. Uses manual attention calculation.

#         Args:
#             query, key, value: Input tensors (B, SeqLen_Window, D).
#             span_scale: Scale factor (scalar or B, 1) influencing effective span.
#             mask: Attention mask for the window.
#             is_causal: Apply causal masking within the window.

#         Returns:
#             Tuple (refined_output (B, SeqLen_Window, D), attention_weights (optional, None here)).
#         """
#         max_iterations = 5
#         iteration = 0
#         prev_attn_out = torch.zeros_like(query)
#         attn_out = torch.zeros_like(query)
#         threshold = self.threshold.item()
#         s_factor = self.s_factor.item()

#         q_current = query

#         while iteration < max_iterations:
#             span_mean = span_scale.mean().item()
#             target_span_len = max(1, int(self.max_span * span_mean))
#             eff_span = min(target_span_len, self.max_dist, q_current.size(1), key.size(1))

#             if eff_span == 0: break

#             q_iter = q_current[:, :eff_span, :]
#             k_iter = key[:, :eff_span, :]
#             v_iter = value[:, :eff_span, :]

#             q_proj = self.attn_local.query_module(q_iter)
#             k_proj = self.attn_local.key_module(k_iter)
#             v_proj = self.attn_local.value_module(v_iter)

#             temperature = (1.0 + self.temp_scale * (1.0 - span_mean)
#                            if self.sharpen
#                            else 0.5 + self.temp_scale * span_mean)
#             temperature = max(temperature, 1e-3)

#             iter_mask = None
#             if mask is not None:
#                 if mask.dim() == 4: iter_mask = mask[:, :, :eff_span, :eff_span]
#                 elif mask.dim() == 2: iter_mask = mask[:eff_span, :eff_span]
#             attn_output_iter, _ = calculate_attention(
#                  q_proj, k_proj, v_proj,
#                  mask=iter_mask,
#                  temperature=temperature,
#                  use_sdpa=False,
#                  is_causal=is_causal,
#                  dropout_p=self.dropout
#             )

#             attn_out_span = self.attn_local._reshape_to_output(attn_output_iter)
#             projected_attn_out_span = self.attn_local.out_proj(attn_out_span)

#             current_iter_out = torch.zeros_like(q_current)
#             current_iter_out[:, :eff_span, :] = projected_attn_out_span

#             diff = torch.abs(current_iter_out - prev_attn_out).mean()
#             dynamic_threshold = threshold + s_factor * diff

#             if diff < dynamic_threshold and iteration > 0:
#                  attn_out = current_iter_out
#                  break

#             prev_attn_out = current_iter_out.clone()
#             q_current = q_current + current_iter_out
#             attn_out = current_iter_out

#             iteration += 1

#         return attn_out, None

#     @autocast('cuda', enabled=torch.cuda.is_available())
#     def _slide_win_local(self, x: Tensor, win_size: int, span_len: int,
#                          span_scale: Tensor, mask: Optional[Tensor] = None,
#                          is_causal: bool = False) -> Tensor:
#         """
#         Process input with sliding window attention, using `_focus` for each window.

#         Args:
#             x: Input tensor (Batch, SeqLen, Dims).
#             win_size: Size of the attention window for queries.
#             span_len: Max length of keys/values relative to query window start.
#             span_scale: Span scale tensor (Batch, 1 or scalar) passed to _focus.
#             mask: Full attention mask.
#             is_causal: Apply causal masking within windows.

#         Returns:
#             Output tensor (Batch, SeqLen, Dims).
#         """
#         batch, ctx, dims = x.size()
#         output = torch.zeros_like(x)

#         num_windows = (ctx + win_size - 1) // win_size

#         for i in range(num_windows):
#             q_start = i * win_size
#             q_end = min(q_start + win_size, ctx)
#             current_window_q_len = q_end - q_start
#             if current_window_q_len == 0: continue

#             kv_start = max(0, q_end - span_len)
#             kv_end = q_end
#             query_win = x[:, q_start:q_end, :]
#             key_win = x[:, kv_start:kv_end, :]
#             value_win = x[:, kv_start:kv_end, :]

#             window_mask = None
#             if mask is not None:
#                 if mask.dim() == 4:
#                     window_mask = mask[:, :, q_start:q_end, kv_start:kv_end]
#                 elif mask.dim() == 2:
#                     window_mask = mask[q_start:q_end, kv_start:kv_end]
#             attn_out_win, _ = self._focus(
#                 query=query_win,
#                 key=key_win,
#                 value=value_win,
#                 span_scale=span_scale,
#                 mask=window_mask,
#                 is_causal=is_causal
#             )

#             output[:, q_start:q_end, :] = attn_out_win

#         return output

# class CTCDecoder(nn.Module):
#     def __init__(self, input_dim: int, vocab_size: int, dims: int = 256, num_layers: int = 2, dropout: float = 0.1):
#         super().__init__()
#         self.input_dim = input_dim
#         self.vocab_size = vocab_size
#         self.dims = dims
        
#         self.projection = nn.Linear(input_dim, dims)
#         self.lstm = nn.LSTM(dims, dims, num_layers, dropout=dropout if num_layers > 1 else 0, batch_first=True, bidirectional=True)
#         self.output = nn.Linear(dims * 2, vocab_size + 1)  # +1 for CTC blank token
#         self.dropout = nn.Dropout(dropout)
        
#     def forward(self, x: Tensor) -> Tensor:
#         x = self.projection(x)  # (batch, seq_len, dims)
#         x = self.dropout(x)
#         x, _ = self.lstm(x)  # (batch, seq_len, dims * 2)
#         x = self.dropout(x)
#         logits = self.output(x)  # (batch, seq_len, vocab_size + 1)
#         return logits

# class CTCWrapper(nn.Module):
#     def __init__(self, model: Model, vocab_size: int, dims: int = 256, num_layers: int = 2):
#         super().__init__()
#         self.model = model
#         self.ctc_decoder = CTCDecoder(
#             input_dim=model.param.dims,
#             vocab_size=vocab_size,
#             dims=dims,
#             num_layers=num_layers
#         )
        
#     def forward(self, input_ids=None, pitch=None, labels=None, input_lengths=None, label_lengths=None):
#         outputs = self.model(input_ids=input_ids, pitch=pitch)
#         x = outputs["logits"]  # (batch, seq_len, vocab_size)
#         ctc_logits = self.ctc_decoder(x)  # (batch, seq_len, vocab_size + 1)
#         loss = None
#         if labels is not None and input_lengths is not None and label_lengths is not None:
#             log_probs = torch.log_softmax(ctc_logits, dim=-1)
#             log_probs = log_probs.transpose(0, 1)
            
#             loss = torch.nn.functional.ctc_loss(
#                 log_probs,
#                 labels,
#                 input_lengths,
#                 label_lengths,
#                 blank=0,
#                 reduction='mean'
#             )
        
#         return {
#             "logits": ctc_logits,
#             "loss": loss,
#             "out": x
#         }
    
#     def decode(self, logits: Tensor, input_lengths: Optional[Tensor] = None) -> List[List[int]]:
#         probs = torch.softmax(logits, dim=-1)  # (batch, seq_len, vocab_size + 1)
#         predictions = torch.argmax(probs, dim=-1)  # (batch, seq_len)
        
#         decoded_sequences = []
#         for i, pred in enumerate(predictions):
#             seq = []
#             prev_token = None
#             for j, token in enumerate(pred):
#                 if input_lengths is not None and j >= input_lengths[i]:
#                     break
#                 if token != 0 and token != prev_token:
#                     seq.append(token.item())
#                 prev_token = token
#             decoded_sequences.append(seq)
#         return decoded_sequences

#     # ctc_model = CTCWrapper(model, vocab_size=40000, dims=256, num_layers=2)

#     # outputs = ctc_model(
#     #     input_ids=input_ids,
#     #     pitch=pitch,
#     #     labels=labels,
#     #     input_lengths=input_lengths,  # Length of each audio sequence
#     #     label_lengths=label_lengths   # Length of each text sequence
#     # )

#     # loss = outputs["loss"]

#     # outputs = ctc_model(input_ids=input_ids, pitch=pitch)
#     # logits = outputs["logits"]

#     # # Decode to text
#     # decoded_sequences = ctc_model.decode(logits, input_lengths=input_lengths)
#     # ctc_model = CTCWrapper(model, vocab_size=param.vocab, dims=256, num_layers=2).to('cuda')
    
#     # print(f"CTC model parameters: {sum(p.numel() for p in ctc_model.parameters() if p.requires_grad):,}")

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

# def compute_metricsB(pred, tokenizer):
#     pred_ids = pred["predictions"]
#     label_ids = pred["label_ids"]
#     if isinstance(pred_ids, tuple):
#         pred_ids = pred_ids[0]
#     else:
#         pred_ids = pred_ids
#     if pred_ids.ndim == 3:
#         pred_ids = np.argmax(pred_ids, axis=-1)
#     label_ids[label_ids == -100] = tokenizer.pad_token_id
#     pred_str = tokenizer.batch_decode(pred_ids, skip_special_tokens=True)
#     label_str = tokenizer.batch_decode(label_ids, skip_special_tokens=True)
#     metrics = evaluate.load(path="wer")
#     wer = metrics.compute(predictions=pred_str, references=label_str)
#     return {"wer": wer}

# def train_and_evaluate(
#     model,
#     tokenizer,
#     train_loader,
#     eval_loader,
#     optimizer,
#     scheduler,
#     loss_fn,
#     max_steps=10000,
#     device="cuda",
#     accumulation_steps=1,
#     clear_cache=True,
#     log_interval=10,
#     eval_interval=100,
#     save_interval=1000,
#     checkpoint_dir="checkpoint_dir",
#     log_dir="log_dir",
# ):
#     model.to(device)
#     global_step = 0
#     scaler = torch.GradScaler()
#     writer = SummaryWriter(log_dir=log_dir)
#     train_iterator = iter(train_loader)
#     total_loss = 0
#     step_in_report = 0
#     dataset_epochs = 0

#     progress_bar = tqdm(
#         total=max_steps, desc="Training Progress", leave=True, colour="green"
#     )

#     model.train()
#     optimizer.zero_grad()

#     while global_step < max_steps:
#         try:
#             batch = next(train_iterator)
#         except StopIteration:
#             train_iterator = iter(train_loader)
#             batch = next(train_iterator)
#             dataset_epochs += 1
#             print(f"Starting dataset epoch {dataset_epochs}")

#             if step_in_report > 0:
#                 avg_loss = total_loss / step_in_report
#                 logging.info(
#                     f"Dataset iteration complete - Steps: {global_step}, Avg Loss: {avg_loss:.4f}"
#                 )
#                 total_loss = 0
#                 step_in_report = 0

#         start_time = time.time()

#         input_features = batch["input_features"].to(device)
#         input_ids = batch["input_ids"].to(device)
#         labels = batch["labels"].long().to(device)

#         with torch.autocast(device_type="cuda"):
#             input_features_encoded = model.encoder(input_features)
#             decoder_output = model.decoder(input_ids, input_features_encoded)
#             logits = decoder_output.view(-1, decoder_output.size(-1))
#             active_logits = logits.view(-1, decoder_output.size(-1))
#             active_labels = labels.view(-1)
#             active_mask = active_labels != -100
#             active_logits = active_logits[active_mask]
#             active_labels = active_labels[active_mask]
#             loss = loss_fn(active_logits, active_labels)
#             # model.adjust_freq(loss=loss.item())
#         total_loss += loss.item()
#         loss = loss / accumulation_steps

#         scaler.scale(loss).backward()

#         if (global_step + 1) % accumulation_steps == 0:
#             scaler.unscale_(optimizer)
#             torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
#             scaler.step(optimizer)
#             scaler.update()
#             optimizer.zero_grad()

#             if clear_cache:
#                 torch.cuda.empty_cache()

#         end_time = time.time()
#         samples_per_sec = len(batch["input_features"]) / (end_time - start_time)

#         if global_step % log_interval == 0:
#             writer.add_scalar(
#                 tag="Loss/train",
#                 scalar_value=total_loss / (global_step + 1),
#                 global_step=global_step,
#             )
#             lr = scheduler.get_last_lr()[0]
#             writer.add_scalar(
#                 tag="LearningRate", scalar_value=lr, global_step=global_step
#             )
#             writer.add_scalar(
#                 tag="SamplesPerSec",
#                 scalar_value=samples_per_sec,
#                 global_step=global_step,
#             )

#         if global_step % eval_interval == 0:
#             model.eval()
#             eval_start_time = time.time()
#             eval_loss = 0
#             all_predictions = []
#             all_labels = []
#             batch_count = 0
#             total_samples = 0

#             with torch.no_grad():
#                 for eval_batch in eval_loader:
#                     # for eval_batch in tqdm(eval_loader, desc=f"Evaluating (Step {global_step})", leave=True, colour='green'):
#                     input_features = eval_batch["input_features"].to(device)
#                     input_ids = eval_batch["input_ids"].to(device)
#                     labels = eval_batch["labels"].long().to(device)

#                     batch = input_features.size(0)
#                     total_samples += batch

#                     input_features_encoded = model.encoder(input_features)
#                     decoder_output = model.decoder(input_ids, input_features_encoded)
#                     logits = decoder_output.view(-1, decoder_output.size(-1))
#                     loss = loss_fn(logits, labels.view(-1))
#                     eval_loss += loss.item()
#                     all_predictions.extend(
#                         torch.argmax(decoder_output, dim=-1).cpu().numpy().tolist()
#                     )
#                     all_labels.extend(labels.cpu().numpy().tolist())
#                     batch_count += 1

#             eval_time = time.time() - eval_start_time
#             loss_avg = eval_loss / batch_count if batch_count > 0 else 0
#             predictions = {
#                 "predictions": np.array(all_predictions, dtype=object),
#                 "label_ids": np.array(all_labels, dtype=object),
#             }
#             metrics = compute_metrics(pred=predictions, tokenizer=tokenizer)

#             writer.add_scalar("Loss/eval", loss_avg, global_step)
#             writer.add_scalar("WER", metrics["wer"], global_step)
#             writer.add_scalar("EvalSamples", total_samples, global_step)
#             writer.add_scalar("EvalTimeSeconds", eval_time, global_step)
#             lr = scheduler.get_last_lr()[0]

#             print(
#                 f"• STEP:{global_step} • samp:{samples_per_sec:.1f} • WER:{metrics['wer']:.2f}% • Loss:{loss_avg:.4f} • LR:{lr:.8f}"
#             )

#             logging.info(
#                 f"EVALUATION STEP {global_step} - WER: {metrics['wer']:.2f}%, Loss: {loss_avg:.4f}, LR: {lr:.8f}"
#             )
#             # scheduler.step()
#             model.train()

#         if global_step % save_interval == 0:
#             checkpoint_path = os.path.join(
#                 checkpoint_dir, f"checkpoint_step_{global_step}.pt"
#             )
#             torch.save(model.state_dict(), checkpoint_path)
#             # print(f"Model saved at step {global_step} to {checkpoint_path}")
#             logging.info(f"Model saved at step {global_step} to {checkpoint_path}")

#         lr = scheduler.get_last_lr()[0]
#         scheduler.step()
#         global_step += 1
#         step_in_report += 1

#         avg_loss = total_loss / (global_step + 1)
#         postfix_dict = {
#             "loss": f"{avg_loss:.4f}",
#             "lr": f"{lr:.6f}",
#             "samp": f"{samples_per_sec:.1f}",
#         }
#         progress_bar.set_postfix(postfix_dict, refresh=True)
#         progress_bar.update(1)

#     final_model_path = os.path.join(checkpoint_dir, "final_model.pt")
#     torch.save(model.state_dict(), final_model_path)
#     print(
#         f"Training completed after {global_step} steps. Final model saved to {final_model_path}"
#     )
#     writer.close()
#     progress_bar.close()

# def mainB():

#     checkpoint_dir = "./output/checkpoints"
#     os.makedirs(checkpoint_dir, exist_ok=True)
#     log_dir = os.path.join("./output/logs", datetime.now().strftime(format="%m-%d_%H"))
#     os.makedirs(name=log_dir, exist_ok=True)

#     logging.basicConfig(
#         filename=os.path.join(log_dir, "training.log"),
#         filemode="w",
#         format="%(asctime)s - %(levelname)s - %(message)s",
#         level=logging.INFO,
#     )

#     token = ""
#     dataset = IterableDatasetDict()
#     dataset["train"] = load_dataset(
#         path="google/fleurs",
#         name="en_us",
#         split="train",
#         streaming=True,
#         token=token,
#         trust_remote_code=True,
#     ).select_columns(column_names=["audio", "transcription"])

#     dataset["test"] = load_dataset(
#         path="google/fleurs",
#         name="en_us",
#         split="test",
#         streaming=True,
#         token=token,
#         trust_remote_code=True,
#     ).select_columns(column_names=["audio", "transcription"])

#     debug = None

#     param = Dimensions(
#         mels=128,
#         audio_ctx=1500,
#         audio_head=4,
#         encoder_idx=4,
#         audio_dims=512,
#         vocab=51865,
#         text_ctx=512,
#         text_head=4,
#         decoder_idx=4,
#         text_dims=512,
#         decoder_start_token_id=0,
#         pad_token_id=0,
#         eos_token_id=0,
#         act="gelu",
#     )

#     model = model

#     Collator = DataCollatorB(
#         tokenizer=tokenizer,
#         audio_ctx=param.audio_ctx,
#         text_ctx=param.text_ctx,
#         mels=param.mels,
#     )

#     train_dataloader = DataLoader(
#         dataset=dataset["train"], batch_size=1, collate_fn=Collator, num_workers=0
#     )

#     eval_dataloader = DataLoader(
#         dataset=dataset["test"], batch_size=1, collate_fn=Collator, num_workers=0
#     )

#     optimizer = torch.optim.AdamW(
#         model.parameters(), lr=5e-4, weight_decay=0.01, eps=1e-6, betas=(0.9, 0.98)
#     )
#     scheduler = torch.optim.lr_scheduler.LinearLR(
#         optimizer, start_factor=0.25, total_iters=10000, last_epoch=-1
#     )

#     loss_fn = torch.nn.CrossEntropyLoss(ignore_index=-100)

#     train_and_evaluate(
#         model=model,
#         tokenizer=tokenizer,
#         train_loader=train_dataloader,
#         eval_loader=eval_dataloader,
#         optimizer=optimizer,
#         scheduler=scheduler,
#         loss_fn=loss_fn,
#         max_steps=10000,
#         device="cuda",
#         accumulation_steps=1,
#         clear_cache=False,
#         log_interval=10,
#         eval_interval=500,
#         save_interval=10000,
#         checkpoint_dir=checkpoint_dir,
#         log_dir=log_dir,
#     )

# def train_and_evaluate(
#     model, tokenizer, train_loader, eval_loader, optimizer, scheduler, loss_fn,
#     max_steps=10000, device='cuda', accumulation_steps=1, clear_cache=True,
#     log_interval=10, eval_interval=100, save_interval=1000,
#     checkpoint_dir="checkpoint_dir", log_dir="log_dir"
# ):
#     model.to(device)
#     global_step = 0
#     scaler = torch.GradScaler()
#     writer = SummaryWriter(log_dir=log_dir)
#     train_iterator = iter(train_loader)
#     total_loss = 0
#     step_in_report = 0
#     dataset_epochs = 0

#     progress_bar = tqdm(total=max_steps, desc="Training Progress", leave=True, colour='green')

#     model.train()
#     optimizer.zero_grad()

#     while global_step < max_steps:
#         try:
#             batch = next(train_iterator)
#         except StopIteration:
#             train_iterator = iter(train_loader)
#             batch = next(train_iterator)
#             dataset_epochs += 1
#             print(f"Starting dataset epoch {dataset_epochs}")

#             if step_in_report > 0:
#                 avg_loss = total_loss / step_in_report
#                 logging.info(f"Dataset iteration complete - Steps: {global_step}, Avg Loss: {avg_loss:.4f}")
#                 total_loss = 0
#                 step_in_report = 0

#         start_time = time.time()

#         batch = {k: v.to(device) if isinstance(v, torch.Tensor) else v for k, v in batch.items()}

#         with torch.autocast(device_type="cuda"):
#             output = model(**batch) if hasattr(model, '__call__') else model.forward(**batch)
#             logits = output["logits"] if isinstance(output, dict) and "logits" in output else output
#             labels = batch["labels"]
#             active_logits = logits.view(-1, logits.size(-1))
#             active_labels = labels.view(-1)
#             active_mask = active_labels != 0
#             active_logits = active_logits[active_mask]
#             active_labels = active_labels[active_mask]
#             loss = loss_fn(active_logits, active_labels)
#         total_loss += loss.item()
#         loss = loss / accumulation_steps

#         scaler.scale(loss).backward()

#         if (global_step + 1) % accumulation_steps == 0:
#             scaler.unscale_(optimizer)
#             torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
#             scaler.step(optimizer)
#             scaler.update()
#             optimizer.zero_grad()
#             if clear_cache:
#                 torch.cuda.empty_cache()

#         end_time = time.time()
#         samples_per_sec = batch["spectrogram"].size(0) / (end_time - start_time)

#         if global_step % log_interval == 0:
#             writer.add_scalar(tag='Loss/train', scalar_value=total_loss / (global_step + 1), global_step=global_step)
#             lr = scheduler.get_last_lr()[0]
#             writer.add_scalar(tag='LearningRate', scalar_value=lr, global_step=global_step)
#             writer.add_scalar(tag='SamplesPerSec', scalar_value=samples_per_sec, global_step=global_step)

#         if global_step % eval_interval == 0:
#             model.eval()
#             eval_start_time = time.time()
#             eval_loss = 0
#             all_predictions = []
#             all_labels = []
#             batch_count = 0
#             total_samples = 0

#             with torch.no_grad():
#                 for eval_batch in eval_loader:
#                     eval_batch = {k: v.to(device) if isinstance(v, torch.Tensor) else v for k, v in eval_batch.items()}
#                     output = model(**eval_batch) if hasattr(model, '__call__') else model.forward(**eval_batch)
#                     logits = output["logits"] if isinstance(output, dict) and "logits" in output else output
#                     labels = eval_batch["labels"]
#                     batch_size = logits.size(0)
#                     total_samples += batch_size
#                     loss = loss_fn(logits.view(-1, logits.size(-1)), labels.view(-1))
#                     eval_loss += loss.item()
#                     all_predictions.extend(torch.argmax(logits, dim=-1).cpu().numpy().tolist())
#                     all_labels.extend(labels.cpu().numpy().tolist())
#                     batch_count += 1

#             eval_time = time.time() - eval_start_time
#             loss_avg = eval_loss / batch_count if batch_count > 0 else 0
#             predictions = {"predictions": np.array(all_predictions, dtype=object), "label_ids": np.array(all_labels, dtype=object)}
#             metrics = compute_metrics(pred=predictions, tokenizer=tokenizer)

#             writer.add_scalar('Loss/eval', loss_avg, global_step)
#             writer.add_scalar('WER', metrics['wer'], global_step)
#             writer.add_scalar('EvalSamples', total_samples, global_step)
#             writer.add_scalar('EvalTimeSeconds', eval_time, global_step)

#             lr = scheduler.get_last_lr()[0]
#             print(f"• STEP:{global_step} • samp:{samples_per_sec:.1f} • WER:{metrics['wer']:.2f}% • Loss:{loss_avg:.4f} • LR:{lr:.8f}")
#             logging.info(f"EVALUATION STEP {global_step} - WER: {metrics['wer']:.2f}%, Loss: {loss_avg:.4f}, LR: {lr:.8f}")
#             model.train()

#         if global_step % save_interval == 0:
#             checkpoint_path = os.path.join(checkpoint_dir, f'checkpoint_step_{global_step}.pt')
#             torch.save(model.state_dict(), checkpoint_path)
#             logging.info(f"Model saved at step {global_step} to {checkpoint_path}")

#         lr = scheduler.get_last_lr()[0]
#         scheduler.step()
#         global_step += 1
#         step_in_report += 1

#         avg_loss = total_loss / (global_step + 1)
#         postfix_dict = {
#             'loss': f'{avg_loss:.4f}',
#             'lr': f'{lr:.6f}',
#             'samp': f'{samples_per_sec:.1f}'
#         }
#         progress_bar.set_postfix(postfix_dict, refresh=True)
#         progress_bar.update(1)

#     final_model_path = os.path.join(checkpoint_dir, 'final_model.pt')
#     torch.save(model.state_dict(), final_model_path)
#     print(f"Training completed after {global_step} steps. Final model saved to {final_model_path}")
#     writer.close()
#     progress_bar.close()

# def get_optimizer(model, lr=5e-4, weight_decay=0.01):
#     return torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=weight_decay, eps=1e-6, betas=(0.9, 0.98))

# def get_scheduler(optimizer, total_steps=10000):
#     return torch.optim.lr_scheduler.LinearLR(optimizer, start_factor=0.25, total_iters=total_steps, last_epoch=-1)

# def get_loss_fn():
#     return torch.nn.CrossEntropyLoss(ignore_index=0)

# def mainc():
#     token = ""
#     log_dir = os.path.join('./output/logs', datetime.now().strftime('%m-%d_%H_%M_%S'))
#     os.makedirs(log_dir, exist_ok=True)
#     tokenizer = setup_tokenizer(token)

#     param = Dimensions(
#         mels=128, aud_ctx=1500, aud_head=4, aud_dims=512, aud_idx=4,
#         vocab=40000, text_ctx=512, text_head=4, text_dims=512, text_idx=4,
#         act="swish", debug={}, cross_attn=True, features=["spectrogram"]
#     )

#     dataset_config = {
#         "spectrogram": True, "waveforms": False, "pitch": False, "downsamples": False,
#         "frequency": True, "hilbert": False, "hop_length": 128, "fmin": 150, "fmax": 2000,
#         "n_mels": 128, "n_fft": 1024, "sampling_rate": 16000, "pad_mode": "constant",
#         "center": True, "power": 2.0, "window_fn": torch.hann_window, "mel_scale": "htk",
#         "norm": None, "normalized": False
#     }

#     model = create_model(param)
#     train_dataset, test_dataset = prepare_datasets(
#         tokenizer=tokenizer, token=token, sanity_check=False, dataset_config=dataset_config
#     )

#     collator = DataCollator(tokenizer=tokenizer)
#     train_loader = DataLoader(train_dataset, batch_size=1, collate_fn=collator, num_workers=0)
#     eval_loader = DataLoader(test_dataset, batch_size=1, collate_fn=collator, num_workers=0)

#     optimizer = get_optimizer(model)
#     scheduler = get_scheduler(optimizer)
#     loss_fn = get_loss_fn()

#     train_and_evaluate(
#         model=model,
#         tokenizer=tokenizer,
#         train_loader=train_loader,
#         eval_loader=eval_loader,
#         optimizer=optimizer,
#         scheduler=scheduler,
#         loss_fn=loss_fn,
#         max_steps=10000,
#         device='cuda',
#         accumulation_steps=1,
#         clear_cache=False,
#         log_interval=10,
#         eval_interval=500,
#         save_interval=10000,
#         checkpoint_dir="./checkpoints",
#         log_dir=log_dir
#     )

# class attention(nn.Module):
#     def __init__(self, dims: int, head: int):
#         super(attention, self).__init__()
#         self.dims = dims
#         self.head = head
#         self.head_dim = dims // head
#         self.q = nn.Linear(dims, dims)
#         self.k = nn.Linear(dims, dims, bias=False)
#         self.v = nn.Linear(dims, dims)
#         self.o = nn.Linear(dims, dims)

#         self.lna = nn.LayerNorm(dims, bias = False)
#         self.lnb = nn.LayerNorm(dims, bias = False)      
#         self.lnc = nn.LayerNorm(self.head_dim, bias = False)
#         self.lnd = nn.LayerNorm(self.head_dim, bias = False)     

#     def _forward(self, x: Tensor, xa = None, mask = None):
#         q = self.q(self.lna(x))
#         k = self.k(self.lnb(x if xa is None else xa))
#         v = self.v(self.lnb(x if xa is None else xa))
#         query = q.view(*q.shape[:2], self.head, -1).permute(0, 2, 1, 3)
#         key = k.view(*k.shape[:2], self.head, -1).permute(0, 2, 1, 3)
#         value = v.view(*v.shape[:2], self.head, -1).permute(0, 2, 1, 3)

#         max_iterations = 5
#         iteration = 0
#         prev_attn_out = torch.zeros_like(query)
#         attn_out = torch.zeros_like(query)
#         threshold = self.threshold.item()
#         s_factor = self.s_factor.item()

#         q_current = query

#         while iteration < max_iterations:
 
#             eff_span = min(x.shape[1], xa.shape[1], q_current.size(1), key.size(1))

#             if eff_span == 0: break

#             q_iter = q_current[:, :eff_span, :]
#             k_iter = key[:, :eff_span, :]
#             v_iter = value[:, :eff_span, :]

#             q_proj = self.attn_local.query_module(q_iter)
#             k_proj = self.attn_local.key_module(k_iter)
#             v_proj = self.attn_local.value_module(v_iter)

#             temperature = (1.0 + self.temp_scale * (1.0 - xa.mean())
#                            if self.sharpen
#                            else 0.5 + self.temp_scale * xa.mean())
#             temperature = max(temperature, 1e-3)

#             iter_mask = None
#             if mask is not None:
#                 if mask.dim() == 4: iter_mask = mask[:, :, :eff_span, :eff_span]
#                 elif mask.dim() == 2: iter_mask = mask[:eff_span, :eff_span]

#             attn_output_iter, _ = calculate_attention(
#                  q_proj, k_proj, v_proj,
#                  mask=iter_mask,
#                  temperature=temperature,
#                  use_sdpa=False,
#                  dropout_p=self.dropout
#             )

#             attn_out_span = self.attn_local._reshape_to_output(attn_output_iter)
#             projected_attn_out_span = self.attn_local.out_proj(attn_out_span)

#             current_iter_out = torch.zeros_like(q_current)
#             current_iter_out[:, :eff_span, :] = projected_attn_out_span

#             diff = torch.abs(current_iter_out - prev_attn_out).mean()
#             dynamic_threshold = threshold + s_factor * diff

#             if diff < dynamic_threshold and iteration > 0:
#                  attn_out = current_iter_out
#                  break

#             prev_attn_out = current_iter_out.clone()
#             q_current = q_current + current_iter_out
#             attn_out = current_iter_out

#             iteration += 1

#         return attn_out, None

#     def _slide_win_local(self, x: Tensor, win_size: int, span_len: int,
#                          span_scale: Tensor, mask: Optional[Tensor] = None,
#                          is_causal: bool = False) -> Tensor:
#         """
#         Process input with sliding window attention, using `_focus` for each window.

#         Args:
#             x: Input tensor (Batch, SeqLen, Dims).
#             win_size: Size of the attention window for queries.
#             span_len: Max length of keys/values relative to query window start.
#             span_scale: Span scale tensor (Batch, 1 or scalar) passed to _focus.
#             mask: Full attention mask.
#             is_causal: Apply causal masking within windows.

#         Returns:
#             Output tensor (Batch, SeqLen, Dims).
#         """
#         batch, ctx, dims = x.size()
#         output = torch.zeros_like(x)

#         num_windows = (ctx + win_size - 1) // win_size

#         for i in range(num_windows):
#             q_start = i * win_size
#             q_end = min(q_start + win_size, ctx)
#             current_window_q_len = q_end - q_start
#             if current_window_q_len == 0: continue

#             kv_start = max(0, q_end - span_len)
#             kv_end = q_end
#             query_win = x[:, q_start:q_end, :]
#             key_win = x[:, kv_start:kv_end, :]
#             value_win = x[:, kv_start:kv_end, :]

#             window_mask = None
#             if mask is not None:
#                 if mask.dim() == 4:
#                     window_mask = mask[:, :, q_start:q_end, kv_start:kv_end]
#                 elif mask.dim() == 2:
#                     window_mask = mask[q_start:q_end, kv_start:kv_end]

#             attn_out_win, _ = self._focus(
#                 query=query_win,
#                 key=key_win,
#                 value=value_win,
#                 span_scale=span_scale,
#                 mask=window_mask,
#                 is_causal=is_causal
#             )

#             output[:, q_start:q_end, :] = attn_out_win

#         return output