| """ |
| 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 |
|
|
| |
| 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 |
|
|
| |
| |
| |
| 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) |
| 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_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 = { |
| "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()) |
|
|
| |
| |
| |
| _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 |
| |
| 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 |
|
|
|
|
| |
| |
| |
| 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 |
| |
| 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()), |
| } |
|
|
|
|
| |
| |
| |
| 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 = {} |
|
|
|
|
| @app.route("/") |
| def index(): |
| return render_template("index.html") |
|
|
|
|
| @app.route("/api/health", methods=["GET"]) |
| def health(): |
| return jsonify({"status": "ok", "device": str(device), "embedding_dim": embedding_dim, "threshold": EVAL_THRESHOLD, "languages": SUPPORTED_LANGUAGES}) |
|
|
|
|
| @app.route("/api/upload_csv", methods=["POST"]) |
| 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 |
|
|
|
|
| @app.route("/api/stream/<job_id>") |
| 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") |
|
|
|
|
| @app.route("/api/predict/audio", methods=["POST"]) |
| 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 |
|
|
|
|
| @app.route("/api/predict/url", methods=["POST"]) |
| 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 |
|
|
|
|
| |
| |
| |
| def admin_required(f): |
| @functools.wraps(f) |
| def decorated(*args, **kwargs): |
| if not session.get("admin_authenticated"): |
| return redirect(url_for("admin_login")) |
| return f(*args, **kwargs) |
| return decorated |
|
|
|
|
| @app.route("/admin/login", methods=["GET", "POST"]) |
| 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) |
|
|
|
|
| @app.route("/admin/logout") |
| def admin_logout(): |
| session.pop("admin_authenticated", None) |
| return redirect(url_for("index")) |
|
|
|
|
| @app.route("/admin") |
| @admin_required |
| def admin_page(): |
| return render_template("admin.html") |
|
|
|
|
| @app.route("/admin/api/status") |
| @admin_required |
| 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}") |
|
|
| |
| _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 |
|
|
| |
| 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 |
|
|
| |
| 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)}") |
|
|
| |
| _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") |
|
|
| |
| with open(EMBED_DIM_FILE, "w") as f: |
| f.write(str(EMB_DIM)) |
|
|
| |
| _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)") |
|
|
| |
| 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]}") |
|
|
| |
| _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}") |
|
|
| |
| 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) |
|
|
| |
| _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}") |
|
|
| |
| 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_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}") |
|
|
| |
| _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}") |
|
|
|
|
| @app.route("/admin/api/train", methods=["POST"]) |
| @admin_required |
| 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 |
|
|
| |
| csv_path = os.path.join(UPLOAD_DIR, f"train_{uuid.uuid4().hex[:8]}.csv") |
| f.save(csv_path) |
|
|
| |
| 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 |
|
|
| |
| _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)}) |
|
|
|
|
| @app.route("/admin/api/stream_train") |
| @admin_required |
| 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") |
|
|
|
|
| @app.route("/admin/api/download_model") |
| @admin_required |
| 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) |
|
|
