Adding train/evaluate/metrics files

evaluate.py

@@ -0,0 +1,32 @@
+from utils.metrics import (
+    calculate_msd,
+    calculate_f0_correlation,
+    calculate_phoneme_accuracy,
+    calculate_spectral_convergence
+)
+from inference import run_tts
+
+def evaluate_bd_tts(model, test_dataset):
+    metrics = {}
+    pred_audio, target_audio = [], []
+
+    for text, target in test_dataset:
+        pred = run_tts(text)
+        pred_audio.append(pred)
+        target_audio.append(target)
+
+    metrics['mel_spectral_distance'] = calculate_msd(pred_audio, target_audio)
+    metrics['f0_correlation'] = calculate_f0_correlation(pred_audio, target_audio)
+    metrics['phoneme_accuracy'] = calculate_phoneme_accuracy(pred_audio, target_audio)
+    metrics['spectral_convergence'] = calculate_spectral_convergence(pred_audio, target_audio)
+
+    # Accent classifier is usually a pretrained model
+    # Placeholder: you’d plug in your Bangla accent classifier here
+    metrics['accent_score'] = 0.85  
+
+    return metrics
+
+if __name__ == "__main__":
+    test_dataset = [("আমি বাংলা বলি।", "reference.wav")]  # dummy dataset
+    print(evaluate_bd_tts(None, test_dataset))
+

metrics.py

@@ -0,0 +1,57 @@
+# utils/metrics.py
+import numpy as np
+import librosa
+
+def calculate_msd(pred_audio, target_audio, sr=22050):
+    """
+    Mel Spectral Distance (MSD) between predicted and target audio.
+    """
+    # Convert to mel-spectrogram
+    pred_mel = librosa.feature.melspectrogram(y=pred_audio, sr=sr)
+    target_mel = librosa.feature.melspectrogram(y=target_audio, sr=sr)
+
+    # Convert to dB
+    pred_db = librosa.power_to_db(pred_mel, ref=np.max)
+    target_db = librosa.power_to_db(target_mel, ref=np.max)
+
+    # Mean squared difference
+    return np.mean((pred_db - target_db) ** 2)
+
+
+def calculate_f0_correlation(pred_audio, target_audio, sr=22050):
+    """
+    Pitch correlation (F0 correlation) between predicted and target.
+    """
+    f0_pred, _, _ = librosa.pyin(pred_audio, fmin=50, fmax=500, sr=sr)
+    f0_target, _, _ = librosa.pyin(target_audio, fmin=50, fmax=500, sr=sr)
+
+    # Remove NaNs
+    mask = ~np.isnan(f0_pred) & ~np.isnan(f0_target)
+    if np.sum(mask) == 0:
+        return 0.0
+    return np.corrcoef(f0_pred[mask], f0_target[mask])[0, 1]
+
+
+def calculate_phoneme_accuracy(pred_phonemes, target_phonemes):
+    """
+    Simple phoneme accuracy metric.
+    (Here, pred_phonemes and target_phonemes are lists of symbols)
+    """
+    if len(target_phonemes) == 0:
+        return 0.0
+    correct = sum(p == t for p, t in zip(pred_phonemes, target_phonemes))
+    return correct / len(target_phonemes)
+
+
+def calculate_spectral_convergence(pred_audio, target_audio, sr=22050):
+    """
+    Spectral convergence: how close the predicted spectrum is to the target.
+    """
+    pred_spec = np.abs(librosa.stft(pred_audio))
+    target_spec = np.abs(librosa.stft(target_audio))
+
+    numerator = np.linalg.norm(target_spec - pred_spec, 'fro')
+    denominator = np.linalg.norm(target_spec, 'fro')
+
+    return numerator / (denominator + 1e-8)
+

train.py

@@ -0,0 +1,59 @@
+import torch
+from torch.utils.data import DataLoader
+from transformers import AdamW, get_linear_schedule_with_warmup
+from utils.dataset import BDTtsDataset
+from inference import tts  # reuse your model
+
+training_config = {
+    "learning_rate": 1e-4,
+    "batch_size": 16,
+    "warmup_steps": 1000,
+    "gradient_accumulation_steps": 4,
+    "mixed_precision": True,
+    "save_strategy": "steps",
+    "save_steps": 500,
+    "eval_steps": 100,
+    "num_epochs": 5
+}
+
+def train():
+    dataset = BDTtsDataset("./data/train")
+    dataloader = DataLoader(dataset, batch_size=training_config["batch_size"], shuffle=True)
+
+    optimizer = AdamW(tts.model.parameters(), lr=training_config["learning_rate"])
+    scheduler = get_linear_schedule_with_warmup(
+        optimizer,
+        num_warmup_steps=training_config["warmup_steps"],
+        num_training_steps=len(dataloader) * training_config["num_epochs"]
+    )
+
+    scaler = torch.cuda.amp.GradScaler() if training_config["mixed_precision"] else None
+    step = 0
+
+    for epoch in range(training_config["num_epochs"]):
+        for batch in dataloader:
+            inputs, targets = batch
+            optimizer.zero_grad()
+
+            with torch.cuda.amp.autocast(enabled=scaler is not None):
+                outputs = tts.model(inputs)
+                loss = outputs.loss if hasattr(outputs, "loss") else torch.nn.functional.mse_loss(outputs, targets)
+
+            if scaler:
+                scaler.scale(loss).backward()
+                scaler.step(optimizer)
+                scaler.update()
+            else:
+                loss.backward()
+                optimizer.step()
+
+            scheduler.step()
+            step += 1
+
+            if step % training_config["save_steps"] == 0:
+                torch.save(tts.model.state_dict(), f"checkpoints/model_step{step}.pth")
+                print(f"Saved checkpoint at step {step}")
+
+if __name__ == "__main__":
+    train()
+