Spaces:
Sleeping
Sleeping
| """ | |
| Extract embeddings for Arabic Native vs Non-Native classifier. | |
| Emphasizes PROSODIC features: pitch, rhythm, timing, fluency — key for Native vs Non-Native. | |
| Uses XLS-R 53 EARLY-TO-MID layers (4,6,8,10): prosody/speaker-rich, less lexical than final layer. | |
| """ | |
| import pandas as pd | |
| import numpy as np | |
| import torch | |
| import librosa | |
| import os | |
| from transformers import Wav2Vec2Model, AutoFeatureExtractor | |
| from tqdm import tqdm | |
| from scipy.signal import butter, lfilter | |
| # XLS-R 53 (multilingual, includes Arabic) | |
| # model_name = "facebook/wav2vec2-large-xlsr-53" | |
| model_name = "facebook/wav2vec2-lv-60-espeak-cv-ft" | |
| # Prosody-focused layers: early-mid (4,6,8,10) encode pitch, rhythm, timing, speaker | |
| # Final layers encode lexicon; we want prosodic cues for nativeness | |
| PROSODY_LAYERS = (4, 6, 8, 10) | |
| # Lazy-loaded so "from extract_embeddings import get_mfcc_from_audio, ..." doesn't load the big model (avoids 2x GPU OOM in compare_models_on_test) | |
| _feature_extractor = None | |
| _model = None | |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| def apply_whisper_filter(audio_array, sr=16000, cutoff_hz=700): | |
| """ | |
| Applies a High-Pass Butterworth filter to strip fundamental pitch | |
| and isolate high-frequency phonetic articulation. | |
| """ | |
| nyq = 0.5 * sr | |
| normal_cutoff = cutoff_hz / nyq | |
| b, a = butter(5, normal_cutoff, btype="high", analog=False) | |
| whispered_audio = lfilter(b, a, audio_array) | |
| return whispered_audio.astype(np.float32) | |
| def _get_wav2vec(): | |
| global _feature_extractor, _model | |
| if _model is None: | |
| _feature_extractor = AutoFeatureExtractor.from_pretrained(model_name) | |
| _model = Wav2Vec2Model.from_pretrained(model_name) | |
| _model.eval() | |
| _model.to(device) | |
| return _feature_extractor, _model | |
| # Embedding dims: paper uses mean + std + F0 → 1024+1024+3 = 2051; optional MFCC → 2090 | |
| WAV2VEC_DIM = 1024 | |
| PITCH_DIM = 3 # mean, std, max of F0 (voiced) | |
| PROSODY_EMBED_DIM = WAV2VEC_DIM * 2 + PITCH_DIM # 2051 (mean + std + pitch) | |
| MFCC_DIM = 39 # 13 + 13 + 13 | |
| EMBED_DIM_WITH_MFCC = PROSODY_EMBED_DIM + MFCC_DIM # 2090 | |
| def get_pitch_stats(speech, sr=16000, fmin=50, fmax=500): | |
| """F0 (pitch) stats for prosody: mean, std, max over voiced frames (paper: explicit pitch helps dialect/nativity).""" | |
| try: | |
| f0 = librosa.yin( | |
| speech.astype(np.float64), fmin=float(fmin), fmax=float(fmax), sr=sr | |
| ) | |
| voiced = (f0 > 0) & np.isfinite(f0) | |
| if np.any(voiced): | |
| f0_v = f0[voiced] | |
| return np.array( | |
| [float(np.mean(f0_v)), float(np.std(f0_v)), float(np.max(f0_v))], | |
| dtype=np.float32, | |
| ) | |
| except Exception: | |
| pass | |
| return np.zeros(PITCH_DIM, dtype=np.float32) | |
| def get_mfcc_from_audio(speech, sr=16000, n_mfcc=13): | |
| """Classical 39-d feature: MFCC mean + delta + delta2 (paper: MFCC >> LPCC/PLP for accent).""" | |
| mfcc = librosa.feature.mfcc(y=speech, sr=sr, n_mfcc=n_mfcc) | |
| mfcc_mean = mfcc.mean(axis=1) | |
| mfcc_delta = librosa.feature.delta(mfcc).mean(axis=1) | |
| mfcc_delta2 = librosa.feature.delta(mfcc, order=2).mean(axis=1) | |
| return np.concatenate([mfcc_mean, mfcc_delta, mfcc_delta2]).astype(np.float32) | |
| def get_embeddings(audio_path, l2_norm=True, return_audio=False): | |
| """Prosody embedding: wav2vec mean + std (temporal variance) + F0 stats (pitch) → 2051-d. | |
| Paper: explicit pitch and time-derivatives (std) help dialect/nativity. Optional return_audio for MFCC.""" | |
| feature_extractor, model = _get_wav2vec() | |
| try: | |
| speech, sr = librosa.load(audio_path, sr=16000) | |
| # 1. Calculate REAL pitch from the unfiltered audio | |
| pitch = get_pitch_stats(speech, sr=sr) | |
| # 2. Apply WHISPER FILTER to blind the model to speaker identity | |
| # whispered_speech = apply_whisper_filter(speech, sr=sr) | |
| whispered_speech = speech | |
| # 3. Pass the WHISPERED speech to Wav2Vec2 | |
| input_values = feature_extractor( | |
| whispered_speech, return_tensors="pt", sampling_rate=16000 | |
| ).input_values.to(device) | |
| with torch.no_grad(): | |
| outputs = model(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] | |
| # (1, T, 1024) after stack+mean over layers; keep time for mean+std | |
| 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) | |
| emb = np.concatenate([emb_mean, emb_std, pitch]) | |
| if l2_norm: | |
| n = np.linalg.norm(emb) | |
| if n > 1e-8: | |
| emb = emb / n | |
| if return_audio: | |
| # Return the WHISPERED speech so the MFCC extractor also uses the filtered version | |
| return emb, whispered_speech, sr | |
| return emb | |
| except Exception as e: | |
| print(f"Error on {audio_path}: {e}") | |
| return None | |
| # CLI: optional (manifest_in, embedding_dir, manifest_out) for e.g. Iqra data | |
| if __name__ == "__main__": | |
| import argparse | |
| ap = argparse.ArgumentParser( | |
| description="Extract XLS-R prosody embeddings from chunk manifest." | |
| ) | |
| ap.add_argument( | |
| "--manifest", | |
| default="manifests/processed_manifest.csv", | |
| help="Input manifest CSV (chunk_path, original_id, label, language)", | |
| ) | |
| ap.add_argument( | |
| "--output-dir", default="embeddings", help="Folder for .npy embeddings" | |
| ) | |
| ap.add_argument( | |
| "--output-manifest", | |
| default=None, | |
| help="Output manifest CSV (default: final_training_manifest.csv or derived from --manifest)", | |
| ) | |
| ap.add_argument( | |
| "--mfcc", | |
| action="store_true", | |
| help="Append 39-d MFCC → 2090-d (default prosody is 2051-d: mean+std+pitch)", | |
| ) | |
| args = ap.parse_args() | |
| manifest_path = args.manifest | |
| output_folder = args.output_dir | |
| use_mfcc = args.mfcc | |
| if args.output_manifest: | |
| output_manifest_path = args.output_manifest | |
| else: | |
| # Main pipeline: processed_manifest -> final_training_manifest; else e.g. iqra_processed_manifest -> iqra_embeddings_manifest | |
| if manifest_path.strip() == "manifests/processed_manifest.csv": | |
| output_manifest_path = "manifests/final_training_manifest.csv" | |
| else: | |
| d = os.path.dirname(manifest_path) | |
| base = ( | |
| os.path.basename(manifest_path) | |
| .replace(".csv", "") | |
| .replace("_processed", "_embeddings") | |
| ) | |
| output_manifest_path = os.path.join(d, base + ".csv") | |
| os.makedirs(output_folder, exist_ok=True) | |
| df = pd.read_csv(manifest_path) | |
| embedding_paths = [] | |
| feat_desc = f"prosody mean+std+pitch ({PROSODY_EMBED_DIM}d)" + ( | |
| f" + MFCC({MFCC_DIM}d)" if use_mfcc else "" | |
| ) | |
| print(f"Extracting embeddings for {len(df)} chunks ({feat_desc}, L2-normalized)...") | |
| for idx, row in tqdm(df.iterrows(), total=len(df)): | |
| if use_mfcc: | |
| out = get_embeddings(row["chunk_path"], l2_norm=False, return_audio=True) | |
| if out is None: | |
| embedding_paths.append(None) | |
| continue | |
| emb, speech, sr = out | |
| mfcc_39 = get_mfcc_from_audio(speech, sr=sr) | |
| emb = np.concatenate([emb, mfcc_39]).astype(np.float32) | |
| n = np.linalg.norm(emb) | |
| if n > 1e-8: | |
| emb = emb / n | |
| else: | |
| emb = get_embeddings(row["chunk_path"], l2_norm=True) | |
| if emb is not None: | |
| emb_filename = os.path.basename(row["chunk_path"]).replace(".wav", ".npy") | |
| emb_path = os.path.join(output_folder, emb_filename) | |
| np.save(emb_path, emb) | |
| embedding_paths.append(emb_path) | |
| else: | |
| embedding_paths.append(None) | |
| df["embedding_path"] = embedding_paths | |
| df.dropna(subset=["embedding_path"], inplace=True) | |
| df.to_csv(output_manifest_path, index=False) | |
| embed_dim = EMBED_DIM_WITH_MFCC if use_mfcc else PROSODY_EMBED_DIM | |
| dim_file = os.path.join(os.path.dirname(output_manifest_path), "embedding_dim.txt") | |
| with open(dim_file, "w") as f: | |
| f.write(str(embed_dim)) | |
| print( | |
| f"Done! Embeddings in {output_folder}, manifest at {output_manifest_path}, dim={embed_dim} (saved to {dim_file})." | |
| ) | |