Spaces:
Sleeping
Sleeping
| """ | |
| Extract SEQUENTIAL embeddings for Arabic Native vs Non-Native 1D-CNN classifier. | |
| Outputs a 2D array (Time_Steps, 1024) instead of a flattened 1D array. | |
| Uses XLS-R 53 EARLY-TO-MID layers (4,6,8,10). | |
| """ | |
| 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 | |
| # Prosody-focused layers | |
| # model_name = "facebook/wav2vec2-large-xlsr-53" | |
| model_name = "facebook/wav2vec2-lv-60-espeak-cv-ft" | |
| PROSODY_LAYERS = (4, 6, 8, 10) | |
| _feature_extractor = None | |
| _model = None | |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| 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 | |
| # The exact feature size of the Wav2Vec2 layers | |
| WAV2VEC_DIM = 1024 | |
| def get_embeddings(audio_path): | |
| """ | |
| Extracts the chronological 2D sequence: (Time, 1024) | |
| No averaging. No squashing. | |
| """ | |
| feature_extractor, model = _get_wav2vec() | |
| try: | |
| speech, sr = librosa.load(audio_path, sr=16000) | |
| input_values = feature_extractor( | |
| 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] | |
| # Stack the target layers and average them together (Layer Average, NOT Time Average) | |
| # Shape becomes (Time, 1024) | |
| time_seq = torch.stack(embs, dim=0).mean(dim=0).squeeze(0) | |
| # Convert directly to numpy and return the 2D matrix | |
| emb = time_seq.cpu().numpy().astype(np.float32) | |
| return emb | |
| except Exception as e: | |
| print(f"Error on {audio_path}: {e}") | |
| return None | |
| def get_embeddings_from_audio(speech, sr=16000): | |
| """Extract 2D sequence (Time, 1024) from raw audio array. For test-time use.""" | |
| feature_extractor, model = _get_wav2vec() | |
| try: | |
| input_values = feature_extractor( | |
| speech, return_tensors="pt", sampling_rate=sr | |
| ).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] | |
| time_seq = torch.stack(embs, dim=0).mean(dim=0).squeeze(0) | |
| return time_seq.cpu().numpy().astype(np.float32) | |
| except Exception: | |
| return None | |
| if __name__ == "__main__": | |
| import argparse | |
| ap = argparse.ArgumentParser( | |
| description="Extract 2D SEQUENTIAL embeddings for 1D-CNN." | |
| ) | |
| 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", | |
| ) | |
| args = ap.parse_args() | |
| manifest_path = args.manifest | |
| output_folder = args.output_dir | |
| if args.output_manifest: | |
| output_manifest_path = args.output_manifest | |
| else: | |
| 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 = [] | |
| print( | |
| f"Extracting 2D Sequential Embeddings for {len(df)} chunks (Shape: Time x {WAV2VEC_DIM})..." | |
| ) | |
| for idx, row in tqdm(df.iterrows(), total=len(df)): | |
| chunk_path = str(row["chunk_path"]).replace("\\", "/") | |
| emb = get_embeddings(chunk_path) | |
| if emb is not None: | |
| emb_filename = os.path.basename(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) | |
| # Save the dimension size so the training script knows what to expect (use separate file if custom manifest to avoid overwriting main embedding_dim.txt) | |
| dim_name = "embedding_dim_cnn.txt" if args.output_manifest else "embedding_dim.txt" | |
| dim_file = os.path.join(os.path.dirname(output_manifest_path), dim_name) | |
| with open(dim_file, "w") as f: | |
| f.write(str(WAV2VEC_DIM)) | |
| print( | |
| f"Done! 2D Embeddings in {output_folder}, manifest at {output_manifest_path}." | |
| ) | |