Spaces:
Sleeping
Sleeping
| """ | |
| Run MLP and MLP+attention on test_data.csv; report metrics for both + ensemble. | |
| Requires: manifests/test_data.csv with nativity_status, both models in models/ | |
| """ | |
| import os | |
| import sys | |
| import pandas as pd | |
| import numpy as np | |
| import torch | |
| import librosa | |
| import requests | |
| import io | |
| from transformers import Wav2Vec2Model, AutoFeatureExtractor | |
| from tqdm import tqdm | |
| from sklearn.metrics import accuracy_score, f1_score, precision_score, recall_score, confusion_matrix | |
| sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) | |
| from model import NativityClassifier, NativityClassifierWithAttention | |
| from extract_embeddings import get_mfcc_from_audio, get_pitch_stats, PROSODY_EMBED_DIM, EMBED_DIM_WITH_MFCC | |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| input_csv = "manifests/test_data.csv" | |
| HIDDEN_DIMS = (384, 256, 128) | |
| DROPOUT = 0.35 | |
| EVAL_THRESHOLD_DEFAULT = 0.35 # fallback if best_threshold.txt missing | |
| PROSODY_LAYERS = (4, 6, 8, 10) | |
| BEST_THRESHOLD_PATH = "models/best_threshold.txt" | |
| CHUNK_SIZE = 160000 | |
| EMBED_DIM_FILE = "manifests/embedding_dim.txt" | |
| THRESHOLD_SWEEP_TEST = [0.25, 0.28, 0.30, 0.32] # report test F1 at these thresholds | |
| # Load wav2vec once | |
| model_name = "facebook/wav2vec2-large-xlsr-53" | |
| feature_extractor = AutoFeatureExtractor.from_pretrained(model_name) | |
| wav2vec = Wav2Vec2Model.from_pretrained(model_name).to(device) | |
| wav2vec.eval() | |
| def get_embedding(audio_np, embed_dim=PROSODY_EMBED_DIM): | |
| """Return 2051-d (mean+std+pitch) or 2090-d (+ MFCC) to match training.""" | |
| with torch.no_grad(): | |
| input_values = feature_extractor( | |
| audio_np, return_tensors="pt", sampling_rate=16000 | |
| ).input_values.to(device) | |
| outputs = wav2vec(input_values, output_hidden_states=True) | |
| h = outputs.hidden_states | |
| num_layers = len(h) - 1 | |
| layer_idx = [min(max(l, 1), num_layers) for l in PROSODY_LAYERS] | |
| embs = [h[i] for i in layer_idx] | |
| time_seq = torch.stack(embs, dim=0).mean(dim=0).squeeze(0) # (T, 1024) | |
| emb_mean = time_seq.mean(dim=0).cpu().numpy().astype(np.float32) | |
| emb_std = time_seq.std(dim=0).cpu().numpy().astype(np.float32) | |
| np.nan_to_num(emb_std, copy=False, nan=0.0, posinf=0.0, neginf=0.0) | |
| pitch = get_pitch_stats(audio_np, sr=16000) | |
| out = np.concatenate([emb_mean, emb_std, pitch]).astype(np.float32) | |
| if embed_dim == EMBED_DIM_WITH_MFCC: | |
| mfcc_39 = get_mfcc_from_audio(audio_np, sr=16000) | |
| out = np.concatenate([out, mfcc_39]).astype(np.float32) | |
| n = np.linalg.norm(out) | |
| if n > 1e-8: | |
| out = out / n | |
| return out | |
| def load_audio(url): | |
| """Load audio from URL; return None if URL returns HTML or unreadable format.""" | |
| try: | |
| r = requests.get(url, timeout=15) | |
| r.raise_for_status() | |
| raw = r.content | |
| if not raw or len(raw) < 100: | |
| return None | |
| ct = (r.headers.get("Content-Type") or "").lower() | |
| if "text/html" in ct or raw.lstrip()[:50].startswith(b"<!") or raw.lstrip()[:20].startswith(b"<?"): | |
| return None | |
| buf = io.BytesIO(raw) | |
| y, _ = librosa.load(buf, sr=16000, mono=True) | |
| return y | |
| except Exception as e: | |
| return None | |
| def predict_file(embs, model_mlp, model_att, lgb_model, threshold): | |
| """embs: (N, D) with D=2051 or 2090. Returns (pred_mlp, pred_att, pred_ens, pred_lgb). Ensemble = avg of MLP + attention + LGB (if present).""" | |
| with torch.no_grad(): | |
| embs_t = torch.tensor(embs, dtype=torch.float32).to(device) | |
| logits_mlp = model_mlp(embs_t) | |
| logits_att = model_att(embs_t) | |
| prob_mlp = torch.softmax(logits_mlp, dim=1).cpu().numpy() | |
| prob_att = torch.softmax(logits_att, dim=1).cpu().numpy() | |
| p1_mlp = prob_mlp[:, 1].mean() | |
| p1_att = prob_att[:, 1].mean() | |
| pred_mlp = 1 if p1_mlp >= threshold else 0 | |
| pred_att = 1 if p1_att >= threshold else 0 | |
| mean_emb = embs.mean(axis=0).astype(np.float32).reshape(1, -1) | |
| p1_lgb = None | |
| pred_lgb = None | |
| if lgb_model is not None: | |
| prob_lgb = lgb_model.predict(mean_emb) # (1, 2) | |
| p1_lgb = float(prob_lgb[0, 1]) | |
| pred_lgb = 1 if p1_lgb >= threshold else 0 | |
| probs_list = [p1_mlp, p1_att] | |
| if p1_lgb is not None: | |
| probs_list.append(p1_lgb) | |
| p1_ens = np.mean(probs_list) | |
| pred_ens = 1 if p1_ens >= threshold else 0 | |
| return pred_mlp, pred_att, pred_ens, pred_lgb, p1_ens | |
| def main(): | |
| df = pd.read_csv(input_csv) | |
| if "nativity_status" not in df.columns: | |
| print("test_data.csv must have 'nativity_status' column") | |
| return | |
| # Match training: 2051 (prosody) or 2090 (prosody+MFCC) | |
| embed_dim = PROSODY_EMBED_DIM | |
| if os.path.isfile(EMBED_DIM_FILE): | |
| with open(EMBED_DIM_FILE) as f: | |
| embed_dim = int(f.read().strip()) | |
| print(f"Embedding dimension: {embed_dim}") | |
| model_mlp = NativityClassifier( | |
| input_dim=embed_dim, hidden_dims=HIDDEN_DIMS, output_dim=2, dropout=DROPOUT | |
| ).to(device) | |
| model_mlp.load_state_dict(torch.load("models/nativity_model.pth", map_location=device)) | |
| model_mlp.eval() | |
| model_att = NativityClassifierWithAttention( | |
| input_dim=embed_dim, hidden_dims=HIDDEN_DIMS, output_dim=2, dropout=DROPOUT, | |
| final_dropout=0.65, reduction_ratio=8 | |
| ).to(device) | |
| model_att.load_state_dict(torch.load("models/nativity_model_attention.pth", map_location=device)) | |
| model_att.eval() | |
| # Best threshold from validation (saved by train_classifier_compare.py) | |
| threshold = EVAL_THRESHOLD_DEFAULT | |
| if os.path.isfile(BEST_THRESHOLD_PATH): | |
| with open(BEST_THRESHOLD_PATH) as f: | |
| threshold = float(f.read().strip()) | |
| print(f"Using threshold from validation: {threshold}") | |
| else: | |
| print(f"Using default threshold: {threshold}") | |
| lgb_model = None | |
| if os.path.isfile("models/lgbm_nativity.txt"): | |
| try: | |
| import lightgbm as lgb | |
| lgb_model = lgb.Booster(model_file="models/lgbm_nativity.txt") | |
| except Exception: | |
| pass | |
| preds_mlp, preds_att, preds_ens, preds_lgb, preds_probs_ens = [], [], [], [], [] | |
| valid_idx = [] | |
| skipped = 0 | |
| print("Running on test set...") | |
| for idx, row in tqdm(df.iterrows(), total=len(df)): | |
| y = load_audio(row["audio_url"]) | |
| if y is None or len(y) == 0: | |
| skipped += 1 | |
| continue | |
| chunks = [ | |
| y[i : i + CHUNK_SIZE] | |
| for i in range(0, len(y), CHUNK_SIZE) | |
| if len(y[i : i + CHUNK_SIZE]) == CHUNK_SIZE | |
| ] | |
| if not chunks: | |
| continue | |
| embs = np.stack([get_embedding(c, embed_dim=embed_dim) for c in chunks]) | |
| p_mlp, p_att, p_ens, p_lgb, p1_ens = predict_file(embs, model_mlp, model_att, lgb_model, threshold) | |
| preds_mlp.append(p_mlp) | |
| preds_att.append(p_att) | |
| preds_ens.append(p_ens) | |
| preds_probs_ens.append(p1_ens) | |
| if p_lgb is not None: | |
| preds_lgb.append(p_lgb) | |
| valid_idx.append(idx) | |
| if not valid_idx: | |
| print("No valid test samples.") | |
| return | |
| if skipped > 0: | |
| print(f"Skipped {skipped} URL(s) (non-audio or load failed).") | |
| y_true = np.array([1 if df.loc[i, "nativity_status"] == "Non-Native" else 0 for i in valid_idx]) | |
| preds_mlp = np.array(preds_mlp) | |
| preds_att = np.array(preds_att) | |
| preds_ens = np.array(preds_ens) | |
| probs_ens = np.array(preds_probs_ens) | |
| preds_lgb = np.array(preds_lgb) if len(preds_lgb) == len(y_true) else None | |
| n_ens = 2 + (1 if preds_lgb is not None else 0) | |
| def row(name, pred): | |
| return { | |
| "name": name, | |
| "acc": accuracy_score(y_true, pred), | |
| "f1": f1_score(y_true, pred, average="macro", zero_division=0), | |
| "prec": precision_score(y_true, pred, average="macro", zero_division=0), | |
| "rec": recall_score(y_true, pred, average="macro", zero_division=0), | |
| "cm": confusion_matrix(y_true, pred), | |
| } | |
| results = [ | |
| row("MLP (baseline)", preds_mlp), | |
| row("MLP + attention (prosody)", preds_att), | |
| row(f"Ensemble ({n_ens} models, t={threshold})", preds_ens), | |
| ] | |
| if preds_lgb is not None: | |
| results.append(row("LightGBM", preds_lgb)) | |
| print("\n" + "=" * 60) | |
| print("COMPARISON ON TEST DATA") | |
| print("=" * 60) | |
| print(f"Test samples: {len(y_true)} (Native=0, Non-Native=1)") | |
| print(f"{'Model':<30} {'Accuracy':>10} {'F1 macro':>10} {'Precision':>10} {'Recall':>10}") | |
| print("-" * 72) | |
| for r in results: | |
| print(f"{r['name']:<30} {r['acc']:>10.4f} {r['f1']:>10.4f} {r['prec']:>10.4f} {r['rec']:>10.4f}") | |
| print("\nConfusion matrices (rows=true, cols=pred):") | |
| for r in results: | |
| print(f" {r['name']}:") | |
| print(f" {r['cm']}") | |
| best = max(results, key=lambda x: x["f1"]) | |
| print(f"\nBest on test (F1): {best['name']} (F1={best['f1']:.4f})") | |
| # Threshold sweep on test (ensemble): see which threshold gives best test F1 | |
| print("\nEnsemble threshold sweep on test:") | |
| print(f" {'t':>6} {'Accuracy':>10} {'F1 macro':>10}") | |
| best_t, best_t_f1 = threshold, 0.0 | |
| for t in THRESHOLD_SWEEP_TEST: | |
| pred_t = (probs_ens >= t).astype(int) | |
| acc_t = accuracy_score(y_true, pred_t) | |
| f1_t = f1_score(y_true, pred_t, average="macro", zero_division=0) | |
| print(f" {t:>6.2f} {acc_t:>10.4f} {f1_t:>10.4f}") | |
| if f1_t > best_t_f1: | |
| best_t_f1 = f1_t | |
| best_t = t | |
| print(f" Best on test: t={best_t} (F1={best_t_f1:.4f})") | |
| if __name__ == "__main__": | |
| main() | |