Update model.py

model.py

@@ -1,8 +1,10 @@
 
-import pyworld as pw
 import os
+import pyworld as pw
 import math
 import warnings
+import time
+import random
 import logging
 import gzip
 import base64
@@ -11,6 +13,7 @@ import torchaudio
 import torch.nn.functional as F
 import torch.nn.init as init
 from torch import nn, Tensor
+from torch.utils.data import Dataset, DataLoader
 import numpy as np
 from einops import rearrange
 import matplotlib.pyplot as plt
@@ -18,16 +21,15 @@ from typing import Optional, Dict, Union, List, Tuple, Any
 from functools import partial
 from datetime import datetime
 from datasets import load_dataset, Audio
-from transformers.trainer_seq2seq import Seq2SeqTrainer
-from transformers.training_args_seq2seq import Seq2SeqTrainingArguments
-import transformers
+from torch.utils.tensorboard import SummaryWriter
+import tqdm
+from tqdm import tqdm
 import evaluate
 from dataclasses import dataclass
-
+import aiohttp    
 torch.backends.cudnn.allow_tf32 = True
 torch.backends.cuda.matmul.allow_tf32 = True
 torch.set_float32_matmul_precision('high')
-transformers.utils.logging.set_verbosity_error()
 
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 dtype = torch.float32
@@ -40,8 +42,6 @@ tokenizer = None
 optimizer = None
 scheduler = None
 model = None
-Residual = None
-MultiheadA = None
 
 @dataclass
 class Dimensions:
@@ -284,22 +284,13 @@ class rotary(nn.Module):
         self.freqs.data.copy_(freqs)
         self.theta.data.copy_(theta)
 
-    def get_bias(self, f0, ctx):
+    def get_pitch_bias(self, f0):
         if f0 is None:
             return None
-        if f0.dim() == 1:
-            length = f0.shape[0]
-            if length == ctx:
-                return f0
-            frames = length / ctx
-            idx = torch.arange(ctx, device=f0.device)
-            idx = (idx * frames).long().clamp(0, length - 1)
-            f0 = f0[idx]
-        f0_norm = (f0 - f0.mean()) / (f0.std() + 1e-8)
-        f0_sim = torch.exp(-torch.cdist(f0_norm.unsqueeze(1),
+        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)))
-        # diff = f0_norm[:, None] - f0_norm[None, :]
-        # f0_sim = torch.exp(-diff.pow(2))                         
         return f0_sim.unsqueeze(0).unsqueeze(0)
 
     def f0proj(self, f0):
@@ -313,7 +304,6 @@ class rotary(nn.Module):
         return f0.to(device=device, dtype=dtype)
 
     def synth_f0(self, f0, ctx):
-        # f0 = self.f0proj(f0)
         if f0.dim() == 1:
             length = f0.shape[0]
             if length == ctx:
@@ -321,7 +311,7 @@ class rotary(nn.Module):
             frames = length / ctx
             idx = torch.arange(ctx, device=f0.device)
             return f0[idx]
-    
+
     def align_f0(self, ctx, f0):
         f0 = self.f0proj(f0)
         if f0.dim() == 3:
@@ -361,26 +351,22 @@ class rotary(nn.Module):
             batch, head, ctx, head_dim = x.shape
         t = torch.arange(ctx, device=device, dtype=dtype)
 
-        f0 = enc.get("f0") if enc is not None else None
         if f0 is not None and f0.dim() == 2:
             if f0.shape[0] == 1: 
                 f0 = f0.squeeze(0)  
             else:
                 f0 = f0.view(-1)        
 
-        if f0 is not None: 
+        if f0 is not None and layer == "encoder":
             f0_mean = f0.mean()
             theta = f0_mean + self.theta
         else:
-            theta = 10000.0  
+            theta = self.theta 
         freqs = (theta / 220.0) * 700 * (torch.pow(10, torch.linspace(0, 2595 * torch.log10(torch.tensor(1 + 8000/700)), 
                 self.dim // 2, device=device, dtype=dtype) / 2595) - 1) / 1000
 
-        if "rot2" in self.debug and self.counter % 100 == 0:
-            print(f" [Rotary] {layer}{self.counter} --- [f0] {f0.shape if f0 is not None else None} [Theta] {theta.item():.2f} [Freqs] {freqs.shape} {freqs.mean():.2f} [ctx] {ctx}")
-        
         freqs = t[:, None] * freqs[None, :]
-        if self.radii and f0 is not None and layer == "encoder":
+        if self.radii and f0 is not None:
             radius = f0.to(device, dtype)
             L = radius.shape[0]
             if L != ctx:
@@ -403,7 +389,6 @@ class rotary(nn.Module):
             theta_value = theta.item() if isinstance(theta, torch.Tensor) else theta
             print(f" [{layer}] [f0] {f0.shape if f0 is not None else None} [Theta] {theta_value:.2f} [Freqs] {freqs.shape} {freqs.mean():.2f} [ctx] {ctx} [Radius] {radius.shape} {radius.mean():.2f}")
 
-
         if "rot3" in self.debug and self.counter % 100 == 0:
             print(f" [Rotary] {layer}{self.counter} --- [f0] {f0.shape if f0 is not None else None} [Theta] {theta.item():.2f} [Freqs] {freqs.shape} {freqs.mean():.2f} [ctx] {ctx} [Radius] {radius.shape} {radius.mean():.2f}")
         
@@ -428,19 +413,6 @@ class rotary(nn.Module):
         x1 = x1.view(orig_shape)
         return torch.cat([x1.type_as(x), x2], dim=-1)
 
-    @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 MultiheadA(nn.Module):
     _seen = set()  
     rbf = False
@@ -472,8 +444,7 @@ class MultiheadA(nn.Module):
                 dims=dims,
                 head=head,
                 debug=debug,
-                radii=True if "radii" in debug else False,
-                use_pbias=True if "pbias" in debug else False,
+                radii=True,
                 )
         else:
             self.rope = None
@@ -525,12 +496,12 @@ class MultiheadA(nn.Module):
         
         qk = (q * scale) @ (k * scale).transpose(-1, -2)
         if self.rope.use_pbias:
-            f0 = enc.get("f0", None) if enc is not None else None                     
-            pbias = self.rope.get_bias(f0, q2)
+            f0 = enc.get("f0", None) if enc is not None else None
+            pbias = self.rope.use_pbias(f0)
             if pbias is not None:
-                qk = qk + pbias
+                qk = qk + pbias[:,:,:q.shape[2],:q.shape[2]]
         token_ids = k[:, :, :, 0]
-        zscale = torch.ones_like(token_ids)
+        zscale = torch.ones_like(token_ids, device=device, dtype=dtype)
         fzero = torch.clamp(F.softplus(self.fzero), self.minz, self.maxz)
         zscale[token_ids.float() == self.pad_token] = fzero
         
@@ -619,6 +590,7 @@ class Residual(nn.Module):
         self.t_gate = tgate
         self.m_gate = mgate
         self.c_gate = cgate
+        self.skip_gates=True
         
         self.blend = nn.Parameter(torch.tensor(0.5))
         
@@ -628,8 +600,8 @@ class Residual(nn.Module):
                   "leaky_relu": nn.LeakyReLU(), "elu": nn.ELU()}
         act_fn = act_map.get(act, nn.GELU())
 
-        self.attna = MultiheadA(dims, head, rotary_emb=True, debug=debug)
-        self.attnb = (MultiheadA(dims, head, rotary_emb=True, debug=debug) if cross_attn else None)
+        self.attna = MultiheadA(dims=dims, head=head, rotary_emb=True, debug=debug)
+        self.attnb = (MultiheadA(dims=dims, head=head, rotary_emb=True, debug=debug) if cross_attn else None)
         
         mlp = dims * 4
         self.mlp = nn.Sequential(Linear(dims, mlp), act_fn, Linear(mlp, dims))
@@ -660,25 +632,28 @@ class Residual(nn.Module):
         
         normx = self.lnc(x)
         mlp_out = self.mlp(normx)
-        
-        if self.t_gate:
-            gate = self.t_gate(normx)
-            x = x + gate * mlp_out
-               
-        elif self.m_gate:
-            gate = self.m_gate(normx)
-            x = x + gate * mlp_out
-        
-        elif self.c_gate:
-            gate_output = self.c_gate(normx, self.features)
-            x = x + gate_output
 
+        if self.skip_gates:
+            x = x + mlp_out
         else:
-            if hasattr(self, 'mlp_gate'):
-                mlp_gate = self.mlp_gate(normx)
-                x = x + mlp_gate * mlp_out
+            if self.t_gate:
+                gate = self.t_gate(normx)
+                x = x + gate * mlp_out
+                
+            elif self.m_gate:
+                gate = self.m_gate(normx)
+                x = x + gate * mlp_out
+            
+            elif self.c_gate:
+                gate_output = self.c_gate(normx, self.features)
+                x = x + gate_output
+
             else:
-                x = x + mlp_out
+                if hasattr(self, 'mlp_gate'):
+                    mlp_gate = self.mlp_gate(normx)
+                    x = x + mlp_gate * mlp_out
+                else:
+                    x = x + mlp_out
                 
         if "residual" in self.debug and self.counter % 100 == 0:
             print(f"Step {self.counter}: Residual block output shape: {x.shape}, xa shape: {xa.shape if xa is not None else None}")        
@@ -897,7 +872,7 @@ class AudioEncoder(nn.Module):
             )
         })
 
-    def forward(self, enc, order=None, layer="encoder"):
+    def forward(self, enc, layer="encoder"):
         enc = dict_to(enc, device, dtype)
         
         if self.counter < 1:
@@ -906,13 +881,10 @@ class AudioEncoder(nn.Module):
             p = default(enc.get("pitch"), enc.get("f0"))
             plot_waveform(x=s, w=w, p=p, hop_length=128)
 
-        if order is None:
-            order = self.features
-
         out = {}
         out.update(enc)
 
-        for f in order:
+        for f in self.features:
             if f in enc and f in self.blocks:
                 x = enc[f]
                 for block in self.blocks[f]:
@@ -921,7 +893,7 @@ class AudioEncoder(nn.Module):
                         
         if "encoder" in self.debug and self.counter % 100 == 0:
             shapes = {k: v.shape for k, v in enc.items()}
-            print(f"Step {self.counter}: mode: {list(enc.keys()) }: shapes: {shapes}, order: {order}")
+            print(f"Step {self.counter}: mode: {list(enc.keys()) }: shapes: {shapes}")
         self.counter += 1
         return out
 
@@ -969,13 +941,12 @@ class TextDecoder(nn.Module):
         mask = self.mask[:x.shape[1], :x.shape[1]]
         x = self.token(x) + self.positional[:x.shape[1]]
         x = F.dropout(x, p=self.dropout, training=self.training)
-
+        
         for block in self.block:
-            x = block(x, xa=None, mask=mask, enc=enc, layer=layer)
+            x = block(x, xa=None, mask=mask, enc=None, layer=layer)
 
         for f in order:
             if f in enc:
-
                 xa = enc[f]
                 for block in self.blocks[f]:
                     out = block(x=x, xa=xa, mask=None, enc=None, layer=layer)
@@ -1029,12 +1000,10 @@ class Echo(nn.Module):
         for name, module in self.encoder.named_modules():
             if isinstance(module, (rotary)):
                 module.update_base(f0)
-                module.return_f0(f0)
 
         for name, module in self.decoder.named_modules():
             if isinstance(module, (rotary)):
                 module.update_base(f0)
-                module.return_f0(f0)
 
     def set_alignment_head(self, dump: bytes):
         array = np.frombuffer(
@@ -1050,16 +1019,19 @@ class Echo(nn.Module):
         return self.decoder(input_ids, encoder_output)
         
     def forward(self,
+        decoder_input_ids=None,
         labels=None,
         waveform: Optional[torch.Tensor]=None,
         input_ids=None,
         spectrogram: torch.Tensor=None,
         pitch: Optional[torch.Tensor]=None,
         f0: Optional[torch.Tensor]=None,
+        f0d: Optional[torch.Tensor]=None,
         envelope: Optional[torch.Tensor]=None,
         phase: Optional[torch.Tensor]=None,
         ) -> Dict[str, torch.Tensor]:
 
+        decoder_input_ids = input_ids
         encoder_inputs = {}
         if spectrogram is not None:
             encoder_inputs["spectrogram"] = spectrogram
@@ -1073,7 +1045,7 @@ class Echo(nn.Module):
             encoder_inputs["phase"] = phase
         if f0 is not None:
             encoder_inputs["f0"] = f0
-
+            
         encoder_outputs = self.encoder(encoder_inputs)
         logits = self.decoder(input_ids, encoder_outputs)
 
@@ -1170,122 +1142,58 @@ class Echo(nn.Module):
         self.counter = 0
         print("Counter reset to 0.")
 
-metric = evaluate.load(path="wer")
-    
-@dataclass
-class DataCollator:
-    tokenizer: Any
-    def __call__(self, features: List[Dict[str, torch.Tensor]]) -> Dict[str, torch.Tensor]:
-        pad_token_id = tokenizer.pad_token_id if hasattr(tokenizer, 'pad_token_id') else 0
-        bos_token_id = tokenizer.bos_token_id if hasattr(tokenizer, 'bos_token_id') else 1
+def ctx_to_samples(audio_ctx, hop_length):
+    samples_token = hop_length * 2
+    n_samples = audio_ctx * samples_token
+    return n_samples
+
+def load_wave(wave_data, sample_rate):
+    if isinstance(wave_data, str):
+        waveform, sr = torchaudio.load(uri=wave_data, normalize=False)
+    elif isinstance(wave_data, dict):
+        waveform = torch.tensor(data=wave_data["array"]).float()
+        sr = wave_data["sampling_rate"]
+    else:
+        raise TypeError("Invalid wave_data format.")
+
+    if sr != sample_rate:
+        original_length = waveform.shape[1]
+        target_length = int(original_length * (sample_rate / sr))
         
-        batch = {}
+        resampler = torchaudio.transforms.Resample(orig_freq=sr, new_freq=sample_rate)
+        waveform = resampler(waveform)
+        
+    return waveform
+
+def pad(array, target_length, axis=-1, dtype: torch.dtype = torch.float32):
+    if isinstance(array, np.ndarray):
+        array = torch.from_numpy(array).to(dtype)
+    if torch.is_tensor(array):
+        if array.shape[axis] > target_length:
+            array = array.index_select(
+                dim=axis,
+                index=torch.arange(
+                    end=target_length, device=array.device, dtype=torch.long
+                ),
+            )
+        if array.shape[axis] < target_length:
+            pad_widths = [(0, 0)] * array.ndim
+            pad_widths[axis] = (0, target_length - array.shape[axis])
+            array = F.pad(
+                input=array, pad=[pad for sizes in pad_widths[::-1] for pad in sizes]
+            )
+        array = array.to(dtype=dtype)
+    else:
+        raise TypeError(
+            f"Unsupported input type: {type(array)}. Expected torch.Tensor or np.ndarray."
+        )
+    return array
 
-        if "spectrogram" in features[0] and features[0]["spectrogram"] is not None:
-            spectrogram_list = [f["spectrogram"] for f in features]
-            max_len_feat = max(f.shape[-1] for f in spectrogram_list)
-            pad_spectrogram = []
-            for feat in spectrogram_list:                
-                current_len = feat.shape[-1]
-                padding = max_len_feat - current_len
-                if padding > 0:
-                    pad_feat = F.pad(feat, (0, padding), mode='constant', value=pad_token_id)
-                else:
-                    pad_feat = feat
-                pad_spectrogram.append(pad_feat)
-            batch["spectrogram"] = torch.stack(pad_spectrogram)
-            
-        if "waveform" in features[0] and features[0]["waveform"] is not None:
-            waveform_list = [f["waveform"] for f in features]
-            max_len_wav = max(w.shape[-1] for w in waveform_list)
-            pad_waveforms = []
-            for wav in waveform_list:
-                current_len = wav.shape[-1]
-                padding = max_len_wav - current_len
-                if padding > 0:
-                    if wav.ndim == 1:
-                        wav = wav.unsqueeze(0)
-                    pad_wav = F.pad(wav, (0, padding), mode='constant', value=pad_token_id)
-                else:
-                    pad_wav = wav
-                pad_waveforms.append(pad_wav)
-            batch["waveform"] = torch.stack(pad_waveforms)
-
-        if "label" in features[0] and features[0]["label"] is not None:
-            labels_list = [f["label"] for f in features]
-            max_len = max(len(l) for l in labels_list)
-            all_ids = []
-            all_labels = []
+def exact_div(x, y):
+    assert x % y == 0
+    return x // y
 
-            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 + [pad_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)
-
-        if "pitch" in features[0] and features[0]["pitch"] is not None:
-            pitch_list = [f["pitch"] for f in features]
-            max_len_pitch = max(e.shape[-1] for e in pitch_list)
-            pad_pitch = []
-            for pitch in pitch_list:
-                current_len = pitch.shape[-1]
-                padding = max_len_pitch - current_len
-                if padding > 0:
-                    pad_pitch_item = F.pad(pitch, (0, padding), mode='constant', value=pad_token_id)
-                else:
-                    pad_pitch_item = pitch
-                pad_pitch.append(pad_pitch_item)
-            batch["pitch"] = torch.stack(pad_pitch)
-      
-        if "f0" in features[0] and features[0]["f0"] is not None:
-            f0_list = [f["f0"] for f in features]
-            max_len_f0 = max(f.shape[-1] for f in f0_list)
-            pad_f0 = []
-            for f0 in f0_list:
-                current_len = f0.shape[-1]
-                padding = max_len_f0 - current_len
-                if padding > 0:
-                    pad_f0_item = F.pad(f0, (0, padding), mode='constant', value=pad_token_id)
-                else:
-                    pad_f0_item = f0
-                pad_f0.append(pad_f0_item)
-            batch["f0"] = torch.stack(pad_f0)
-
-        if "envelope" in features[0] and features[0]["envelope"] is not None:
-            env_list = [f["envelope"] for f in features]
-            max_len = max(f.shape[-1] for f in env_list)
-            pad_env = []
-            for feat in env_list:                
-                current_len = feat.shape[-1]
-                padding = max_len - current_len
-                if padding > 0:
-                    pad_feat = F.pad(feat, (0, padding), mode='constant', value=pad_token_id)
-                else:
-                    pad_feat = feat
-                pad_env.append(pad_feat)
-            batch["envelope"] = torch.stack(pad_env)
-
-        if "phase" in features[0] and features[0]["phase"] is not None:
-            ph_list = [f["phase"] for f in features]
-            max_len = max(f.shape[-1] for f in ph_list)
-            pad_ph = []
-            for feat in ph_list:                
-                current_len = feat.shape[-1]
-                padding = max_len - current_len
-                if padding > 0:
-                    pad_feat = F.pad(feat, (0, padding), mode='constant', value=pad_token_id)
-                else:
-                    pad_feat = feat
-                pad_ph.append(pad_feat)
-            batch["phase"] = torch.stack(pad_ph) 
-        return batch
+metrics = evaluate.load(path="wer")
 
 def hilbert_transform(x):
     N = x.shape[-1]
@@ -1338,26 +1246,51 @@ def process_spectrogram_with_hilbert(spec):
     phase = torch.angle(analytic)
     return envelope, phase
         
-def load_wave(wave_data, sample_rate):
-    if isinstance(wave_data, str):
-        waveform, sr = torchaudio.load(uri=wave_data, normalize=False)
-    elif isinstance(wave_data, dict):
-        waveform = torch.tensor(data=wave_data["array"]).float()
-        sr = wave_data["sampling_rate"]
-    else:
-        raise TypeError("Invalid wave_data format.")
-    
-    if waveform.dim() == 1:
-        waveform = waveform.unsqueeze(0)
-    
-    if sr != sample_rate:
-        original_length = waveform.shape[1]
-        target_length = int(original_length * (sample_rate / sr))
-        
-        resampler = torchaudio.transforms.Resample(orig_freq=sr, new_freq=sample_rate)
-        waveform = resampler(waveform)
-        
-    return waveform.flatten()
+@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)
+
+        for key in all_keys:
+            if key == "label":
+                labels_list = [f["label"] for f in features]
+                max_len = max(len(l) for l in labels_list)
+                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 + [pad_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", "f0", "envelope", "phase"]:
+                items = [f[key] for f in features if key in f]
+                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 extract_features(batch, tokenizer, spectrogram, waveforms, pitch, frequency=False,
                      hop_length=128, fmin=0, fmax=8000, n_mels=128, n_fft=1024, sampling_rate=16000,
@@ -1443,72 +1376,20 @@ def extract_features(batch, tokenizer, spectrogram, waveforms, pitch, frequency=
     batch["label"] = tokenizer.encode(batch["transcription"], add_special_tokens=False)
     return batch
 
-def compute_metrics(eval_pred, compute_result: bool = True, 
-                    print_pred: bool = False, num_samples: int = 0, tokenizer=None, pitch=None, model=None):
-    
-    pred_logits = eval_pred.predictions
-    label_ids = eval_pred.label_ids
-
-    if hasattr(pred_logits, "cpu"):
-        pred_logits = pred_logits.cpu()
-    if hasattr(label_ids, "cpu"):
-        label_ids = label_ids.cpu()
-    if isinstance(pred_logits, tuple):
-        pred_ids = pred_logits[0]
+def compute_metrics(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_logits
-    if hasattr(pred_ids, "ndim") and pred_ids.ndim == 3:
-        if not isinstance(pred_ids, torch.Tensor):
-            pred_ids = torch.tensor(pred_ids)
-        pred_ids = pred_ids.argmax(dim=-1)
-        pred_ids = pred_ids.tolist()
-            
-    if hasattr(label_ids, "tolist"):
-        label_ids = label_ids.tolist()
-        
-    label_ids = [[0 if token == -100 else token for token in seq] for seq in label_ids]
-    pred_str = tokenizer.batch_decode(pred_ids, skip_special_tokens=False)
-    label_str = tokenizer.batch_decode(label_ids, skip_special_tokens=False)
-
-    if print_pred:
-        for i in range(min(num_samples, len(pred_str))):
-            print(f"Preds: {pred_str[i]}")
-            print(f"Label: {label_str[i]}")
-            print(f"preds: {pred_ids[i]}")
-            print(f"label: {label_ids[i]}")
-            print("--------------------------------")  
-    
+        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)
-    wer = 100 * metric.compute(predictions=pred_str, references=label_str)
-    
-    if model is None:
-        global global_model
-        if 'global_model' in globals():
-            model = global_model
-    
-    if model is not None:
-        trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad) / 1_000_000
-        if trainable_params > 0:
-            efficiency_score = (100 - wer) / trainable_params
-        else:
-            print("Warning: Zero trainable parameters detected")
-            efficiency_score = 0.0
-    else:
-        print("Warning: Model not available for parameter counting")
-        trainable_params = 0.0
-        efficiency_score = 0.0
-    
-    if hasattr(wer, "item"):
-        wer = wer.item()
-    
-    metrics = {
-        "wer": float(wer),
-        "trainable_params_M": float(trainable_params),
-        "efficiency_score": float(efficiency_score),
-    }
-    
-    return metrics
+    wer = metrics.compute(predictions=pred_str, references=label_str)
+    return {"wer": wer}
 
 logger = logging.getLogger(__name__)
 
@@ -1533,13 +1414,16 @@ def setup_tokenizer(token: str, local_tokenizer_path: str = "D:/newmodel/model/t
             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, skip_special_tokens=True):
         results = []
         for ids in ids_list:
+            if not isinstance(ids, list):
+                ids = ids.tolist()
             if skip_special_tokens:
                 ids = [id for id in ids if id not in [0, 1, 2]]
             results.append(tokenizer.decode(ids))
-        return results
+        return results        
     def save_pretrained(save_dir):
         os.makedirs(save_dir, exist_ok=True)
         tokenizer.save(f"{save_dir}/tokenizer.json")
@@ -1552,229 +1436,324 @@ def setup_tokenizer(token: str, local_tokenizer_path: str = "D:/newmodel/model/t
     return tokenizer
 
 def prepare_datasets(tokenizer, token: str, sanity_check: bool = False, dataset_config: Optional[Dict] = None) -> Tuple[any, any]:
-    if dataset_config is None:
-        dataset_config = {
-            "spectrogram": True,
-            "waveforms": True,
-            "pitch": True,
-            "frequency": True,
-            "downsamples": True,
-            "hop_length": 128,
-            "fmin": 50,
-            "fmax": 2000,
-            "n_mels": 128,
-            "n_fft": 1024,
-            "sampling_rate": 16000,
-        }
-    
-    dataset = load_dataset(
-        "google/fleurs", 
-        "en_us", 
-        token=token, 
-        trust_remote_code=True,
-        streaming=False)
-    
-    dataset = dataset.cast_column(column="audio", feature=Audio(sampling_rate=16000)).select_columns(["audio", "transcription"])
-    
+
     if sanity_check:
+
+        dataset = load_dataset(
+            "./librispeech_asr.py", "clean", "train.100",
+            storage_options={'client_kwargs': {'timeout': aiohttp.ClientTimeout(total=3600)}},
+            token=token, trust_remote_code=True, streaming=False)
+
+        dataset = dataset.rename_column("text", "transcription")
+        dataset = dataset.cast_column(column="audio", feature=Audio(sampling_rate=16000)).select_columns(["audio", "transcription"])
+
         dataset = dataset["test"].take(10)
         dataset = dataset.select_columns(["audio", "transcription"])
-        logger.info(f"Sanity dataset size: {dataset.num_rows}") 
-        print(f"Sanity dataset size: {dataset.num_rows}")
         prepare_fn = partial(extract_features, tokenizer=tokenizer, **dataset_config)
-        
-        dataset = dataset.map(
-            function=prepare_fn, 
-            remove_columns=["audio", "transcription"]
-        ).with_format(type="torch")
+        dataset = dataset.map(function=prepare_fn, remove_columns=["audio", "transcription"]).with_format(type="torch")
         train_dataset = dataset
         test_dataset = dataset
     else:
+        cache_dir = "./processed_datasets"
+        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 dataset_config is None:
+            dataset_config = {
+                "spectrogram": True,
+                "waveforms": True,
+                "pitch": True,
+                "frequency": True,
+                "downsamples": True,
+                "hop_length": 128,
+                "fmin": 50,
+                "fmax": 2000,
+                "n_mels": 128,
+                "n_fft": 1024,
+                "sampling_rate": 16000,
+            }
+
+        dataset = load_dataset(
+            "./librispeech_asr.py", "clean", "train.100",
+            storage_options={'client_kwargs': {'timeout': aiohttp.ClientTimeout(total=3600)}},
+            token=token, trust_remote_code=True, streaming=False)
+
+        dataset = dataset.rename_column("text", "transcription")
+        dataset = dataset.cast_column(column="audio", feature=Audio(sampling_rate=16000)).select_columns(["audio", "transcription"])
+
         def filter_func(x):
             return (0 < len(x["transcription"]) < 512 and
                    len(x["audio"]["array"]) > 0 and
                    len(x["audio"]["array"]) < 1500 * 160)
         
-        dataset = dataset.filter(filter_func).shuffle(seed=4)
-        logger.info(f"Dataset size: {dataset['train'].num_rows}, {dataset['test'].num_rows}")
-        print(f"Dataset size: {dataset['train'].num_rows}, {dataset['test'].num_rows}")
+        dataset = dataset.filter(filter_func)
         prepare_fn = partial(extract_features, tokenizer=tokenizer, **dataset_config)
-        columns_to_remove = list(next(iter(dataset.values())).features)
-        train_dataset = dataset["train"]
-        test_dataset = dataset["test"].take(50)
-        logger.info(f"Train dataset size: {train_dataset.num_rows}, Test dataset size: {test_dataset.num_rows}")
-        
+
+        train_dataset = dataset["train.100"].take(10000)
+        test_dataset = dataset["test"].take(1000)
         train_dataset = train_dataset.map(
             function=prepare_fn, 
-            remove_columns=columns_to_remove
+            remove_columns=["audio", "transcription"]
         ).with_format(type="torch")
         
         test_dataset = test_dataset.map(
             function=prepare_fn, 
-            remove_columns=columns_to_remove
+            remove_columns=["audio", "transcription"]
         ).with_format(type="torch")
-    
+
+        train_dataset.save_to_disk(cache_file_train)
+        test_dataset.save_to_disk(cache_file_test)
+
     return train_dataset, test_dataset
 
-def get_training_args(
-    log_dir: str,
-    batch_eval_metrics: bool = False,
-    max_steps: int = 10,
-    save_steps: int = 1000,
-    eval_steps: int = 1,
-    warmup_steps: int = 0,
-    num_train_epochs: int = 1,
-    logging_steps: int = 1,
-    eval_on_start: bool = False,
-    learning_rate: float = 1e-4,
-    weight_decay: float = 0.01,
-    max_grad_norm: float = 1.0,
-) -> Seq2SeqTrainingArguments:
-
-    return Seq2SeqTrainingArguments(
-        output_dir=log_dir,
-        per_device_train_batch_size=1,
-        per_device_eval_batch_size=1,
-        gradient_accumulation_steps=1,
-        eval_accumulation_steps=None,
-        eval_strategy="steps",
-        save_strategy="no",
-        max_steps=max_steps,
-        save_steps=save_steps,
-        eval_steps=eval_steps,
-        warmup_steps=warmup_steps,
-        num_train_epochs=num_train_epochs,
-        logging_steps=logging_steps,
-        logging_dir=log_dir,
-        logging_strategy="steps",
-        report_to=["tensorboard"],
-        push_to_hub=False,
-        disable_tqdm=False,
-        save_total_limit=1,
-        label_names=["labels"],
-        optim="adamw_torch",
-        lr_scheduler_type="cosine",
-        learning_rate=learning_rate,
-        weight_decay=weight_decay,
-        save_safetensors=False,
-        eval_on_start=eval_on_start,
-        batch_eval_metrics=batch_eval_metrics,
-        max_grad_norm=max_grad_norm,     
-    )
+@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)
+
+        for key in all_keys:
+            if key == "label":
+                labels_list = [f["label"] for f in features]
+                max_len = max(len(l) for l in labels_list)
+                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 + [pad_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", "f0", "envelope", "phase"]:
+                items = [f[key] for f in features if key in f]
+                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 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 main():
-     
     token = ""
-    log_dir = os.path.join('./output/logs', datetime.now().strftime(format='%m-%d_%H_%M_%S'))
-    os.makedirs(name=log_dir, exist_ok=True)
+    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)
 
-    def sanity(sanity: bool):
-
-        if sanity:
-            training_args = get_training_args(
-            log_dir,
-            batch_eval_metrics = False,
-            max_steps = 10,
-            save_steps = 0,
-            eval_steps = 1,
-            warmup_steps = 0,
-            logging_steps = 1,
-            eval_on_start = False,
-            learning_rate = 5e-6,
-            weight_decay = 0.01,
-            )
-        else:
-            training_args = get_training_args(
-            log_dir,
-            batch_eval_metrics = False,
-            max_steps = 1000,
-            save_steps = 1005,
-            eval_steps = 100,   
-            warmup_steps = 100,
-            logging_steps = 10,
-            eval_on_start = False,
-            learning_rate = 2.5e-4,
-            weight_decay = 0.01,
-            )
-
-        return training_args
-        
     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"]
-        )
-
-    sanity_check = False
+        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"]
+    )
 
-    training_args = sanity(sanity_check)
     dataset_config = {
-        "spectrogram": True,
-        "waveforms": True,
-        "pitch": False,
-        "downsamples": False,
-        "frequency": False,
-        "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}
-    
+        "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)
-    
-    global global_model
-    global_model = model
-    
-    metrics_fn = partial(compute_metrics, print_pred=False, num_samples=5, 
-                    tokenizer=tokenizer, model=model)
-    
-    print(f"{'Sanity check' if sanity_check else 'Training'} mode")
     train_dataset, test_dataset = prepare_datasets(
-        tokenizer=tokenizer,
-        token=token,
-        sanity_check=sanity_check,
-        dataset_config=dataset_config)
+        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)
 
-    trainer = Seq2SeqTrainer(
-        args=training_args,
+    optimizer = get_optimizer(model)
+    scheduler = get_scheduler(optimizer)
+    loss_fn = get_loss_fn()
+
+    train_and_evaluate(
         model=model,
-        train_dataset=train_dataset,
-        eval_dataset=test_dataset,
-        data_collator=DataCollator(tokenizer=tokenizer),
-        compute_metrics=metrics_fn,
-        ) 
-       
-    model.init_weights()
-    trainer.train()
+        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
+    )
 
 if __name__ == "__main__":
     main()
 
-# from tensorboard import program
-# log_dir = "./output/logs" 
-# tb = program.TensorBoard()
-# tb.configure(argv=[None, '--logdir', log_dir])
-# url = tb.launch()
-# print(f"TensorBoard started at {url}")
-