Spaces:
Sleeping
Sleeping
| """ | |
| Train Native vs Non-Native classifier on Wav2Vec embeddings. | |
| Uses early stopping, LR scheduler, class weights, and saves best model by val F1. | |
| """ | |
| import os | |
| import sys | |
| import random | |
| # Allow importing model from same directory when run from project root | |
| sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) | |
| import pandas as pd | |
| import numpy as np | |
| import torch | |
| import torch.nn as nn | |
| import torch.optim as optim | |
| from torch.utils.data import DataLoader | |
| from sklearn.model_selection import train_test_split | |
| from sklearn.preprocessing import LabelEncoder | |
| from sklearn.utils.class_weight import compute_class_weight | |
| from sklearn.metrics import ( | |
| accuracy_score, | |
| f1_score, | |
| precision_score, | |
| recall_score, | |
| confusion_matrix, | |
| ) | |
| from model import NativityClassifier | |
| # ----------------------------------------------------------------------------- | |
| # Config (change these if needed) | |
| # ----------------------------------------------------------------------------- | |
| SEED = 42 | |
| MANIFEST_PATH = "manifests/final_training_manifest.csv" | |
| MODEL_SAVE_DIR = "models" | |
| BEST_MODEL_PATH = os.path.join(MODEL_SAVE_DIR, "nativity_model.pth") | |
| # With more compute you can safely raise these | |
| EPOCHS = 200 # max epochs; early stopping will usually stop earlier | |
| BATCH_SIZE = 32 | |
| LR = 1e-3 | |
| WEIGHT_DECAY = 5e-4 # stronger L2 to reduce overfitting | |
| PATIENCE = 15 # stricter early stopping (was 30) — stop sooner | |
| MIN_LR = 1e-6 | |
| # Model architecture (must match model.py and embeddings dim) | |
| HIDDEN_DIMS = (384, 256, 128) | |
| DROPOUT = 0.5 # increased from 0.35 to reduce overfitting | |
| INPUT_DIM = 1024 # overridden in main() from first .npy (1063 when using --mfcc) | |
| OUTPUT_DIM = 2 | |
| # ----------------------------------------------------------------------------- | |
| # Reproducibility | |
| # ----------------------------------------------------------------------------- | |
| def set_seed(seed): | |
| random.seed(seed) | |
| np.random.seed(seed) | |
| torch.manual_seed(seed) | |
| if torch.cuda.is_available(): | |
| torch.cuda.manual_seed_all(seed) | |
| torch.backends.cudnn.deterministic = True | |
| torch.backends.cudnn.benchmark = False | |
| else: | |
| torch.backends.cudnn.benchmark = True | |
| # ----------------------------------------------------------------------------- | |
| # Dataset with optional L2 normalization | |
| # ----------------------------------------------------------------------------- | |
| class EmbeddingDataset(torch.utils.data.Dataset): | |
| def __init__(self, df, embedding_path_col="embedding_path", label_col="label_idx", l2_norm=True): | |
| self.paths = [str(p).replace("\\", "/") for p in df[embedding_path_col].tolist()] | |
| self.labels = df[label_col].values | |
| self.l2_norm = l2_norm | |
| def __len__(self): | |
| return len(self.labels) | |
| def __getitem__(self, idx): | |
| emb = np.load(self.paths[idx]).astype(np.float32) | |
| if self.l2_norm: | |
| n = np.linalg.norm(emb) | |
| if n > 1e-8: | |
| emb = emb / n | |
| return torch.tensor(emb, dtype=torch.float32), torch.tensor( | |
| self.labels[idx], dtype=torch.long | |
| ) | |
| # ----------------------------------------------------------------------------- | |
| # Training | |
| # ----------------------------------------------------------------------------- | |
| def main(): | |
| set_seed(SEED) | |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| print(f"Device: {device}") | |
| os.makedirs(MODEL_SAVE_DIR, exist_ok=True) | |
| df = pd.read_csv(MANIFEST_PATH) | |
| input_dim = int(np.load(str(df["embedding_path"].iloc[0]).replace("\\", "/")).shape[0]) | |
| print(f"Embedding dimension: {input_dim}") | |
| le = LabelEncoder() | |
| df["label_idx"] = le.fit_transform(df["label"]) | |
| print(f"Classes: {le.classes_} (Native=0, Non-Native=1)") | |
| print(f"Total samples: {len(df)}") | |
| weights = compute_class_weight( | |
| "balanced", classes=np.unique(df["label_idx"]), y=df["label_idx"] | |
| ) | |
| class_weights = torch.tensor(weights, dtype=torch.float32).to(device) | |
| # Split by speaker (original_id) so the same speaker is never in both train and val. | |
| speaker_df = df.groupby("original_id")["label_idx"].first().reset_index() | |
| train_speakers, val_speakers = train_test_split( | |
| speaker_df["original_id"], | |
| test_size=0.2, | |
| stratify=speaker_df["label_idx"], | |
| random_state=SEED, | |
| ) | |
| train_df = df[df["original_id"].isin(train_speakers)].copy() | |
| val_df = df[df["original_id"].isin(val_speakers)].copy() | |
| print(f"Train: {len(train_df)} chunks ({len(train_speakers)} speakers), Val: {len(val_df)} chunks ({len(val_speakers)} speakers)") | |
| train_set = EmbeddingDataset(train_df, l2_norm=True) | |
| val_set = EmbeddingDataset(val_df, l2_norm=True) | |
| train_loader = DataLoader( | |
| train_set, | |
| batch_size=BATCH_SIZE, | |
| shuffle=True, | |
| num_workers=0, | |
| pin_memory=(device.type == "cuda"), | |
| drop_last=True, # avoid batch size 1 with BatchNorm in training | |
| ) | |
| val_loader = DataLoader(val_set, batch_size=BATCH_SIZE, shuffle=False, num_workers=0) | |
| model = NativityClassifier( | |
| input_dim=input_dim, | |
| hidden_dims=HIDDEN_DIMS, | |
| output_dim=OUTPUT_DIM, | |
| dropout=DROPOUT, | |
| ).to(device) | |
| criterion = nn.CrossEntropyLoss(weight=class_weights) | |
| optimizer = optim.AdamW(model.parameters(), lr=LR, weight_decay=WEIGHT_DECAY) | |
| scheduler = optim.lr_scheduler.ReduceLROnPlateau( | |
| optimizer, mode="max", factor=0.5, patience=5, min_lr=MIN_LR | |
| ) | |
| best_val_f1 = 0.0 | |
| best_epoch = 0 | |
| patience_counter = 0 | |
| for epoch in range(EPOCHS): | |
| model.train() | |
| total_loss = 0.0 | |
| for batch_x, batch_y in train_loader: | |
| batch_x, batch_y = batch_x.to(device), batch_y.to(device) | |
| optimizer.zero_grad() | |
| logits = model(batch_x) | |
| loss = criterion(logits, batch_y) | |
| loss.backward() | |
| torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) | |
| optimizer.step() | |
| total_loss += loss.item() | |
| train_loss = total_loss / len(train_loader) | |
| # Train metrics (to monitor overfitting: train_F1 >> val_F1 = overfitting) | |
| model.eval() | |
| train_preds, train_labels = [], [] | |
| with torch.no_grad(): | |
| for batch_x, batch_y in train_loader: | |
| batch_x = batch_x.to(device) | |
| preds = model(batch_x).argmax(dim=1).cpu().numpy() | |
| train_preds.extend(preds) | |
| train_labels.extend(batch_y.numpy()) | |
| train_f1 = f1_score(train_labels, train_preds, average="macro", zero_division=0) | |
| # Validation | |
| all_preds = [] | |
| all_labels = [] | |
| val_loss = 0.0 | |
| with torch.no_grad(): | |
| for batch_x, batch_y in val_loader: | |
| batch_x, batch_y = batch_x.to(device), batch_y.to(device) | |
| logits = model(batch_x) | |
| val_loss += criterion(logits, batch_y).item() | |
| preds = logits.argmax(dim=1) | |
| all_preds.extend(preds.cpu().numpy()) | |
| all_labels.extend(batch_y.cpu().numpy()) | |
| all_preds = np.array(all_preds) | |
| all_labels = np.array(all_labels) | |
| val_loss /= len(val_loader) | |
| acc = accuracy_score(all_labels, all_preds) | |
| f1 = f1_score(all_labels, all_preds, average="macro", zero_division=0) | |
| prec = precision_score(all_labels, all_preds, average="macro", zero_division=0) | |
| rec = recall_score(all_labels, all_preds, average="macro", zero_division=0) | |
| scheduler.step(f1) | |
| if f1 > best_val_f1: | |
| best_val_f1 = f1 | |
| best_epoch = epoch + 1 | |
| patience_counter = 0 | |
| torch.save(model.state_dict(), BEST_MODEL_PATH) | |
| print(f" -> New best model saved (val F1={f1:.4f})") | |
| else: | |
| patience_counter += 1 | |
| print( | |
| f"Epoch {epoch + 1:3d} | loss {train_loss:.4f} | " | |
| f"train_F1 {train_f1:.4f} | val_F1 {f1:.4f} | val_acc {acc:.4f} | " | |
| f"(gap={(train_f1-f1):.3f})" | |
| ) | |
| if patience_counter >= PATIENCE: | |
| print(f"Early stopping at epoch {epoch + 1} (no improvement for {PATIENCE} epochs).") | |
| break | |
| print(f"\nBest val F1: {best_val_f1:.4f} (epoch {best_epoch})") | |
| # Load best and print final val metrics + confusion matrix | |
| model.load_state_dict(torch.load(BEST_MODEL_PATH, map_location=device)) | |
| model.eval() | |
| all_preds = [] | |
| all_labels = [] | |
| with torch.no_grad(): | |
| for batch_x, batch_y in val_loader: | |
| batch_x = batch_x.to(device) | |
| logits = model(batch_x) | |
| all_preds.extend(logits.argmax(dim=1).cpu().numpy()) | |
| all_labels.extend(batch_y.numpy()) | |
| all_preds = np.array(all_preds) | |
| all_labels = np.array(all_labels) | |
| print("\nFinal validation metrics (best model):") | |
| print(f" Accuracy: {accuracy_score(all_labels, all_preds):.4f}") | |
| print(f" F1 macro: {f1_score(all_labels, all_preds, average='macro', zero_division=0):.4f}") | |
| print(f" Precision: {precision_score(all_labels, all_preds, average='macro', zero_division=0):.4f}") | |
| print(f" Recall: {recall_score(all_labels, all_preds, average='macro', zero_division=0):.4f}") | |
| print(" Confusion matrix (rows=true, cols=pred):") | |
| cm = confusion_matrix(all_labels, all_preds) | |
| print(f" {cm}") | |
| print(f" (Native=0, Non-Native=1)") | |
| print(f"\nModel saved to {BEST_MODEL_PATH}") | |
| if __name__ == "__main__": | |
| main() | |