Delete echoutils.py

echoutils.py

@@ -1,1593 +0,0 @@
-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
-from datasets import 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 Any, List, Dict, 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
-
-# def shape(tensor: torch.Tensor, head: int, head_dim: int, batch: int, ctx: int): 
-#     return tensor.view(batch, ctx, head, head_dim).transpose(1, 2).contiguous()
-
-# def reshape_to_output(attn_output, head: int, head_dim: int, batch: int, ctx: int, dims: int):
-#     return attn_output.permute(0, 2, 1, 3).reshape(batch, ctx, dims).contiguous()
-
-def shape(self, tensor: torch.Tensor, ctx: int, batch: int):
-    return tensor.view(batch, ctx, self.head, self.head_dim).transpose(1, 2).contiguous()
-
-def reshape_to_output(self, attn_output, batch, ctx):
-    return attn_output.permute(0, 2, 1, 3).reshape(batch, ctx, self.dims).contiguous()
-
-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.unsqueeze(0).unsqueeze(0).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 | (~padding_mask)
-    return mask
-
-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 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 qkv_init(dims: int, head: int):
-    head_dim = dims // head
-    scale = head_dim ** -0.5
-    q = nn.Linear(dims, dims)
-    k = nn.Linear(dims, dims, bias=False)
-    v = nn.Linear(dims, dims)
-    o = nn.Linear(dims, dims)
-    return q, k, v, o, scale
-
-def create_qkv(q, k, v, x, xa=None, head=8):
-    head_dim = q.out_features // head
-    scale = head_dim ** -0.5
-    q = q(x) * scale
-    k = k(xa if xa is not None else x) * scale
-    v = v(xa if xa is not None else x)
-    batch, ctx, _ = q.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, temperature=1.0, is_causal=True):
-    # q, k, v = create_qkv(q, k, v, dims, head)
-
-    batch, head, ctx, dims = q.shape
-    attn_mask = None
-    if mask is not None:
-        if mask.dim() <= 3:
-            attn_mask = create_attention_mask(
-                batch_size=batch,
-                ctx=ctx,
-                is_causal=is_causal,
-                padding_mask=mask if mask.dim() > 1 else None,
-                device=device)
-        else:
-            attn_mask = mask
-    scaled_q = q
-    if temperature != 1.0 and temperature > 0:
-        scaled_q = q * (1.0 / temperature)**.5
-    a = scaled_dot_product_attention(scaled_q, k, v, attn_mask=attn_mask, is_causal=is_causal if attn_mask is None else False)
-    out = a.permute(0, 2, 1, 3).flatten(start_dim=2)
-    return out, None
-
-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 rotary(nn.Module):
-#     def __init__(self, dims, head, max_ctx=1500, radii=False, debug: List[str] = [], use_pbias=False, axial=False, spec_shape=None):
-
-#         super(rotary, self).__init__()
-#         self.use_pbias = use_pbias
-#         self.dims = dims
-#         self.head = head
-#         self.head_dim = dims // head
-#         self.radii = radii
-#         self.debug = debug
-#         self.counter = 0
-#         self.last_theta = None
-#         self.axial = axial
-
-#         self.bias = nn.Parameter(torch.zeros(max_ctx, dims // 2), requires_grad=True if use_pbias else False)
-#         theta = (torch.tensor(10000, device=device, dtype=dtype))
-#         self.theta = nn.Parameter(theta, requires_grad=True)    
-#         self.theta_values = []
-
-#         if axial and spec_shape is not None:
-#             time_frames, freq_bins = spec_shape
-#             self.time_frames = time_frames
-#             self.freq_bins = freq_bins
-            
-#             time_theta = 50.0
-#             time_freqs = 1.0 / (time_theta ** (torch.arange(0, dims, 4)[:(dims // 4)].float() / dims))
-#             self.register_buffer('time_freqs', time_freqs)
-            
-#             freq_theta = 100.0
-#             freq_freqs = 1.0 / (freq_theta ** (torch.arange(0, dims, 4)[:(dims // 4)].float() / dims))
-#             self.register_buffer('freq_freqs', freq_freqs)
-
-#     def pitch_bias(self, f0):
-#         if f0 is None:
-#             return None
-#         f0_flat = f0.squeeze().float()
-#         f0_norm = (f0_flat - f0_flat.mean()) / (f0_flat.std() + 1e-8)
-#         f0_sim = torch.exp(-torch.cdist(f0_norm.unsqueeze(1), 
-#                                     f0_norm.unsqueeze(1)))
-#         return f0_sim.unsqueeze(0).unsqueeze(0)
-
-#     def theta_freqs(self, theta):
-#         if theta.dim() == 0:
-#             theta = theta.unsqueeze(0)
-#         freq = (theta.unsqueeze(-1) / 220.0) * 700 * (
-#             torch.pow(10, torch.linspace(0, 2595 * torch.log10(torch.tensor(1 + 8000/700)), 
-#                     self.head_dim // 2, device=theta.device, dtype=theta.dtype) / 2595) - 1) / 1000
-#         return freq
-
-#     def _apply_radii(self, freqs, f0, ctx):
-#         if self.radii and f0 is not None:
-#             radius = f0.to(device, dtype)
-#             L = radius.shape[0]
-#             if L != ctx:
-#                 feature = L / ctx
-#                 idx = torch.arange(ctx, device=f0.device)
-#                 idx = (idx * feature).long().clamp(0, L - 1)
-#                 radius = radius[idx]
-#                 return torch.polar(radius.unsqueeze(-1), freqs), radius
-#             else:
-#                 return torch.polar(radius.unsqueeze(-1), freqs), radius
-#         else:
-#             return torch.polar(torch.ones_like(freqs), freqs), None
-
-#     def check_f0(self, f0, f0t, ctx):
-#         if f0 is not None and f0.shape[1] == ctx:
-#             return f0
-#         elif f0t is not None and f0t.shape[1] == ctx:
-#             return f0t
-#         else:
-#             return None         
-
-#     def axial_freqs(self, ctx):
-#         if not self.axial:
-#             return None
-#         time_frames = self.time_frames
-#         freq_bins = self.freq_bins
-
-#         t = torch.arange(ctx, device=device, dtype=dtype)
-#         t_x = (t % time_frames).float()
-#         t_y = torch.div(t, time_frames, rounding_mode='floor').float()
-#         freqs_x = torch.outer(t_x, self.time_freqs)
-#         freqs_y = torch.outer(t_y, self.freq_freqs)
-#         freqs_cis_x = torch.polar(torch.ones_like(freqs_x), freqs_x)
-#         freqs_cis_y = torch.polar(torch.ones_like(freqs_y), freqs_y)
-#         return torch.cat([freqs_cis_x, freqs_cis_y], dim=-1)
-
-#     def forward(self, x=None, feats=None, feature=None, layer=None) -> Tensor:
-#         ctx=x
-#         f0 = feats.get("f0") if feats is not None else None 
-#         f0t = feats.get("f0t") if feats is not None else None 
-
-#         f0 = self.check_f0(f0, f0t, ctx)
-#         if f0 is not None:
-#             # if f0.dim() == 2:
-#             #     f0 = f0.squeeze(0) 
-#             theta = f0 + self.theta  
-#         else:
-#             theta = self.theta 
-#         freqs = self.theta_freqs(theta)
-#         t = torch.arange(ctx, device=device, dtype=dtype) # type: ignore
-#         freqs = t[:, None] * freqs
-#         freqs, radius = self._apply_radii(freqs, f0, ctx)
-
-#         if self.axial and feature == "spectrogram":
-#             freqs_2d = self.axial_freqs(ctx)
-#             if freqs_2d is not None:
-#                 return freqs_2d.unsqueeze(0)
-
-#         if "radius" in self.debug and self.counter == 10:
-#             print(f"  [{layer}] [Radius] {radius.shape if radius is not None else None} {radius.mean() if radius is not None else None} [Theta] {theta.mean() if theta is not None else None} [f0] {f0.shape if f0 is not None else None} [Freqs] {freqs.shape} {freqs.mean():.2f} [ctx] {ctx}")
-#         self.counter += 1
-#         return freqs.unsqueeze(0)
-
-#     @staticmethod
-#     def split(X: Tensor):
-#         half_dim = X.shape[-1] // 2
-#         return X[..., :half_dim], X[..., half_dim:]
-
-#     @staticmethod
-#     def apply_rotary(x, freqs):
-#         x1 = x[..., :freqs.shape[-1]*2]
-#         x2 = x[..., freqs.shape[-1]*2:]
-#         orig_shape = x1.shape
-#         if x1.ndim == 2:
-#             x1 = x1.unsqueeze(0)
-#         x1 = x1.float().reshape(*x1.shape[:-1], -1, 2).contiguous()
-#         x1 = torch.view_as_complex(x1) * freqs
-#         x1 = torch.view_as_real(x1).flatten(-2)
-#         x1 = x1.view(orig_shape)
-#         return torch.cat([x1.type_as(x), x2], dim=-1)
-
-
-# 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)  
-#         else:
-#             self.rope = None 
-#             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 plot_waveform(x=None, w=None, p=None, per=None, sample_idx=0, sr=16000, hop_length=160, 
-                                 title="", markers=None, marker_labels=None, 
-                                 show_voiced_regions=True, show_energy=False):
-    num_plots = sum([x is not None, w is not None, p is not None, per is not None])
-    if num_plots == 0:
-        raise ValueError("No data to plot. Please provide at least one input tensor.")
-    t_spans = []
-    
-    if w is not None:
-        w_np = w[sample_idx].detach().cpu().numpy()
-        if w_np.ndim > 1:
-            w_np = w_np.squeeze()
-        t_spans.append(len(w_np) / sr)
-    if x is not None:
-        x_np = x[sample_idx].detach().cpu().numpy()
-        if x_np.shape[0] < x_np.shape[1]:
-            x_np = x_np.T
-        t_spans.append(x_np.shape[0] * hop_length / sr)
-    if p is not None:
-        p_np = p[sample_idx].detach().cpu().numpy()
-        if p_np.ndim > 1:
-            p_np = p_np.squeeze()
-        t_spans.append(len(p_np) * hop_length / sr)
-    if per is not None:
-        per_np = per[sample_idx].detach().cpu().numpy()
-        if per_np.ndim > 1:
-            per_np = per_np.squeeze()
-        t_spans.append(len(per_np) * hop_length / sr)
-    max_t = max(t_spans) if t_spans else 0
-    fig, axs = plt.subplots(num_plots, 1, figsize=(14, 4*num_plots), sharex=True)
-    if num_plots == 1:
-        axs = [axs]
-    if show_voiced_regions and per is not None:
-        per_np = per[sample_idx].detach().cpu().numpy()
-        if per_np.ndim > 1:
-            per_np = per_np.squeeze()
-        t_per = np.arange(len(per_np)) * hop_length / sr
-        threshold = 0.5
-        for ax in axs:
-            for i in range(len(per_np)-1):
-                if per_np[i] > threshold:
-                    ax.axvspan(t_per[i], t_per[i+1], color='lightblue', alpha=0.2, zorder=0)
-    cu_ax = 0
-    if w is not None:
-        w_np = w[sample_idx].detach().cpu().numpy()
-        if w_np.ndim > 1:
-            w_np = w_np.squeeze()
-        t = np.arange(len(w_np)) / sr
-        axs[cu_ax].plot(t, w_np, color="tab:blue")
-        
-        if show_energy:
-            frame_length = hop_length
-            hop_length_energy = hop_length // 2
-            energy = []
-            for i in range(0, len(w_np)-frame_length, hop_length_energy):
-                frame = w_np[i:i+frame_length]
-                energy.append(np.sqrt(np.mean(frame**2)))
-            energy = np.array(energy)
-            energy = energy / np.max(energy) * 0.8 * max(abs(w_np.min()), abs(w_np.max()))  
-            t_energy = np.arange(len(energy)) * hop_length_energy / sr
-            axs[cu_ax].plot(t_energy, energy, color="red", alpha=0.7, label="Energy")
-            axs[cu_ax].legend(loc='upper right')
-        axs[cu_ax].set_title("Waveform")
-        axs[cu_ax].set_ylabel("Amplitude")
-        axs[cu_ax].set_xlim([0, max_t])
-        axs[cu_ax].grid(True, axis='x', linestyle='--', alpha=0.3)
-        cu_ax += 1
-    
-    if x is not None:
-        x_np = x[sample_idx].detach().cpu().numpy()
-        if x_np.shape[0] < x_np.shape[1]:
-            x_np = x_np.T
-        axs[cu_ax].imshow(x_np.T, aspect="auto", origin="lower", cmap="magma", 
-                                   extent=[0, x_np.shape[0]*hop_length/sr, 0, x_np.shape[1]])
-        axs[cu_ax].set_title("Spectrogram")
-        axs[cu_ax].set_ylabel("Mel Bin")
-        axs[cu_ax].set_xlim([0, max_t])
-        axs[cu_ax].grid(True, axis='x', linestyle='--', alpha=0.3)
-        cu_ax += 1
-    
-    if p is not None:
-        p_np = p[sample_idx].detach().cpu().numpy()
-        if p_np.ndim > 1:
-            p_np = p_np.squeeze()
-        t_p = np.arange(len(p_np)) * hop_length / sr
-        axs[cu_ax].plot(t_p, p_np, color="tab:green")
-        axs[cu_ax].set_title("Pitch")
-        axs[cu_ax].set_ylabel("Frequency (Hz)")
-        axs[cu_ax].set_xlim([0, max_t])
-        axs[cu_ax].grid(True, axis='both', linestyle='--', alpha=0.3)
-        axs[cu_ax].set_ylim([0, min(1000, p_np.max() * 1.2)])
-        cu_ax += 1
-    
-    if per is not None:
-        per_np = per[sample_idx].detach().cpu().numpy()
-        if per_np.ndim > 1:
-            per_np = per_np.squeeze()
-        t_per = np.arange(len(per_np)) * hop_length / sr
-        axs[cu_ax].plot(t_per, per_np, color="tab:red")
-        axs[cu_ax].set_title("Period (Voice Activity)")
-        axs[cu_ax].set_ylabel("periodocity")
-        axs[cu_ax].set_xlim([0, max_t])
-        axs[cu_ax].grid(True, axis='both', linestyle='--', alpha=0.3)
-        axs[cu_ax].set_ylim([-0.05, 1.05])
-        axs[cu_ax].axhline(y=0.5, color='k', linestyle='--', alpha=0.3)
-    
-    if markers is not None:
-        for i, t in enumerate(markers):
-            label = marker_labels[i] if marker_labels and i < len(marker_labels) else None
-            for ax in axs:
-                ax.axvline(x=t, color='k', linestyle='-', alpha=0.7, label=label if i == 0 else None)
-        if marker_labels:
-            axs[0].legend(loc='upper right', fontsize='small')
-    axs[-1].set_xlabel("t (s)")
-    fig.suptitle(title, fontsize=16)
-    plt.tight_layout(rect=[0, 0, 1, 0.97]) # type: ignore
-    plt.show()
-    return fig
-
-def valid(default_value, *items):
-    """Get first non-None item"""
-    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()}
-
-class Sinusoids(nn.Module):
-    def __init__(self, ctx: int, dims: int):
-        super().__init__()
-
-        position = torch.arange(start=0, end=ctx, dtype=dtype).unsqueeze(dim=1)
-        div_term = torch.exp(input=torch.arange(start=0, end=dims, step=2, dtype=dtype) * -(math.log(10000.0) / dims))
-        features = torch.zeros(ctx, dims)
-        features[:, 0::2] = torch.sin(position * div_term)
-        features[:, 1::2] = torch.cos(position* div_term)
-        self.register_buffer('sinusoid', tensor=features)
-        self.positional_embeddings = nn.Parameter(self.sinusoid.clone()) # type: ignore
-    def forward(self, positions):
-        position_embeddings = self.positional_embeddings[positions]
-        return position_embeddings
-
-def sinusoids(length, channels, max_tscale=10000):
-    assert channels % 2 == 0
-    log_tscale_increment = torch.log(torch.tensor(float(max_tscale))) / (channels // 2 - 1)
-    inv_tscales = torch.exp(-log_tscale_increment * torch.arange(channels // 2, device=device, dtype=torch.float32))
-    scaled_t = torch.arange(length, device=device, dtype=torch.float32).unsqueeze(1) * inv_tscales.unsqueeze(0)
-    return torch.cat([torch.sin(scaled_t), torch.cos(scaled_t)], dim=1)
-
-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, f0=False, f0t=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 torchaudio
-    # import torchaudio.functional
-    # import torchaudio.transforms
-
-    # 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)),
-    #         "f0t": 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"]
-    labels = tokenizer.encode(batch["transcription"])
-    wav = load_wave(wave_data=audio, sample_rate=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,
-    }
-
-    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, n_fft=1024, hop_length=256, window_fn=torch.hann_window):
-    #     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
-    #     spectrogram_tensor = torch.tensor(spectrogram_tensor)
-    #     return spectrogram_tensor
-    if spec: 
-        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
-        spectrogram_tensor = torch.tensor(spectrogram_tensor)
-    
-
-
-    # if spec:    
-        # if isinstance(wav, torch.Tensor):
-        #     wav = wav.to(device)
-        # spectrogram_tensor = mel_spectrogram(wav, sample_rate, **spectrogram_config)
-        # spectrogram_tensor = spectrogram_tensor.permute(1, 0)
-
-
-    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):
-        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 compute_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 f0 or f0t or pitch 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 f0:
-        f0_tensor = torch.from_numpy(f0_np)
-    else:
-        f0_tensor = None
-
-    if f0t:
-        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
-        f0t_tensor = torch.cat([torch.tensor([bos_pitch]), pitch_tok])
-        f0t_tensor = torch.where(f0t_tensor == 0.0, torch.zeros_like(f0t_tensor), (f0t_tensor - 71.0) / (500.0 - 71.0))
-    else:
-        f0t_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 = compute_pitch(wav, sample_rate, hop_length=hop_length)[0]
-        p_tensor = torch.from_numpy(p_tensor)
-        p_tensor = p_tensor.unsqueeze(0) 
-        # p_tensor = torch.from_numpy(f0_np)
-    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 f0_tensor is not None:
-            dummy_tensor = torch.ones_like(f0_tensor)
-        elif f0t_tensor is not None:
-            dummy_tensor = torch.ones_like(f0t_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"['f0']: {f0_tensor.shape if f0 else None}") 
-        print(f"['f0t']: {f0t_tensor.shape if f0t 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,
-        "f0": f0_tensor if f0 else None,
-        "f0t": f0t_tensor if f0t 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,
-    }
-
-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,
-        "f0t": 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(
-            "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)
-        raw_test = raw_test.filter(filter_func)
-        raw_train = raw_train.cast_column("audio", Audio(sampling_rate=sample_rate))
-        raw_test = raw_test.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
-
-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, layer=layer, act=act, feature=feature, features=features)
-    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, kernel_size, 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)
-
-        # self.spectrogram_encoder = nn.Sequential(
-        #     Conv1d(input_dims, 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)
-
-        # self.waveform_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:
-                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)
-        # elif feature == "spectrogram":
-        #     x = self.spectrogram_encoder(x).permute(0, 2, 1)
-        # elif feature == "waveform":
-        #     x = self.waveform_encoder(x).permute(0, 2, 1)
-
-        # 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.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", "f0t", "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
-
-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):
-    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