Download webapp/app.py from neev2211/arabic-native-classifier: direct link, hf CLI and curl.
- Browser
- Download file 35.6 kB
-
https://huggingface.co/spaces/neev2211/arabic-native-classifier/resolve/main/webapp/app.py
- Command line
-
hf download hf://spaces/neev2211/arabic-native-classifier/webapp/app.py
-
curl -L -o app.py https://huggingface.co/spaces/neev2211/arabic-native-classifier/resolve/main/webapp/app.py
35.6 kB
| """ | |
| Flask web interface for Arabic Native vs Non-Native audio classification. | |
| Supports: CSV upload, custom CSV column mapping, and direct audio file upload. | |
| Pipeline: Audio -> Chunk (10s) -> Wav2Vec2 (prosody layers) -> MLP Classifier -> Results | |
| Uses Server-Sent Events (SSE) for real-time pipeline progress. | |
| Includes admin panel for model training. | |
| """ | |
| import os | |
| import sys | |
| import io | |
| import uuid | |
| import json | |
| import time | |
| import traceback | |
| import threading | |
| import hashlib | |
| import functools | |
| import tempfile | |
| import zipfile | |
| import numpy as np | |
| import pandas as pd | |
| import torch | |
| import librosa | |
| import requests | |
| import soundfile as sf | |
| import joblib | |
| from flask import Flask, render_template, request, jsonify, Response, stream_with_context, redirect, url_for, session, send_file | |
| from flask_cors import CORS | |
| from werkzeug.utils import secure_filename | |
| from sklearn.metrics import ( | |
| accuracy_score, | |
| f1_score, | |
| precision_score, | |
| recall_score, | |
| confusion_matrix, | |
| ) | |
| from sklearn.decomposition import PCA | |
| from sklearn.preprocessing import LabelEncoder | |
| from sklearn.model_selection import train_test_split | |
| # Add scripts/ to path for model imports | |
| SCRIPT_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "scripts") | |
| sys.path.insert(0, SCRIPT_DIR) | |
| from model import NativityClassifier | |
| from extract_embeddings import get_mfcc_from_audio, get_pitch_stats, PROSODY_EMBED_DIM, EMBED_DIM_WITH_MFCC | |
| # ============================================================================= | |
| # Config | |
| # ============================================================================= | |
| BASE_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..") | |
| MODELS_DIR = os.path.join(BASE_DIR, "models") | |
| MANIFESTS_DIR = os.path.join(BASE_DIR, "manifests") | |
| UPLOAD_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "uploads") | |
| os.makedirs(UPLOAD_DIR, exist_ok=True) | |
| EMBED_DIM_FILE = os.path.join(MANIFESTS_DIR, "embedding_dim.txt") | |
| BEST_THRESHOLD_PATH = os.path.join(MODELS_DIR, "best_threshold.txt") | |
| MODEL_PATH = os.path.join(MODELS_DIR, "nativity_model.pth") | |
| HIDDEN_DIMS = (256, 128) # Must match train_classifier_compare.py | |
| DROPOUT = 0.4 | |
| PCA_COMPONENTS = 256 | |
| PCA_PATH = os.path.join(MODELS_DIR, "pca_transform.pkl") | |
| EVAL_THRESHOLD_DEFAULT = 0.35 | |
| PROSODY_LAYERS = (4, 6, 8, 10) | |
| SUPPORTED_LANGUAGES = ["Arabic_SA", "Arabic_AE", "Arabic_QA", "Arabic_KW"] | |
| # Admin config | |
| ADMIN_PASSWORD = os.environ.get("ADMIN_PASSWORD", "hackenza2026") | |
| EMBEDDINGS_DIR = os.path.join(BASE_DIR, "embeddings") | |
| AUDIO_CHUNKS_DIR = os.path.join(BASE_DIR, "data", "processed_audio_chunks") | |
| os.makedirs(EMBEDDINGS_DIR, exist_ok=True) | |
| os.makedirs(AUDIO_CHUNKS_DIR, exist_ok=True) | |
| # Training state (shared across threads) | |
| _training_state = { | |
| "active": False, | |
| "logs": [], | |
| "progress": 0, | |
| "total": 0, | |
| "phase": "", | |
| "metrics": None, | |
| "error": None, | |
| "done": False, | |
| } | |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| embedding_dim = PROSODY_EMBED_DIM | |
| if os.path.isfile(EMBED_DIM_FILE): | |
| with open(EMBED_DIM_FILE) as f: | |
| embedding_dim = int(f.read().strip()) | |
| EVAL_THRESHOLD = EVAL_THRESHOLD_DEFAULT | |
| if os.path.isfile(BEST_THRESHOLD_PATH): | |
| with open(BEST_THRESHOLD_PATH) as f: | |
| EVAL_THRESHOLD = float(f.read().strip()) | |
| # ============================================================================= | |
| # Lazy model loading | |
| # ============================================================================= | |
| _wav2vec_model = None | |
| _feature_extractor = None | |
| _classifier = None | |
| _pca = None | |
| _models_loaded = False | |
| def get_wav2vec(): | |
| global _wav2vec_model, _feature_extractor | |
| if _wav2vec_model is None: | |
| from transformers import Wav2Vec2Model, AutoFeatureExtractor | |
| # MUST match extract_embeddings.py model used for training embeddings | |
| model_name = "facebook/wav2vec2-lv-60-espeak-cv-ft" | |
| _feature_extractor = AutoFeatureExtractor.from_pretrained(model_name) | |
| _wav2vec_model = Wav2Vec2Model.from_pretrained(model_name).to(device) | |
| _wav2vec_model.eval() | |
| return _feature_extractor, _wav2vec_model | |
| def get_pca(): | |
| global _pca | |
| if _pca is None: | |
| import joblib | |
| _pca = joblib.load(PCA_PATH) | |
| return _pca | |
| def get_classifier(): | |
| global _classifier | |
| if _classifier is None: | |
| _classifier = NativityClassifier( | |
| input_dim=PCA_COMPONENTS, hidden_dims=HIDDEN_DIMS, | |
| output_dim=2, dropout=DROPOUT, | |
| ).to(device) | |
| _classifier.load_state_dict(torch.load(MODEL_PATH, map_location=device)) | |
| _classifier.eval() | |
| return _classifier | |
| def ensure_models_loaded(): | |
| global _models_loaded | |
| if not _models_loaded: | |
| get_wav2vec() | |
| get_pca() | |
| get_classifier() | |
| _models_loaded = True | |
| # ============================================================================= | |
| # Pipeline | |
| # ============================================================================= | |
| def get_embedding(audio_np, embed_dim=None): | |
| if embed_dim is None: | |
| embed_dim = embedding_dim | |
| fe, model = get_wav2vec() | |
| with torch.no_grad(): | |
| iv = fe(audio_np, return_tensors="pt", sampling_rate=16000).input_values.to(device) | |
| out = model(iv, output_hidden_states=True) | |
| h = out.hidden_states | |
| nl = len(h) - 1 | |
| li = [min(max(l, 1), nl) for l in PROSODY_LAYERS] | |
| ts = torch.stack([h[i] for i in li], dim=0).mean(dim=0).squeeze(0) | |
| em = ts.mean(dim=0).cpu().numpy().astype(np.float32) | |
| es = ts.std(dim=0).cpu().numpy().astype(np.float32) | |
| np.nan_to_num(es, copy=False, nan=0.0, posinf=0.0, neginf=0.0) | |
| pitch = get_pitch_stats(audio_np, sr=16000) | |
| vec = np.concatenate([em, es, pitch]).astype(np.float32) | |
| if embed_dim == EMBED_DIM_WITH_MFCC: | |
| vec = np.concatenate([vec, get_mfcc_from_audio(audio_np, sr=16000)]).astype(np.float32) | |
| n = np.linalg.norm(vec) | |
| if n > 1e-8: | |
| vec = vec / n | |
| # Apply PCA to match training pipeline (2051-d -> 256-d) | |
| pca = get_pca() | |
| vec_pca = pca.transform(vec.reshape(1, -1)).astype(np.float32) | |
| return torch.tensor(vec_pca, dtype=torch.float32, device=device) | |
| def chunk_audio(y, sr=16000, chunk_sec=10): | |
| cs = sr * chunk_sec | |
| chunks = [y[i:i+cs] for i in range(0, len(y), cs) if len(y[i:i+cs]) == cs] | |
| if not chunks and len(y) >= sr: | |
| chunks = [y] | |
| return chunks | |
| def classify_chunks(chunks, sr=16000): | |
| clf = get_classifier() | |
| preds, confs, details = [], [], [] | |
| with torch.no_grad(): | |
| for idx, c in enumerate(chunks): | |
| emb = get_embedding(c, embed_dim=embedding_dim) | |
| logits = clf(emb) | |
| probs = torch.softmax(logits, dim=1) | |
| pn, pnn = probs[0, 0].item(), probs[0, 1].item() | |
| pred = 1 if pnn >= EVAL_THRESHOLD else 0 | |
| conf = pnn if pred == 1 else pn | |
| preds.append(pred) | |
| confs.append(conf) | |
| details.append({ | |
| "chunk_idx": idx, "duration_sec": round(len(c)/sr, 1), | |
| "prediction": "Non-Native" if pred == 1 else "Native", | |
| "confidence": round(conf, 4), | |
| "prob_native": round(pn, 4), "prob_non_native": round(pnn, 4), | |
| }) | |
| w = np.array(confs, dtype=np.float64) | |
| w = w / (w.sum() + 1e-10) | |
| ws = np.sum(np.array(preds) * w) | |
| return { | |
| "prediction": "Non-Native" if ws >= EVAL_THRESHOLD else "Native", | |
| "confidence": round(float(np.average(confs, weights=w)), 4), | |
| "num_chunks": len(chunks), "threshold_used": EVAL_THRESHOLD, | |
| "weighted_score": round(float(ws), 4), "chunk_details": details, | |
| } | |
| def load_audio_from_url(url): | |
| try: | |
| r = requests.get(url, timeout=30) | |
| r.raise_for_status() | |
| raw = r.content | |
| if not raw or len(raw) < 100: | |
| return None, "Empty response" | |
| ct = (r.headers.get("Content-Type") or "").lower() | |
| if "text/html" in ct or raw.lstrip()[:50].startswith(b"<!"): | |
| return None, "URL returned HTML" | |
| y, sr = librosa.load(io.BytesIO(raw), sr=16000, mono=True) | |
| return y, None | |
| except Exception as e: | |
| return None, str(e) | |
| def load_audio_from_file(fp): | |
| try: | |
| y, sr = librosa.load(fp, sr=16000, mono=True) | |
| return y, None | |
| except Exception as e: | |
| return None, str(e) | |
| def compute_metrics(y_true, y_pred, y_confs): | |
| fn = lambda s: 1 if s == "Non-Native" else 0 | |
| yt, yp = np.array([fn(s) for s in y_true]), np.array([fn(s) for s in y_pred]) | |
| cm = confusion_matrix(yt, yp).tolist() | |
| return { | |
| "accuracy": round(accuracy_score(yt, yp), 4), | |
| "f1_macro": round(f1_score(yt, yp, average="macro", zero_division=0), 4), | |
| "f1_weighted": round(f1_score(yt, yp, average="weighted", zero_division=0), 4), | |
| "precision_macro": round(precision_score(yt, yp, average="macro", zero_division=0), 4), | |
| "recall_macro": round(recall_score(yt, yp, average="macro", zero_division=0), 4), | |
| "avg_confidence": round(float(np.mean(y_confs)), 4) if y_confs else 0.0, | |
| "confusion_matrix": cm, | |
| "per_class": { | |
| "Native": {"f1": round(f1_score(yt, yp, pos_label=0, zero_division=0), 4), "precision": round(precision_score(yt, yp, pos_label=0, zero_division=0), 4), "recall": round(recall_score(yt, yp, pos_label=0, zero_division=0), 4)}, | |
| "Non-Native": {"f1": round(f1_score(yt, yp, pos_label=1, zero_division=0), 4), "precision": round(precision_score(yt, yp, pos_label=1, zero_division=0), 4), "recall": round(recall_score(yt, yp, pos_label=1, zero_division=0), 4)}, | |
| }, | |
| "total_samples": len(yt), "native_count": int((yt==0).sum()), | |
| "non_native_count": int((yt==1).sum()), "correct_predictions": int((yt==yp).sum()), | |
| } | |
| # ============================================================================= | |
| # Flask | |
| # ============================================================================= | |
| app = Flask(__name__) | |
| CORS(app) | |
| app.config["MAX_CONTENT_LENGTH"] = 500 * 1024 * 1024 | |
| app.secret_key = os.environ.get("FLASK_SECRET_KEY", "nativity-admin-secret-key-change-me") | |
| _pending_jobs = {} | |
| def index(): | |
| return render_template("index.html") | |
| def health(): | |
| return jsonify({"status": "ok", "device": str(device), "embedding_dim": embedding_dim, "threshold": EVAL_THRESHOLD, "languages": SUPPORTED_LANGUAGES}) | |
| def upload_csv(): | |
| """Upload CSV, get job_id for SSE streaming.""" | |
| try: | |
| if "file" not in request.files: | |
| return jsonify({"error": "No CSV file uploaded"}), 400 | |
| f = request.files["file"] | |
| if f.filename == "": | |
| return jsonify({"error": "Empty filename"}), 400 | |
| df = pd.read_csv(f) | |
| job_id = uuid.uuid4().hex[:12] | |
| _pending_jobs[job_id] = { | |
| "df": df, | |
| "col_id": request.form.get("col_id", "").strip(), | |
| "col_audio": request.form.get("col_audio", "").strip(), | |
| "col_nativity": request.form.get("col_nativity", "").strip(), | |
| "col_language": request.form.get("col_language", "").strip(), | |
| } | |
| return jsonify({"job_id": job_id, "total": len(df), "columns": list(df.columns)}) | |
| except Exception as e: | |
| traceback.print_exc() | |
| return jsonify({"error": str(e)}), 500 | |
| def stream_csv(job_id): | |
| """SSE endpoint for step-by-step CSV progress.""" | |
| job = _pending_jobs.pop(job_id, None) | |
| if not job: | |
| return jsonify({"error": "Job not found"}), 404 | |
| def sse(data): | |
| return f"data: {json.dumps(data)}\n\n" | |
| def generate(): | |
| df = job["df"] | |
| cols = list(df.columns) | |
| ca, ci, cn, cl = job["col_audio"], job["col_id"], job["col_nativity"], job["col_language"] | |
| audio_col = ca if ca and ca in cols else None | |
| id_col = ci if ci and ci in cols else None | |
| nat_col = cn if cn and cn in cols else None | |
| lang_col = cl if cl and cl in cols else None | |
| if not audio_col: | |
| for c in cols: | |
| if any(k in c.lower() for k in ["audio", "url", "path"]): | |
| audio_col = c; break | |
| if not audio_col: | |
| yield sse({"type": "error", "message": "No audio URL column found in: " + str(cols)}) | |
| return | |
| if not id_col: | |
| for c in cols: | |
| if "id" in c.lower(): id_col = c; break | |
| if not nat_col: | |
| for c in cols: | |
| if any(k in c.lower() for k in ["nativity", "native", "label"]): nat_col = c; break | |
| if not lang_col: | |
| for c in cols: | |
| if "lang" in c.lower(): lang_col = c; break | |
| has_gt = nat_col and nat_col in cols | |
| total = len(df) | |
| yield sse({"type": "init", "total": total, "cols": {"id": id_col, "audio": audio_col, "nativity": nat_col, "language": lang_col}}) | |
| yield sse({"type": "step", "message": "Loading Wav2Vec2 encoder & classifier..."}) | |
| ensure_models_loaded() | |
| yield sse({"type": "step", "message": "Models ready. Starting predictions..."}) | |
| results, errors = [], [] | |
| y_true, y_pred, y_confs = [], [], [] | |
| for idx, row in df.iterrows(): | |
| i = int(idx) if not isinstance(idx, int) else idx | |
| url = str(row[audio_col]).strip() | |
| sid = str(row[id_col]) if id_col and id_col in cols else str(i) | |
| gt = str(row[nat_col]).strip() if has_gt else None | |
| lang = str(row[lang_col]).strip() if lang_col and lang_col in cols else "Unknown" | |
| yield sse({"type": "progress", "n": i+1, "total": total, "id": sid, "phase": "download", "msg": f"Downloading audio for #{sid}..."}) | |
| audio_np, err = load_audio_from_url(url) | |
| if err: | |
| errors.append(sid) | |
| results.append({"id": sid, "language": lang, "ground_truth": gt, "prediction": "Error", "confidence": 0, "error": err}) | |
| yield sse({"type": "sample_error", "n": i+1, "id": sid, "error": err}) | |
| continue | |
| dur = round(len(audio_np) / 16000, 1) | |
| chunks = chunk_audio(audio_np) | |
| nc = len(chunks) | |
| yield sse({"type": "progress", "n": i+1, "total": total, "id": sid, "phase": "chunk", "msg": f"#{sid}: {dur}s audio — chunked into {nc} segment(s)"}) | |
| if not chunks: | |
| errors.append(sid) | |
| results.append({"id": sid, "language": lang, "ground_truth": gt, "prediction": "Error", "confidence": 0, "error": "Too short"}) | |
| yield sse({"type": "sample_error", "n": i+1, "id": sid, "error": "Too short"}) | |
| continue | |
| yield sse({"type": "progress", "n": i+1, "total": total, "id": sid, "phase": "embed", "msg": f"#{sid}: Extracting Wav2Vec2 embeddings ({nc} chunks)..."}) | |
| res = classify_chunks(chunks) | |
| is_correct = None | |
| if gt and gt in ("Native", "Non-Native"): | |
| is_correct = res["prediction"] == gt | |
| y_true.append(gt); y_pred.append(res["prediction"]); y_confs.append(res["confidence"]) | |
| row_r = {"id": sid, "language": lang, "ground_truth": gt, "prediction": res["prediction"], | |
| "confidence": res["confidence"], "num_chunks": res["num_chunks"], | |
| "weighted_score": res["weighted_score"], "is_correct": is_correct, "chunk_details": res["chunk_details"]} | |
| results.append(row_r) | |
| yield sse({"type": "sample_done", "n": i+1, "total": total, "result": row_r}) | |
| metrics = compute_metrics(y_true, y_pred, y_confs) if y_true else None | |
| yield sse({"type": "done", "total": total, "processed": len(results)-len(errors), | |
| "errors_count": len(errors), "has_ground_truth": bool(has_gt), | |
| "metrics": metrics, "results": results}) | |
| return Response(stream_with_context(generate()), content_type="text/event-stream") | |
| def predict_audio(): | |
| try: | |
| if "audio" not in request.files: | |
| return jsonify({"error": "No audio file"}), 400 | |
| f = request.files["audio"] | |
| if f.filename == "": | |
| return jsonify({"error": "Empty filename"}), 400 | |
| lang = request.form.get("language", "").strip() or "Unknown" | |
| ext = os.path.splitext(f.filename)[1] or ".wav" | |
| tmp = os.path.join(UPLOAD_DIR, f"{uuid.uuid4().hex}{ext}") | |
| f.save(tmp) | |
| try: | |
| y, err = load_audio_from_file(tmp) | |
| if err: | |
| return jsonify({"error": err}), 400 | |
| dur = round(len(y)/16000, 1) | |
| chunks = chunk_audio(y) | |
| if not chunks: | |
| return jsonify({"error": "Audio too short"}), 400 | |
| res = classify_chunks(chunks) | |
| return jsonify({"success": True, "filename": f.filename, "language": lang, | |
| "duration_sec": dur, **res}) | |
| finally: | |
| if os.path.exists(tmp): os.remove(tmp) | |
| except Exception as e: | |
| traceback.print_exc() | |
| return jsonify({"error": str(e)}), 500 | |
| def predict_url(): | |
| try: | |
| data = request.get_json() | |
| if not data or "audio_url" not in data: | |
| return jsonify({"error": "Missing audio_url"}), 400 | |
| url = data["audio_url"].strip() | |
| lang = data.get("language", "").strip() or "Unknown" | |
| y, err = load_audio_from_url(url) | |
| if err: | |
| return jsonify({"error": err}), 400 | |
| dur = round(len(y)/16000, 1) | |
| chunks = chunk_audio(y) | |
| if not chunks: | |
| return jsonify({"error": "Audio too short"}), 400 | |
| res = classify_chunks(chunks) | |
| return jsonify({"success": True, "audio_url": url, "language": lang, | |
| "duration_sec": dur, **res}) | |
| except Exception as e: | |
| traceback.print_exc() | |
| return jsonify({"error": str(e)}), 500 | |
| # ============================================================================= | |
| # Admin Panel | |
| # ============================================================================= | |
| def admin_required(f): | |
| def decorated(*args, **kwargs): | |
| if not session.get("admin_authenticated"): | |
| return redirect(url_for("admin_login")) | |
| return f(*args, **kwargs) | |
| return decorated | |
| def admin_login(): | |
| if request.method == "POST": | |
| pw = request.form.get("password", "") | |
| if pw == ADMIN_PASSWORD: | |
| session["admin_authenticated"] = True | |
| return redirect(url_for("admin_page")) | |
| return render_template("admin_login.html", error="Invalid password") | |
| return render_template("admin_login.html", error=None) | |
| def admin_logout(): | |
| session.pop("admin_authenticated", None) | |
| return redirect(url_for("index")) | |
| def admin_page(): | |
| return render_template("admin.html") | |
| def admin_status(): | |
| """Current training state.""" | |
| return jsonify({ | |
| "active": _training_state["active"], | |
| "done": _training_state["done"], | |
| "phase": _training_state["phase"], | |
| "progress": _training_state["progress"], | |
| "total": _training_state["total"], | |
| "logs": _training_state["logs"][-100:], | |
| "metrics": _training_state["metrics"], | |
| "error": _training_state["error"], | |
| }) | |
| def _log(msg): | |
| _training_state["logs"].append({"time": time.strftime("%H:%M:%S"), "msg": msg}) | |
| print(f"[TRAIN] {msg}") | |
| def _training_worker(csv_path): | |
| """Background training pipeline: download → chunk → embed → PCA → train.""" | |
| global _classifier, _pca, _models_loaded | |
| try: | |
| import torch.nn as nn | |
| import torch.optim as optim | |
| from model import NativityClassifier, NativityClassifierWithAttention | |
| SEED = 42 | |
| np.random.seed(SEED) | |
| torch.manual_seed(SEED) | |
| _training_state["phase"] = "loading_csv" | |
| _log("Reading CSV...") | |
| df = pd.read_csv(csv_path) | |
| total_speakers = len(df) | |
| _training_state["total"] = total_speakers | |
| _log(f"Found {total_speakers} audio files") | |
| labels = df["nativity_status"].value_counts().to_dict() | |
| _log(f"Labels: {labels}") | |
| # ---- Phase 1: Download & Chunk ---- | |
| _training_state["phase"] = "downloading" | |
| _log("Phase 1: Downloading audio & chunking...") | |
| preprocessor_sr = 16000 | |
| chunk_duration = 3 | |
| stride_duration = 1 | |
| samples_per_chunk = preprocessor_sr * chunk_duration | |
| samples_per_stride = preprocessor_sr * stride_duration | |
| chunk_records = [] | |
| for idx, row in df.iterrows(): | |
| dp_id = row["dp_id"] | |
| url = str(row["audio_url"]).strip() | |
| _training_state["progress"] = idx + 1 | |
| _log(f"[{idx+1}/{total_speakers}] Downloading dp_id={dp_id}...") | |
| try: | |
| resp = requests.get(url, timeout=30) | |
| resp.raise_for_status() | |
| audio_bytes = io.BytesIO(resp.content) | |
| y, sr = librosa.load(audio_bytes, sr=preprocessor_sr, mono=True) | |
| if len(y) == 0: | |
| _log(f" dp_id={dp_id}: empty audio, skipped") | |
| continue | |
| # Trim silence | |
| if len(y) > samples_per_chunk: | |
| y_trim, _ = librosa.effects.trim(y, top_db=25, frame_length=2048, hop_length=512) | |
| if len(y_trim) >= samples_per_chunk: | |
| y = y_trim | |
| # Overlapping chunks | |
| chunk_count = 0 | |
| for i in range(0, len(y), samples_per_stride): | |
| chunk = y[i:i + samples_per_chunk].copy() | |
| if len(chunk) < int(preprocessor_sr * 2.5): | |
| continue | |
| peak = np.abs(chunk).max() | |
| if peak > 1e-8: | |
| chunk = chunk / peak | |
| chunk_fname = f"{dp_id}_chunk{chunk_count}.wav" | |
| chunk_path = os.path.join(AUDIO_CHUNKS_DIR, chunk_fname) | |
| sf.write(chunk_path, chunk, preprocessor_sr) | |
| chunk_records.append({ | |
| "chunk_path": chunk_path, | |
| "original_id": dp_id, | |
| "label": row["nativity_status"], | |
| "language": row["language"], | |
| }) | |
| chunk_count += 1 | |
| _log(f" dp_id={dp_id}: {chunk_count} chunks") | |
| except Exception as e: | |
| _log(f" dp_id={dp_id}: FAILED - {e}") | |
| if not chunk_records: | |
| _training_state["error"] = "No audio could be downloaded/chunked" | |
| _training_state["done"] = True | |
| _training_state["active"] = False | |
| return | |
| chunk_df = pd.DataFrame(chunk_records) | |
| manifest_path = os.path.join(MANIFESTS_DIR, "processed_manifest.csv") | |
| chunk_df.to_csv(manifest_path, index=False) | |
| _log(f"Total chunks: {len(chunk_df)}") | |
| # ---- Phase 2: Extract Embeddings ---- | |
| _training_state["phase"] = "embedding" | |
| _training_state["progress"] = 0 | |
| _training_state["total"] = len(chunk_df) | |
| _log("Phase 2: Extracting Wav2Vec2 embeddings...") | |
| from extract_embeddings import get_embeddings, PROSODY_EMBED_DIM as EMB_DIM | |
| emb_paths = [] | |
| for eidx, row in chunk_df.iterrows(): | |
| _training_state["progress"] = eidx + 1 | |
| if (eidx + 1) % 50 == 0 or eidx == 0: | |
| _log(f" Embedding {eidx+1}/{len(chunk_df)}...") | |
| emb = get_embeddings(row["chunk_path"], l2_norm=True) | |
| if emb is not None: | |
| emb_fname = os.path.basename(row["chunk_path"]).replace(".wav", ".npy") | |
| emb_path = os.path.join(EMBEDDINGS_DIR, emb_fname) | |
| np.save(emb_path, emb) | |
| emb_paths.append(emb_path) | |
| else: | |
| emb_paths.append(None) | |
| chunk_df["embedding_path"] = emb_paths | |
| chunk_df.dropna(subset=["embedding_path"], inplace=True) | |
| final_manifest = os.path.join(MANIFESTS_DIR, "final_training_manifest.csv") | |
| chunk_df.to_csv(final_manifest, index=False) | |
| _log(f"Embeddings extracted: {len(chunk_df)} valid") | |
| # Save embedding dim | |
| with open(EMBED_DIM_FILE, "w") as f: | |
| f.write(str(EMB_DIM)) | |
| # ---- Phase 3: Train ---- | |
| _training_state["phase"] = "training" | |
| _training_state["progress"] = 0 | |
| _log("Phase 3: Training classifier...") | |
| le = LabelEncoder() | |
| chunk_df["label_idx"] = le.fit_transform(chunk_df["label"]) | |
| speaker_df = chunk_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 = chunk_df[chunk_df["original_id"].isin(train_speakers)].copy() | |
| val_df = chunk_df[chunk_df["original_id"].isin(val_speakers)].copy() | |
| _log(f"Train: {len(train_df)} chunks ({train_df.original_id.nunique()} speakers)") | |
| _log(f"Val: {len(val_df)} chunks ({val_df.original_id.nunique()} speakers)") | |
| # Load embeddings | |
| def load_embs(sub_df): | |
| X, Y = [], [] | |
| for _, r in sub_df.iterrows(): | |
| e = np.load(r["embedding_path"]).astype(np.float32) | |
| n = np.linalg.norm(e) | |
| if n > 1e-8: | |
| e = e / n | |
| X.append(e) | |
| Y.append(r["label_idx"]) | |
| return np.stack(X), np.array(Y) | |
| X_train_raw, y_train = load_embs(train_df) | |
| X_val_raw, y_val = load_embs(val_df) | |
| _log(f"Raw embedding dim: {X_train_raw.shape[1]}") | |
| # PCA | |
| _log(f"Fitting PCA ({X_train_raw.shape[1]}d -> {PCA_COMPONENTS}d)...") | |
| pca = PCA(n_components=PCA_COMPONENTS, random_state=SEED) | |
| X_train = pca.fit_transform(X_train_raw).astype(np.float32) | |
| X_val = pca.transform(X_val_raw).astype(np.float32) | |
| joblib.dump(pca, PCA_PATH) | |
| _log(f"PCA variance explained: {pca.explained_variance_ratio_.sum():.4f}") | |
| # DataLoaders | |
| from torch.utils.data import DataLoader, TensorDataset | |
| from sklearn.utils.class_weight import compute_class_weight | |
| train_set = TensorDataset(torch.tensor(X_train), torch.tensor(y_train, dtype=torch.long)) | |
| val_set = TensorDataset(torch.tensor(X_val), torch.tensor(y_val, dtype=torch.long)) | |
| train_loader = DataLoader(train_set, batch_size=32, shuffle=True, drop_last=True) | |
| val_loader = DataLoader(val_set, batch_size=32, shuffle=False) | |
| weights = compute_class_weight("balanced", classes=np.unique(y_train), y=y_train) | |
| class_weights = torch.tensor(weights, dtype=torch.float32).to(device) | |
| criterion = nn.CrossEntropyLoss(weight=class_weights, label_smoothing=0.15) | |
| # Train MLP | |
| _log("Training MLP baseline...") | |
| model_mlp = NativityClassifier( | |
| input_dim=PCA_COMPONENTS, hidden_dims=HIDDEN_DIMS, | |
| output_dim=2, dropout=DROPOUT, | |
| ).to(device) | |
| optimizer = optim.AdamW(model_mlp.parameters(), lr=1e-3, weight_decay=1e-4) | |
| scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode="max", factor=0.5, patience=5, min_lr=1e-6) | |
| best_f1 = 0.0 | |
| patience_counter = 0 | |
| max_epochs = 100 | |
| patience_limit = 15 | |
| min_epochs = 40 | |
| _training_state["total"] = max_epochs | |
| for epoch in range(max_epochs): | |
| _training_state["progress"] = epoch + 1 | |
| model_mlp.train() | |
| t_loss, t_correct, t_total = 0, 0, 0 | |
| for bx, by in train_loader: | |
| bx, by = bx.to(device), by.to(device) | |
| if model_mlp.training: | |
| bx = bx + torch.randn_like(bx) * 0.02 | |
| optimizer.zero_grad() | |
| logits = model_mlp(bx) | |
| loss = criterion(logits, by) | |
| loss.backward() | |
| torch.nn.utils.clip_grad_norm_(model_mlp.parameters(), 1.0) | |
| optimizer.step() | |
| t_loss += loss.item() * bx.size(0) | |
| t_correct += (logits.argmax(1) == by).sum().item() | |
| t_total += bx.size(0) | |
| model_mlp.eval() | |
| v_preds, v_labels = [], [] | |
| with torch.no_grad(): | |
| for bx, by in val_loader: | |
| bx, by = bx.to(device), by.to(device) | |
| preds = model_mlp(bx).argmax(1) | |
| v_preds.extend(preds.cpu().numpy()) | |
| v_labels.extend(by.cpu().numpy()) | |
| v_f1 = f1_score(v_labels, v_preds, average="macro", zero_division=0) | |
| v_acc = accuracy_score(v_labels, v_preds) | |
| scheduler.step(v_f1) | |
| if (epoch + 1) % 10 == 0 or epoch == 0: | |
| _log(f" Epoch {epoch+1}: loss={t_loss/max(1,t_total):.4f} train_acc={t_correct/max(1,t_total):.4f} val_acc={v_acc:.4f} val_f1={v_f1:.4f}") | |
| if v_f1 > best_f1: | |
| best_f1 = v_f1 | |
| torch.save(model_mlp.state_dict(), MODEL_PATH) | |
| patience_counter = 0 | |
| else: | |
| patience_counter += 1 | |
| if patience_counter >= patience_limit and epoch >= min_epochs: | |
| _log(f" Early stopping at epoch {epoch+1}") | |
| break | |
| _log(f"Best MLP F1: {best_f1:.4f}") | |
| # Threshold sweep | |
| model_mlp.load_state_dict(torch.load(MODEL_PATH, map_location=device)) | |
| model_mlp.eval() | |
| all_probs = [] | |
| with torch.no_grad(): | |
| for bx, _ in val_loader: | |
| bx = bx.to(device) | |
| probs = torch.softmax(model_mlp(bx), dim=1).cpu().numpy() | |
| all_probs.append(probs) | |
| all_probs = np.concatenate(all_probs) | |
| best_thresh = 0.4 | |
| best_f1_t = 0 | |
| for t in [0.25, 0.28, 0.30, 0.32, 0.35, 0.40, 0.45, 0.50]: | |
| p = (all_probs[:, 1] >= t).astype(int) | |
| f = f1_score(y_val, p, average="macro", zero_division=0) | |
| if f > best_f1_t: | |
| best_f1_t = f | |
| best_thresh = t | |
| with open(BEST_THRESHOLD_PATH, "w") as f: | |
| f.write(str(best_thresh)) | |
| _log(f"Best threshold: {best_thresh} (F1={best_f1_t:.4f})") | |
| # Final metrics | |
| final_preds = (all_probs[:, 1] >= best_thresh).astype(int) | |
| acc = accuracy_score(y_val, final_preds) | |
| f1 = f1_score(y_val, final_preds, average="macro", zero_division=0) | |
| prec = precision_score(y_val, final_preds, average="macro", zero_division=0) | |
| rec = recall_score(y_val, final_preds, average="macro", zero_division=0) | |
| cm = confusion_matrix(y_val, final_preds).tolist() | |
| metrics = { | |
| "accuracy": round(acc, 4), "f1_macro": round(f1, 4), | |
| "precision": round(prec, 4), "recall": round(rec, 4), | |
| "confusion_matrix": cm, "threshold": best_thresh, | |
| "train_chunks": len(train_df), "val_chunks": len(val_df), | |
| "total_speakers": chunk_df.original_id.nunique(), | |
| "total_chunks": len(chunk_df), | |
| } | |
| _training_state["metrics"] = metrics | |
| _log(f"Final: Acc={acc:.4f} F1={f1:.4f} Prec={prec:.4f} Rec={rec:.4f}") | |
| _log(f"Confusion matrix: {cm}") | |
| # Reload models in memory | |
| _classifier = None | |
| _pca = None | |
| _models_loaded = False | |
| _log("Models saved. Classifier will reload on next prediction.") | |
| _training_state["phase"] = "complete" | |
| _training_state["done"] = True | |
| _training_state["active"] = False | |
| _log("Training complete!") | |
| except Exception as e: | |
| traceback.print_exc() | |
| _training_state["error"] = str(e) | |
| _training_state["done"] = True | |
| _training_state["active"] = False | |
| _log(f"ERROR: {e}") | |
| def admin_train(): | |
| """Start training from uploaded CSV.""" | |
| if _training_state["active"]: | |
| return jsonify({"error": "Training already in progress"}), 409 | |
| if "file" not in request.files: | |
| return jsonify({"error": "No CSV file uploaded"}), 400 | |
| f = request.files["file"] | |
| if f.filename == "": | |
| return jsonify({"error": "Empty filename"}), 400 | |
| # Save CSV to temp | |
| csv_path = os.path.join(UPLOAD_DIR, f"train_{uuid.uuid4().hex[:8]}.csv") | |
| f.save(csv_path) | |
| # Validate CSV format | |
| try: | |
| df = pd.read_csv(csv_path) | |
| required = {"dp_id", "audio_url", "nativity_status", "language"} | |
| missing = required - set(df.columns) | |
| if missing: | |
| os.remove(csv_path) | |
| return jsonify({"error": f"Missing columns: {missing}. Required: {required}"}), 400 | |
| except Exception as e: | |
| os.remove(csv_path) | |
| return jsonify({"error": f"Invalid CSV: {e}"}), 400 | |
| # Reset state & start background thread | |
| _training_state.update({ | |
| "active": True, "logs": [], "progress": 0, "total": len(df), | |
| "phase": "starting", "metrics": None, "error": None, "done": False, | |
| }) | |
| thread = threading.Thread(target=_training_worker, args=(csv_path,), daemon=True) | |
| thread.start() | |
| return jsonify({"success": True, "total_rows": len(df)}) | |
| def admin_stream_train(): | |
| """SSE endpoint for training progress.""" | |
| def generate(): | |
| last_log_idx = 0 | |
| while True: | |
| logs = _training_state["logs"] | |
| new_logs = logs[last_log_idx:] | |
| last_log_idx = len(logs) | |
| data = { | |
| "active": _training_state["active"], | |
| "done": _training_state["done"], | |
| "phase": _training_state["phase"], | |
| "progress": _training_state["progress"], | |
| "total": _training_state["total"], | |
| "logs": new_logs, | |
| "metrics": _training_state["metrics"], | |
| "error": _training_state["error"], | |
| } | |
| yield f"data: {json.dumps(data)}\n\n" | |
| if _training_state["done"]: | |
| break | |
| time.sleep(1) | |
| return Response(stream_with_context(generate()), content_type="text/event-stream") | |
| def admin_download_model(): | |
| """Download trained model files as a zip (model + PCA + threshold).""" | |
| files_to_zip = [ | |
| ("nativity_model.pth", MODEL_PATH), | |
| ("pca_transform.pkl", PCA_PATH), | |
| ("best_threshold.txt", BEST_THRESHOLD_PATH), | |
| ] | |
| buf = io.BytesIO() | |
| with zipfile.ZipFile(buf, "w", zipfile.ZIP_DEFLATED) as zf: | |
| for name, path in files_to_zip: | |
| if os.path.isfile(path): | |
| zf.write(path, name) | |
| buf.seek(0) | |
| return send_file(buf, mimetype="application/zip", | |
| as_attachment=True, download_name="trained_model.zip") | |
| if __name__ == "__main__": | |
| print(f"Device: {device} | Dim: {embedding_dim} | Threshold: {EVAL_THRESHOLD}") | |
| app.run(host="0.0.0.0", port=5000, debug=False, threaded=True) | |