train_model.py

import pandas as pd
import numpy as np
import librosa
import librosa.display
import soundfile

import torch
import torch.utils.data as tud
import torch.nn as nn
import torch.nn.functional as F
from torchsummary import summary
from torchaudio.transforms import MelSpectrogram, ComputeDeltas

from torch.optim.adamw import AdamW
from torch.optim.lr_scheduler import CosineAnnealingWarmRestarts

import textgrid

import re
import json

import os
from os import listdir
from os.path import isfile, join

import math
import random

from pathlib import Path
from IPython.display import Audio

from tqdm import tqdm
tqdm.pandas()

import matplotlib.pyplot as plt

import pyaudio
import wave

from fastprogress import master_bar, progress_bar

from google.cloud import texttospeech

import warnings
warnings.simplefilter("ignore", UserWarning)

import sys
import time

with open('your_config.json', 'r') as archivo_json:
  config_datos = json.load(archivo_json)

wake_words = config_datos['wake_words']
wake_words_sequence = []
for indice, elemento in enumerate(wake_words):
  wake_words_sequence.append(str(indice))
wake_word_seq_map = dict(zip(wake_words, wake_words_sequence))

sr = 16000

dataset_language = config_datos['dataset_language']
add_vanilla_noise_to_negative_dataset = config_datos['add_vanilla_noise_to_negative_dataset']
generateVoicesWithGoogle = config_datos['voices_generation_with_google']
windowSizeFromConfig = config_datos['window_size_ms']

path_to_dataset = 'dataset'
path_to_dataset_w = path_to_dataset + '/'

# ------------

ttsConfig = config_datos['tts_generated_clips']

if (ttsConfig['rate']['start'] >= 0.25 and ttsConfig['rate']['start'] <= 4.0 and ttsConfig['rate']['stop'] >= 0.25 and ttsConfig['rate']['stop'] <= 4.0 and ttsConfig['rate']['start'] <= ttsConfig['rate']['stop']) and (ttsConfig['pitch']['start'] >= -20.0 and ttsConfig['pitch']['start'] <= 20.0 and ttsConfig['pitch']['stop'] >= -20.0 and ttsConfig['pitch']['stop'] <= 20.0 and ttsConfig['pitch']['start'] <= ttsConfig['pitch']['stop']):
  i7512 = 1
else:
  print('your_config.json > tts_generated_clips invalid values. rate must be in the range [0.25, 4.0] and pitch must be in the range [-20.0, 20.0], and start must be lower than stop.')
  sys.exit()

print("NOTE: Running this file may take several minutes.")

if not torch.cuda.is_available():
  print('CUDA is not available.  Train on CPU.')
else:
  print('CUDA is available!  Train on GPU.')

# sys.exit()

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

wake_words_withOOV = wake_words[:]
wake_words_withOOV.append("oov")

def list_files(mypath):
  return [mypath + f for f in listdir(mypath) if isfile(join(mypath, f))]

def getWavAudioDuration(nombre_archivo):
  with wave.open(nombre_archivo, 'rb') as archivo_audio:
    # Obtén la frecuencia de muestreo (número de muestras por segundo)
    frecuencia_muestreo = archivo_audio.getframerate()
    # Obtén el número total de frames (muestras)
    num_frames = archivo_audio.getnframes()
    # Calcula la duración en segundos
    duracion = num_frames / frecuencia_muestreo
  return duracion

noise_test = list_files('noise/noise_test/')
noise_train_complete = list_files('noise/noise_train/')

regex_pattern = r'\b(?:{})\b'.format('|'.join(map(re.escape, wake_words)))
pattern = re.compile(regex_pattern, flags=re.IGNORECASE)
def wake_words_search(pattern, word):
  try:
    return bool(pattern.search(word))
  except TypeError:
    return False

# Dataset checkpoint
positive_train_data = pd.read_csv(path_to_dataset_w+'positive/train.csv')
positive_dev_data = pd.read_csv(path_to_dataset_w+'positive/dev.csv')
positive_test_data = pd.read_csv(path_to_dataset_w+'positive/test.csv')

negative_train_data = pd.read_csv(path_to_dataset_w+'negative/train.csv')
negative_dev_data = pd.read_csv(path_to_dataset_w+'negative/dev.csv')
negative_test_data = pd.read_csv(path_to_dataset_w+'negative/test.csv')

negative_train_data['sentence'] = negative_train_data['sentence'].apply(
  lambda x: x if isinstance(x, str) and x.strip() != "" else "Ouhbg agihugy uagoboiuslv bofieglk"
)
negative_dev_data['sentence'] = negative_dev_data['sentence'].apply(
  lambda x: x if isinstance(x, str) and x.strip() != "" else "Ouhbg agihugy uagoboiuslv bofieglk"
)
negative_test_data['sentence'] = negative_test_data['sentence'].apply(
  lambda x: x if isinstance(x, str) and x.strip() != "" else "Ouhbg agihugy uagoboiuslv bofieglk"
)

# Add vanilla noise to negative dataset
max_noise_duration = 90000

if add_vanilla_noise_to_negative_dataset:

  for noiseItemPath in noise_train_complete:
    noiseItemDuration = round(getWavAudioDuration(noiseItemPath) * 1000, 1)
    if noiseItemDuration <= max_noise_duration:
      negative_train_data = pd.concat([negative_train_data, pd.DataFrame([{
        'path': noiseItemPath,
        'sentence': 'Hsdflkjhsdf lhskldhfansljf sndlkfjnsdf',
        'timestamps': {},
        'duration': noiseItemDuration
      }])], ignore_index=True)

  for noiseItemPath in noise_test:
    noiseItemDuration = round(getWavAudioDuration(noiseItemPath) * 1000, 1)
    if noiseItemDuration <= max_noise_duration:
      negative_test_data = pd.concat([negative_test_data, pd.DataFrame([{
        'path': noiseItemPath,
        'sentence': 'Hsdflkjhsdf lhskldhfansljf sndlkfjnsdf',
        'timestamps': {},
        'duration': noiseItemDuration
      }])], ignore_index=True)

# max duration in positive dataset
print(f"Max duration in positive train {positive_train_data['duration'].max()}")
print(f"Min duration in positive train {positive_train_data['duration'].min()}")
print(f"Max duration in positive dev {positive_dev_data['duration'].max()}")
print(f"Min duration in positive dev {positive_dev_data['duration'].min()}")
print(f"Max duration in positive test {positive_test_data['duration'].max()}")
print(f"Min duration in positive test {positive_test_data['duration'].min()}")

# max duration in negative dataset
print(f"Max duration in negative train {negative_train_data['duration'].max()}")
print(f"Min duration in negative train {negative_train_data['duration'].min()}")
print(f"Max duration in negative dev {negative_dev_data['duration'].max()}")
print(f"Min duration in negative dev {negative_dev_data['duration'].min()}")
print(f"Max duration in negative test {negative_test_data['duration'].max()}")
print(f"Min duration in negative test {negative_test_data['duration'].min()}")

print('positive_train_data head:')
print(positive_train_data.head(10))
print('negative_train_data head:')
print(negative_train_data.head(10))

print('positive_train_data tail:')
print(positive_train_data.tail(10))
print('negative_train_data tail:')
print(negative_train_data.tail(10))

print(f"positive_train_data size {positive_train_data.shape}")
print(f"positive_dev_data size {positive_dev_data.shape}")
print(f"positive_test_data size {positive_test_data.shape}")
print(f"negative_train_data size {negative_train_data.shape}")
print(f"negative_dev_data size {negative_dev_data.shape}")
print(f"negative_test_data size {negative_test_data.shape}")

train_ds = pd.concat([positive_train_data , negative_train_data]).sample(frac=1).reset_index(drop=True)
dev_ds = pd.concat([positive_dev_data , negative_dev_data]).sample(frac=1).reset_index(drop=True)
test_ds = pd.concat([positive_test_data , negative_test_data]).sample(frac=1).reset_index(drop=True)

print(f"Training dataset size {train_ds.shape}")
print(f"Validation dataset size {dev_ds.shape}")
print(f"Test dataset size {test_ds.shape}")

# checking pattern spread on train_ds
for word in wake_words:
  word_pattern = re.compile(r'\b'+word+r'\b', flags=re.IGNORECASE)
  print(word + f" Total word {(train_ds[[wake_words_search(word_pattern, sentence) for sentence in train_ds['sentence']]].size/train_ds.size) * 100} %")

generated_data = path_to_dataset_w + 'generated'
Path(f"{generated_data}").mkdir(parents=True, exist_ok=True)

os.environ["GOOGLE_APPLICATION_CREDENTIALS"]=config_datos['google_credentials_file']

client = texttospeech.TextToSpeechClient()

def generate_voices(word):
  Path(f"{generated_data}/{word}").mkdir(parents=True, exist_ok=True)
  # Set the text input to be synthesized
  synthesis_input = texttospeech.SynthesisInput(text=word)
  # Performs the list voices request
  voices = client.list_voices()
  # Get english voices

  en_voices =  [voice.name for voice in voices.voices if voice.name.split("-")[0] == dataset_language]
  speaking_rates = np.arange(ttsConfig['rate']['start'], ttsConfig['rate']['stop'], ttsConfig['rate']['step']).tolist()
  pitches = np.arange(ttsConfig['pitch']['start'], ttsConfig['pitch']['stop'], ttsConfig['pitch']['step']).tolist()

  file_count = 0
  start = time.time()

  for voi in en_voices:
    for sp_rate in speaking_rates:
      for pit in pitches:
        file_name = f'{generated_data}/{word}/{voi}_{sp_rate}_{pit}.wav'
        voice = texttospeech.VoiceSelectionParams(language_code=voi[:5], name=voi)
        # Select the type of audio file you want returned
        audio_config = texttospeech.AudioConfig(
          # format of the audio byte stream.
          audio_encoding=texttospeech.AudioEncoding.LINEAR16,

          #Speaking rate/speed, in the range [0.25, 4.0]. 1.0 is the normal native speed
          speaking_rate=sp_rate,

          #Speaking pitch, in the range [-20.0, 20.0]. 20 means increase 20 semitones from the original pitch. -20 means decrease 20 semitones from the original pitch.
          pitch=pit # [-10, -5, 0, 5, 10]
        )
        response = client.synthesize_speech(
          request={"input": synthesis_input, "voice": voice, "audio_config": audio_config}
        )
        # The response's audio_content is binary.
        with open(file_name, "wb") as out:
          out.write(response.audio_content)
          file_count+=1
        if file_count%100 == 0:
          end = time.time()
          print(f"generated {file_count} files in {end-start} seconds")

# Voices generation with Google Cloud text-to-speech API
if generateVoicesWithGoogle:
  print("Generating audios with Google Cloud text-to-speech API:")
  for word in wake_words:
    generate_voices(word)

for word in wake_words:
  d = {}
  d['path'] = [f"{generated_data}/{word}/{file_name}" for file_name in os.listdir(f"{generated_data}/{word}")]
  d['sentence'] = [word] * len(d['path'])
  pd.DataFrame(data=d).to_csv(f"{generated_data}/{word}.csv", index=False)

word_cols = {'path' : [], 'sentence': []}
train, dev, test = pd.DataFrame(word_cols), pd.DataFrame(word_cols), pd.DataFrame(word_cols)
for word in wake_words:
  word_df = pd.read_csv(f"{generated_data}/{word}.csv")
  tra, val, te =  np.split(word_df.sample(frac=1, random_state=42),  [int(.6*len(word_df)), int(.8*len(word_df))])
  train = pd.concat([train , tra]).sample(frac=1).reset_index(drop=True)
  dev = pd.concat([dev , val]).sample(frac=1).reset_index(drop=True)
  test = pd.concat([test , te]).sample(frac=1).reset_index(drop=True)

# Checkpoint save
train.to_csv(f"{generated_data}/train.csv", index=False)
dev.to_csv(f"{generated_data}/dev.csv", index=False)
test.to_csv(f"{generated_data}/test.csv", index=False)

# add dummy values for these columns for generated data
train['timestamps'] = ''
train['duration'] = ''

dev['timestamps'] = ''
dev['duration'] = ''

test['timestamps'] = ''
test['duration'] = ''

train_ds = pd.concat([train_ds , train]).sample(frac=1).reset_index(drop=True)
dev_ds = pd.concat([dev_ds , dev]).sample(frac=1).reset_index(drop=True)
test_ds = pd.concat([test_ds , test]).sample(frac=1).reset_index(drop=True)

print('train_ds.shape', train_ds.shape)

print(f"Training dataset size {train_ds.shape}")
print(f"Validation dataset size {dev_ds.shape}")
print(f"Test dataset size {test_ds.shape}")

# now verify how much data we have for train set
for word in wake_words:
  word_pattern = re.compile(r'\b'+word+r'\b', flags=re.IGNORECASE)
  print(word + f" (2) Total word {(train_ds[[wake_words_search(word_pattern, sentence) for sentence in train_ds['sentence']]].size/train_ds.size) * 100} %")

print('oov' + f" (2) Total {(train_ds[[not wake_words_search(pattern, sentence) for sentence in train_ds['sentence']]].size/train_ds.size) * 100} %")

print('train_ds.tail:')
print(train_ds.tail(20))

# --- Add noise
noise_train = noise_train_complete[:int(len(noise_train_complete) * 0.8)]
noise_dev = noise_train_complete[int(len(noise_train_complete) * 0.8):]

# random.randint(0,len(noise_dev))
# print noise data stats
print(f"Train noise dataset {len(noise_train)}")
print(f"Validation noise dataset {len(noise_dev)}")
print(f"Test noise dataset {len(noise_test)}")

# Detener aquí para revisión del dataset, después de esto ya empieza entrenamiento.
# sys.exit()

key_pattern = re.compile("\'(?P<k>[^ ]+)\'")
def compute_labels(metadata, audio_data):
  label = len(wake_words) # by default negative label

  # if it is generated data then
  if metadata['sentence'].lower() in wake_words and not metadata['timestamps']:
    label = int(wake_word_seq_map[metadata['sentence'].lower()])
  else:
    # if the sentence has one wakeword get label and end timestamp
    wake_word_found = False
    for word in metadata['sentence'].lower().split():
      word = re.sub('\W+', '', word)
      if word in wake_words:
        wake_word_found = True
        break

    if wake_word_found:
      label = int(wake_word_seq_map[word])
      if word in  metadata['timestamps']:
        timestamps = metadata['timestamps']
        if type(timestamps) == str:
          timestamps = json.loads(key_pattern.sub(r'"\g<k>"', timestamps))
        word_ts = timestamps[word]
        audio_start_idx = int((word_ts['start'] * 1000) * sr / 1000)
        audio_end_idx = int((word_ts['end'] * 1000) * sr / 1000)
        audio_data = audio_data[audio_start_idx:audio_end_idx]
      else: # if there are issues with word alignment, we might not get ts
        label = len(wake_words)  # mark them for negative
    else: # is negative data
      if metadata['timestamps'] and type(metadata['timestamps']) == str: # si tiene timestamps recorto una palabra aleatoria
        timestamps = json.loads(key_pattern.sub(r'"\g<k>"', metadata['timestamps']))
        if timestamps:
          randomWord = random.choice(list(timestamps.keys()))
          word_ts = timestamps[randomWord]
          audio_start_idx = int((word_ts['start'] * 1000) * sr / 1000)
          audio_end_idx = int((word_ts['end'] * 1000) * sr / 1000)
          audio_data = audio_data[audio_start_idx:audio_end_idx]

  return label, audio_data

class AudioCollator(object):
  def __init__(self, noise_set=None, device='cpu'):
    self.noise_set = noise_set
    self.device = device

  def __call__(self, batch):
    batch_tensor = {}
    window_size_ms = windowSizeFromConfig
    max_length = int(window_size_ms/1000 * sr)
    audio_tensors = []
    labels = []
    for sample in batch:
      # get audio_data in tensor format
      audio_data = librosa.core.load(sample['path'], sr=sr, mono=True)[0]
      # get the label and its audio
      label, audio_data = compute_labels(sample, audio_data)
      audio_data_length = audio_data.size / sr * 1000 #ms

      # below is to make sure that we always got length of 12000
      # i.e 750 ms with sr 16000
      # trim to max_length
      if audio_data_length > window_size_ms:
        # randomly trim either at start and end
        if random.random() < 0.5:
          audio_data = audio_data[:max_length]
        else:
          audio_data = audio_data[audio_data.size-max_length:]

      # pad with zeros
      if audio_data_length < window_size_ms:
        # randomly either append or prepend
        if random.random() < 0.5:
          audio_data = np.append(audio_data, np.zeros(int(max_length - audio_data.size)))
        else:
          audio_data = np.append(np.zeros(int(max_length - audio_data.size)), audio_data)

      # Add noise
      if self.noise_set:
        noise_level =  random.randint(5, 30)/100 # 5 to 30%
        noise_sample = librosa.core.load(self.noise_set[random.randint(0,len(self.noise_set)-1)], sr=sr, mono=True)[0]
        # randomly select first or last seq of noise
        if random.random() < 0.5:
          audio_data = (1 - noise_level) * audio_data +  noise_level * noise_sample[:max_length]
        else:
          audio_data = (1 - noise_level) * audio_data +  noise_level * noise_sample[-max_length:]

      # Convert to tensor and send to device immediately
      audio_tensor = torch.from_numpy(audio_data).to(self.device)
      audio_tensors.append(audio_tensor)
      labels.append(label)

    batch_tensor = {
        'audio': torch.stack(audio_tensors),  # Already on correct device
        'labels': torch.tensor(labels, device=self.device)
    }

    return batch_tensor

# --- Prepare for train
batch_size = 16
num_workers = 0

# Pass device to collators so they create tensors on GPU directly
train_audio_collator = AudioCollator(noise_set=noise_train, device=device)
train_dl = tud.DataLoader(train_ds.to_dict(orient='records'),
                  batch_size=batch_size,
                  drop_last=True,
                  shuffle=True,
                  num_workers=num_workers,
                  collate_fn=train_audio_collator)

dev_audio_collator = AudioCollator(noise_set=noise_dev, device=device)
dev_dl = tud.DataLoader(dev_ds.to_dict(orient='records'),
                  batch_size=batch_size,
                  num_workers=num_workers,
                  collate_fn=dev_audio_collator)

test_audio_collator = AudioCollator(noise_set=noise_test, device=device)
test_dl = tud.DataLoader(test_ds.to_dict(orient='records'),
                  batch_size=batch_size,
                  num_workers=num_workers,
                  collate_fn=test_audio_collator)

# For ZMUV, use CPU since we're just calculating statistics
zmuv_audio_collator = AudioCollator(device='cpu')
zmuv_dl = tud.DataLoader(train_ds.to_dict(orient='records'),
                  batch_size=1,
                  num_workers=num_workers,
                  collate_fn=zmuv_audio_collator)

# IMPROVED MODEL WITH DEPTHWISE SEPARABLE CONVOLUTIONS AND RESIDUAL CONNECTIONS
class DepthwiseSeparableConv2d(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, padding=0, stride=1, bias=True):
        super(DepthwiseSeparableConv2d, self).__init__()
        self.depthwise = nn.Conv2d(in_channels, in_channels, kernel_size=kernel_size, 
                                  padding=padding, stride=stride, groups=in_channels, bias=bias)
        self.pointwise = nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=bias)
    
    def forward(self, x):
        x = self.depthwise(x)
        x = self.pointwise(x)
        return x

class ResidualBlock(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, padding, stride, use_depthwise=True):
        super(ResidualBlock, self).__init__()
        
        if use_depthwise:
            self.conv = DepthwiseSeparableConv2d(in_channels, out_channels, kernel_size, padding, stride)
        else:
            self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, padding=padding, stride=stride, bias=True)
        
        self.bn = nn.BatchNorm2d(out_channels, affine=True)
        self.relu = nn.ReLU()
        self.dropout = nn.Dropout2d(0.3)  # More aggressive dropout
        
        # Skip connection adjustment if dimensions change
        self.skip_connection = None
        if in_channels != out_channels or stride != 1:
            self.skip_connection = nn.Sequential(
                nn.Conv2d(in_channels, out_channels, 1, stride=stride, bias=False),
                nn.BatchNorm2d(out_channels)
            )
    
    def forward(self, x):
        identity = x
        
        out = self.conv(x)
        out = self.bn(out)
        out = self.relu(out)
        out = self.dropout(out)
        
        # Adjust skip connection if needed
        if self.skip_connection is not None:
            identity = self.skip_connection(identity)
        elif identity.shape != out.shape:
            # Simple padding/cropping for shape mismatch
            if identity.shape[2] > out.shape[2]:  # Height
                diff = identity.shape[2] - out.shape[2]
                identity = identity[:, :, diff//2:diff//2+out.shape[2], :]
            if identity.shape[3] > out.shape[3]:  # Width
                diff = identity.shape[3] - out.shape[3]
                identity = identity[:, :, :, diff//2:diff//2+out.shape[3]]
        
        # Only add if shapes match exactly
        if identity.shape == out.shape:
            out = out + identity
        
        return out

class ImprovedCNN(nn.Module):
    def __init__(self, num_labels, num_maps1, num_maps2, num_hidden_input, hidden_size):
        super(ImprovedCNN, self).__init__()
        self.num_hidden_input = num_hidden_input
        
        # First residual block
        self.encoder1 = nn.Sequential(
            ResidualBlock(1, num_maps1, (8, 16), padding=(4, 0), stride=(2, 2)),
            nn.MaxPool2d(2)
        )
        
        # Second residual block  
        self.encoder2 = nn.Sequential(
            ResidualBlock(num_maps1, num_maps2, (5, 5), padding=2, stride=(2, 1)),
            nn.MaxPool2d(2)
        )
        
        # Output layers with more aggressive dropout
        self.output = nn.Sequential(
            nn.Linear(num_hidden_input, hidden_size),
            nn.ReLU(),
            nn.Dropout(0.4),  # Increased dropout
            nn.Linear(hidden_size, hidden_size // 2),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(hidden_size // 2, num_labels)
        )
    
    def forward(self, x):
        x = x[:, :1]  # log_mels only
        x = x.permute(0, 1, 3, 2)  # (time, n_mels)
        
        # Pass through residual blocks
        x1 = self.encoder1(x)
        x2 = self.encoder2(x1)
        
        # Flattening
        x = x2.view(-1, self.num_hidden_input)
        return self.output(x)

class CNN_Cal(nn.Module):
  def __init__(self, num_labels, num_maps1, num_maps2, num_hidden_input, hidden_size):
    super(CNN_Cal, self).__init__()
    conv0 = nn.Conv2d(1, num_maps1, (8, 16), padding=(4, 0), stride=(2, 2), bias=True)
    pool = nn.MaxPool2d(2)
    conv1 = nn.Conv2d(num_maps1, num_maps2, (5, 5), padding=2, stride=(2, 1), bias=True)
    self.encoder1 = nn.Sequential(conv0,
                                  nn.ReLU(),
                                  pool,
                                  nn.BatchNorm2d(num_maps1, affine=True))
    self.encoder2 = nn.Sequential(conv1,
                                  nn.ReLU(),
                                  pool,
                                  nn.BatchNorm2d(num_maps2, affine=True))

  def forward(self, x):
    x = x[:, :1] # log_mels only
    x = x.permute(0, 1, 3, 2)  # change to (time, n_mels)
    # pass through first conv layer
    x1 = self.encoder1(x)
    # pass through second conv layer
    x2 = self.encoder2(x1)
    # flattening - keep first dim batch same, flatten last 3 dims
    x = x2.view(x2.size(0), -1)
    return x

num_labels = len(wake_words) + 1 # oov
num_maps1  = 48
num_maps2  = 64
num_hidden_input =  768
hidden_size = 128

# Use improved model
model = ImprovedCNN(num_labels, num_maps1, num_maps2, num_hidden_input, hidden_size)

model.to(device)
print(model)

model_calc = CNN_Cal(num_labels, num_maps1, num_maps2, num_hidden_input, hidden_size)
model_calc.to(device)
print(model_calc)

print(summary(model_calc, input_size=(1,80,61)))
print(summary(model, input_size=(1,40,61)))

from typing import Iterable

class ZmuvTransform(nn.Module):
  def __init__(self):
    super().__init__()
    self.register_buffer('total', torch.zeros(1))
    self.register_buffer('mean', torch.zeros(1))
    self.register_buffer('mean2', torch.zeros(1))

  def update(self, data, mask=None):
    with torch.no_grad():
      if mask is not None:
        data = data * mask
        mask_size = mask.sum().item()
      else:
        mask_size = data.numel()
      self.mean = (data.sum() + self.mean * self.total) / (self.total + mask_size)
      self.mean2 = ((data ** 2).sum() + self.mean2 * self.total) / (self.total + mask_size)
      self.total += mask_size

  def initialize(self, iterable: Iterable[torch.Tensor]):
    for ex in iterable:
      self.update(ex)

  @property
  def std(self):
    return (self.mean2 - self.mean ** 2).sqrt()

  def forward(self, x):
    return (x - self.mean) / self.std

# Per-utterance normalization transform
class PerUtteranceNorm(nn.Module):
    def __init__(self, eps=1e-8):
        super().__init__()
        self.eps = eps
    
    def forward(self, x):
        # x shape: (batch, channels, n_mels, time)
        # Normalize over n_mels and time dimensions for each utterance
        mean = x.mean(dim=(2, 3), keepdim=True)
        std = x.std(dim=(2, 3), keepdim=True)
        return (x - mean) / (std + self.eps)

zmuv_transform = ZmuvTransform().to(device)
per_utterance_norm = PerUtteranceNorm().to(device)

if Path(path_to_dataset_w + "zmuv.pt.bin").exists():
  zmuv_transform.load_state_dict(torch.load(str(path_to_dataset_w + "zmuv.pt.bin")))
else:
  for idx, batch in enumerate(tqdm(zmuv_dl, desc="Constructing ZMUV")):
    zmuv_transform.update(batch['audio'].to(device))
  print(dict(zmuv_mean=zmuv_transform.mean, zmuv_std=zmuv_transform.std))
  torch.save(zmuv_transform.state_dict(), str(path_to_dataset_w + "zmuv.pt.bin"))

print(f"Mean is {zmuv_transform.mean.item():0.6f}")
print(f"Standard Deviation is {zmuv_transform.std.item():0.6f}")

zmuv_mean = zmuv_transform.mean.item()
zmuv_std = zmuv_transform.std.item()

# IMPROVED TRAINING CONFIGURATION
learning_rate = 0.001
# Adaptive weight decay - different rates for conv and linear layers
conv_weight_decay = 0.0001
linear_weight_decay = 0.001

# Label smoothing CrossEntropyLoss
criterion = nn.CrossEntropyLoss(label_smoothing=0.1)

# Separate parameter groups for adaptive weight decay
conv_params = []
linear_params = []

for name, param in model.named_parameters():
    if param.requires_grad:
        if 'conv' in name or 'depthwise' in name or 'pointwise' in name:
            conv_params.append(param)
        else:
            linear_params.append(param)

optimizer = AdamW([
    {'params': conv_params, 'weight_decay': conv_weight_decay},
    {'params': linear_params, 'weight_decay': linear_weight_decay}
], lr=learning_rate)

# Cosine Annealing Warm Restarts scheduler
scheduler = CosineAnnealingWarmRestarts(optimizer, T_0=10, T_mult=2, eta_min=1e-6)

log_offset = 1e-7
num_mels = 40 # https://en.wikipedia.org/wiki/Mel_scale
num_fft = 512 # window length - Fast Fourier Transform
hop_length = 200  # making hops of size hop_length each time to sample the next window

# CREATE MEL SPECTROGRAM TRANSFORM ONCE AND KEEP IN GPU
mel_spectrogram_transform = MelSpectrogram(n_mels=num_mels,
                                          sample_rate=sr,
                                          n_fft=num_fft,
                                          hop_length=hop_length,
                                          norm='slaney').to(device)

def audio_transform(audio_data):
  # Ensure audio_data is on GPU
  if not audio_data.is_cuda:
    audio_data = audio_data.to(device)
  
  # Use pre-initialized transform that's already on GPU
  log_mels = mel_spectrogram_transform(audio_data.float()).add_(log_offset).log_().contiguous()
  # returns (channel, n_mels, time) - already on GPU
  return log_mels

# Early stopping class
class EarlyStopping:
    def __init__(self, patience=7, min_delta=0.001, restore_best_weights=True):
        self.patience = patience
        self.min_delta = min_delta
        self.restore_best_weights = restore_best_weights
        self.best_loss = None
        self.counter = 0
        self.best_weights = None
        
    def __call__(self, val_loss, model):
        if self.best_loss is None:
            self.best_loss = val_loss
            self.save_checkpoint(model)
        elif val_loss < self.best_loss - self.min_delta:
            self.best_loss = val_loss
            self.counter = 0
            self.save_checkpoint(model)
        else:
            self.counter += 1
            
        if self.counter >= self.patience:
            if self.restore_best_weights:
                model.load_state_dict(self.best_weights)
            return True
        return False
    
    def save_checkpoint(self, model):
        self.best_weights = model.state_dict().copy()

# --- IMPROVED TRAINING LOOP

epochs = config_datos['train_epochs']

# Verify GPU usage
print(f"Device being used: {device}")
print(f"Model is on GPU: {next(model.parameters()).is_cuda}")
print(f"ZMUV transform is on GPU: {zmuv_transform.mean.is_cuda}")

# Early stopping
early_stopping = EarlyStopping(patience=10, min_delta=0.001)

# config for progress bar
mb = master_bar(range(epochs))
mb.names = ['Training loss', 'Validation loss']
x = []

training_losses = []
validation_losses = []
training_accuracies = []
validation_accuracies = []

valid_mean_min = np.inf

for epoch in mb:
  x.append(epoch)
  
  # Training phase
  model.train()
  total_loss = torch.Tensor([0.0]).to(device)
  correct_predictions = 0
  total_predictions = 0
  
  for batch in progress_bar(train_dl, parent=mb):
    # Data is already on GPU from AudioCollator
    audio_data = batch['audio']  # Already on device
    labels = batch['labels']      # Already on device
    
    # Verify data is on GPU
    if not audio_data.is_cuda:
        audio_data = audio_data.to(device)
    if not labels.is_cuda:
        labels = labels.to(device)
    
    # get mel spectograms - audio_transform now handles GPU properly
    mel_audio_data = audio_transform(audio_data)
    
    # Apply global ZMUV normalization
    mel_audio_data = zmuv_transform(mel_audio_data)
    
    # Apply per-utterance normalization
    mel_audio_data = per_utterance_norm(mel_audio_data.unsqueeze(1))
    
    predicted_scores = model(mel_audio_data)
    
    # get loss
    loss = criterion(predicted_scores, labels)

    optimizer.zero_grad()
    
    # backward propagation
    loss.backward()
    
    # Gradient clipping
    torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
    
    optimizer.step()

    with torch.no_grad():
      total_loss += loss
      # Calculate accuracy
      _, predicted = torch.max(predicted_scores, 1)
      correct_predictions += (predicted == labels).sum().item()
      total_predictions += labels.size(0)

  # Update scheduler
  scheduler.step()

  train_mean = total_loss / len(train_dl)
  train_accuracy = correct_predictions / total_predictions
  training_losses.append(train_mean.detach().cpu().item())  # Properly detach before CPU
  training_accuracies.append(train_accuracy)

  # Validation phase
  model.eval()
  validation_loss = torch.Tensor([0.0]).to(device)
  correct_val_predictions = 0
  total_val_predictions = 0
  
  with torch.no_grad():
    for batch in progress_bar(dev_dl, parent=mb):
      # Data is already on GPU from AudioCollator
      audio_data = batch['audio']
      labels = batch['labels']
      
      # Verify data is on GPU
      if not audio_data.is_cuda:
          audio_data = audio_data.to(device)
      if not labels.is_cuda:
          labels = labels.to(device)
      
      # get mel spectograms
      mel_audio_data = audio_transform(audio_data)
      
      # Apply global ZMUV normalization
      mel_audio_data = zmuv_transform(mel_audio_data)
      
      # Apply per-utterance normalization
      mel_audio_data = per_utterance_norm(mel_audio_data.unsqueeze(1))
      
      predicted_scores = model(mel_audio_data)
      
      # get loss
      loss = criterion(predicted_scores, labels)
      validation_loss += loss
      
      # Calculate accuracy
      _, predicted = torch.max(predicted_scores, 1)
      correct_val_predictions += (predicted == labels).sum().item()
      total_val_predictions += labels.size(0)

  val_mean = validation_loss / len(dev_dl)
  val_accuracy = correct_val_predictions / total_val_predictions
  validation_losses.append(val_mean.detach().cpu().item())  # Properly detach before CPU
  validation_accuracies.append(val_accuracy)

  # Update training chart
  mb.update_graph([[x, training_losses], [x, validation_losses]], [0, epochs])
  
  current_lr = optimizer.param_groups[0]['lr']
  mb.write(f"\nEpoch {epoch}: Train loss {train_mean.item():.6f} (acc: {train_accuracy:.4f}) | "
           f"Val loss {val_mean.item():.6f} (acc: {val_accuracy:.4f}) | LR: {current_lr:.6f}")

  # save model if validation loss has decreased
  if val_mean.item() <= valid_mean_min:
    print('Validation loss decreased ({:.6f} --> {:.6f}).  Saving model ...'.format(
    valid_mean_min, val_mean.item()))
    torch.save(model.state_dict(), path_to_dataset_w + 'model_trained.pt')
    valid_mean_min = val_mean.item()
  
  # Early stopping check
  if early_stopping(val_mean.item(), model):
    print(f"Early stopping triggered at epoch {epoch}")
    break

# Final evaluation on test set
model.eval()
test_loss = torch.Tensor([0.0]).to(device)
correct_test_predictions = 0
total_test_predictions = 0
all_predictions = []
all_labels = []

with torch.no_grad():
    for batch in tqdm(test_dl, desc="Testing"):
        audio_data = batch['audio'].to(device)
        labels = batch['labels'].to(device)
        
        mel_audio_data = audio_transform(audio_data)
        mel_audio_data = zmuv_transform(mel_audio_data)
        mel_audio_data = per_utterance_norm(mel_audio_data.unsqueeze(1))
        
        predicted_scores = model(mel_audio_data)
        loss = criterion(predicted_scores, labels)
        test_loss += loss
        
        _, predicted = torch.max(predicted_scores, 1)
        correct_test_predictions += (predicted == labels).sum().item()
        total_test_predictions += labels.size(0)
        
        all_predictions.extend(predicted.cpu().numpy())
        all_labels.extend(labels.cpu().numpy())

test_mean = test_loss / len(test_dl)
test_accuracy = correct_test_predictions / total_test_predictions

print(f"\nFinal Test Results:")
print(f"Test Loss: {test_mean.item():.6f}")
print(f"Test Accuracy: {test_accuracy:.4f}")

# Calculate per-class metrics
from sklearn.metrics import classification_report, confusion_matrix
import numpy as np

# Classification report
print("\nClassification Report:")
print(classification_report(all_labels, all_predictions, target_names=wake_words_withOOV))

# Confusion matrix
print("\nConfusion Matrix:")
cm = confusion_matrix(all_labels, all_predictions)
print(cm)

# Calculate False Accept Rate (FAR) and False Reject Rate (FRR) for each wake word
print("\nPer-class FAR and FRR:")
for i, word in enumerate(wake_words_withOOV):
    if i < len(wake_words):  # Only for actual wake words, not OOV
        tp = cm[i, i]
        fn = cm[i, :].sum() - tp
        fp = cm[:, i].sum() - tp
        tn = cm.sum() - tp - fn - fp
        
        far = fp / (fp + tn) if (fp + tn) > 0 else 0
        frr = fn / (fn + tp) if (fn + tp) > 0 else 0
        
        print(f"{word}: FAR={far:.4f}, FRR={frr:.4f}")

# check if CUDA is available
train_on_gpu = torch.cuda.is_available()

if not train_on_gpu:
  print('CUDA is not available.  Training on CPU ...')
else:
  print('CUDA is available!  Training on GPU ...')

# Save final model
torch.save(model.state_dict(), path_to_dataset_w + 'model_trained_improved.pt')

# Save training history
training_history = {
    'training_losses': [loss for loss in training_losses],
    'validation_losses': [loss for loss in validation_losses],
    'training_accuracies': training_accuracies,
    'validation_accuracies': validation_accuracies,
    'test_loss': test_mean.item(),
    'test_accuracy': test_accuracy,
    'epochs_trained': len(training_losses),
    'training_dataset_size': train_ds.shape[0],
    'validation_dataset_size': dev_ds.shape[0],
    'test_dataset_size': test_ds.shape[0]
}

with open(path_to_dataset_w + 'training_history.json', 'w') as f:
    json.dump(training_history, f, indent=2)

# Save improved model metadata
with open(path_to_dataset_w + 'model_data_improved.json', 'w') as archivo:
  archivo.write(json.dumps({
    "zmuv_mean": zmuv_mean,
    "zmuv_std": zmuv_std,
    "window_size": windowSizeFromConfig,
    "hop_length": hop_length,
    "num_mels": num_mels,
    "num_fft": num_fft,
    "sample_rate": sr,
    "log_offset": log_offset,
    "train_epochs": len(training_losses),
    "original_path": path_to_dataset_w + 'model_trained_improved.pt',
    "final_validation_loss": valid_mean_min,
    "final_test_loss": test_mean.item(),
    "final_test_accuracy": test_accuracy,
    "classes": wake_words_withOOV,
    "classes_base": wake_words,
    "vanilla_noise_in_negative_dataset": add_vanilla_noise_to_negative_dataset,
    "improvements": [
        "depthwise_separable_convolutions",
        "residual_connections", 
        "cosine_annealing_scheduler",
        "gradient_clipping",
        "label_smoothing",
        "per_utterance_normalization",
        "aggressive_dropout",
        "adaptive_weight_decay",
        "early_stopping"
    ]
  }, indent=2))

print(f"\nTraining completed! Models and history saved.")
print(f"Best validation loss: {valid_mean_min:.6f}")
print(f"Final test accuracy: {test_accuracy:.4f}")