File size: 5,456 Bytes
637a5d2 25f659b 637a5d2 25f659b 637a5d2 25f659b 637a5d2 25f659b 637a5d2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 | """
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}."
)
|