Spaces:

neev2211
/

arabic-native-classifier

Sleeping

arabic-native-classifier / scripts /make_predictions.py

Fix model architecture: add PCA transform (2051->256d) to match training pipeline

6a10fdd 8 months ago

7.97 kB

"""
Run inference on test CSV: download audio, chunk, embed with Wav2Vec (middle layers), classify.
Arabic Native vs Non-Native. Uses language-agnostic middle-layer embeddings.
"""
import os
import sys
import pandas as pd
import numpy as np
import torch
import librosa
import requests
import io
from transformers import Wav2Vec2Model, AutoFeatureExtractor
from tqdm import tqdm
from sklearn.metrics import (
    accuracy_score,
    f1_score,
    precision_score,
    recall_score,
    confusion_matrix,
)

# Import shared classifier (same architecture as training)
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from model import NativityClassifier
from extract_embeddings import get_mfcc_from_audio, get_pitch_stats, PROSODY_EMBED_DIM, EMBED_DIM_WITH_MFCC

# -----------------------------------------------------------------------------
# Config
# -----------------------------------------------------------------------------
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model_path = "models/nativity_model.pth"
input_csv = "manifests/test_data.csv"
output_csv = "final_predictions.csv"
EMBED_DIM_FILE = "manifests/embedding_dim.txt"

# Model arch must match train_classifier_compare / model.py
HIDDEN_DIMS = (256, 128)  # Must match train_classifier_compare.py
DROPOUT = 0.4
PCA_COMPONENTS = 256
PCA_PATH = "models/pca_transform.pkl"
EVAL_THRESHOLD_DEFAULT = 0.35
BEST_THRESHOLD_PATH = "models/best_threshold.txt"
PROSODY_LAYERS = (4, 6, 8, 10)  # must match extract_embeddings.py (prosody-focused)

# Embedding dim: 2051 (prosody) or 2090 (prosody+MFCC)
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())

# Use validation-picked threshold if available (from train_classifier_compare.py)
if os.path.isfile(BEST_THRESHOLD_PATH):
    with open(BEST_THRESHOLD_PATH) as f:
        EVAL_THRESHOLD = float(f.read().strip())
else:
    EVAL_THRESHOLD = EVAL_THRESHOLD_DEFAULT

# -----------------------------------------------------------------------------
# Load models
# -----------------------------------------------------------------------------
import joblib

model_name = "facebook/wav2vec2-large-xlsr-53"
feature_extractor = AutoFeatureExtractor.from_pretrained(model_name)
wav2vec = Wav2Vec2Model.from_pretrained(model_name).to(device)
wav2vec.eval()

pca = joblib.load(PCA_PATH)

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()


def get_embedding(audio_np, embed_dim=None):
    """Prosody embedding: mean+std+pitch (2051-d) or +MFCC (2090-d); L2 norm to match training."""
    if embed_dim is None:
        embed_dim = embedding_dim
    with torch.no_grad():
        input_values = feature_extractor(
            audio_np, return_tensors="pt", sampling_rate=16000
        ).input_values.to(device)
        outputs = wav2vec(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)  # (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)
    pitch = get_pitch_stats(audio_np, sr=16000)
    out = np.concatenate([emb_mean, emb_std, pitch]).astype(np.float32)
    if embed_dim == EMBED_DIM_WITH_MFCC:
        mfcc_39 = get_mfcc_from_audio(audio_np, sr=16000)
        out = np.concatenate([out, mfcc_39]).astype(np.float32)
    n = np.linalg.norm(out)
    if n > 1e-8:
        out = out / n
    out = out.reshape(1, -1)
    # Apply PCA to match training pipeline (2051-d -> 256-d)
    out = pca.transform(out).astype(np.float32)
    return torch.tensor(out, dtype=torch.float32, device=device)


def process_file_and_predict(url):
    """
    Download, chunk (10s), embed each chunk, then aggregate with confidence-weighted voting.
    Returns (label_str, confidence).
    """
    try:
        response = requests.get(url, timeout=15)
        response.raise_for_status()
        raw = response.content
        if not raw or len(raw) < 100:
            return "Error", 0.0
        ct = (response.headers.get("Content-Type") or "").lower()
        if "text/html" in ct or raw.lstrip()[:50].startswith(b"<!") or raw.lstrip()[:20].startswith(b"<?"):
            return "Error", 0.0
        audio_bytes = io.BytesIO(raw)
        y, sr = librosa.load(audio_bytes, sr=16000, mono=True)
    except Exception:
        return "Error", 0.0

    chunk_size = 160000  # 10 sec at 16 kHz
    chunks = [
        y[i : i + chunk_size]
        for i in range(0, len(y), chunk_size)
        if len(y[i : i + chunk_size]) == chunk_size
    ]
    if not chunks:
        return "Error", 0.0

    chunk_preds = []   # 0 = Native, 1 = Non-Native
    chunk_confs = []   # probability of predicted class

    with torch.no_grad():
        for chunk in chunks:
            emb = get_embedding(chunk, embed_dim=embedding_dim)
            logits = classifier(emb)
            probs = torch.softmax(logits, dim=1)
            p1 = probs[0, 1].item()
            pred = 1 if p1 >= EVAL_THRESHOLD else 0
            conf = p1 if pred == 1 else (1.0 - p1)
            chunk_preds.append(pred)
            chunk_confs.append(conf)

    # Confidence-weighted voting: weight each chunk by its confidence
    weights = np.array(chunk_confs, dtype=np.float64)
    weights = weights / (weights.sum() + 1e-10)
    weighted_score = np.sum(np.array(chunk_preds) * weights)
    final_label = "Non-Native" if weighted_score >= EVAL_THRESHOLD else "Native"
    # Final confidence = weighted average of confidences (or max)
    final_conf = float(np.average(chunk_confs, weights=weights))

    return final_label, final_conf


def main():
    df = pd.read_csv(input_csv)
    results = []
    print(f"Predicting for {len(df)} files (device={device})...")
    for _, row in tqdm(df.iterrows(), total=len(df)):
        label, conf = process_file_and_predict(row["audio_url"])
        results.append({
            "dp_id": row["dp_id"],
            "prediction": label,
            "confidence_score": round(conf, 4),
        })
    out_df = pd.DataFrame(results)
    out_df.to_csv(output_csv, index=False)
    n_ok = (out_df["prediction"] != "Error").sum()
    n_fail = len(out_df) - n_ok
    print(f"Done. Saved to {output_csv} ({n_ok} predicted, {n_fail} failed to load)")

    # If ground-truth labels exist (nativity_status), print test metrics like validation
    if "nativity_status" in df.columns:
        mask = out_df["prediction"] != "Error"
        y_true_str = df.loc[mask, "nativity_status"].values
        y_pred_str = out_df.loc[mask, "prediction"].values
        map_fn = lambda s: 1 if s == "Non-Native" else 0
        y_true = np.array([map_fn(s) for s in y_true_str])
        y_pred = np.array([map_fn(s) for s in y_pred_str])

        acc = accuracy_score(y_true, y_pred)
        f1 = f1_score(y_true, y_pred, average="macro", zero_division=0)
        prec = precision_score(y_true, y_pred, average="macro", zero_division=0)
        rec = recall_score(y_true, y_pred, average="macro", zero_division=0)
        cm = confusion_matrix(y_true, y_pred)

        print("\nTest metrics (from test_data.csv, using nativity_status):")
        print(f"  Accuracy:  {acc:.4f}")
        print(f"  F1 macro:  {f1:.4f}")
        print(f"  Precision: {prec:.4f}")
        print(f"  Recall:    {rec:.4f}")
        print("  Confusion matrix (rows=true, cols=pred):")
        print(f"    {cm}")
        print("  (Native=0, Non-Native=1)")


if __name__ == "__main__":
    main()

	"""
	Run inference on test CSV: download audio, chunk, embed with Wav2Vec (middle layers), classify.
	Arabic Native vs Non-Native. Uses language-agnostic middle-layer embeddings.
	"""
	import os
	import sys
	import pandas as pd
	import numpy as np
	import torch
	import librosa
	import requests
	import io
	from transformers import Wav2Vec2Model, AutoFeatureExtractor
	from tqdm import tqdm
	from sklearn.metrics import (
	accuracy_score,
	f1_score,
	precision_score,
	recall_score,
	confusion_matrix,
	)

	# Import shared classifier (same architecture as training)
	sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
	from model import NativityClassifier
	from extract_embeddings import get_mfcc_from_audio, get_pitch_stats, PROSODY_EMBED_DIM, EMBED_DIM_WITH_MFCC

	# -----------------------------------------------------------------------------
	# Config
	# -----------------------------------------------------------------------------
	device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
	model_path = "models/nativity_model.pth"
	input_csv = "manifests/test_data.csv"
	output_csv = "final_predictions.csv"
	EMBED_DIM_FILE = "manifests/embedding_dim.txt"

	# Model arch must match train_classifier_compare / model.py
	HIDDEN_DIMS = (256, 128) # Must match train_classifier_compare.py
	DROPOUT = 0.4
	PCA_COMPONENTS = 256
	PCA_PATH = "models/pca_transform.pkl"
	EVAL_THRESHOLD_DEFAULT = 0.35
	BEST_THRESHOLD_PATH = "models/best_threshold.txt"
	PROSODY_LAYERS = (4, 6, 8, 10) # must match extract_embeddings.py (prosody-focused)

	# Embedding dim: 2051 (prosody) or 2090 (prosody+MFCC)
	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())

	# Use validation-picked threshold if available (from train_classifier_compare.py)
	if os.path.isfile(BEST_THRESHOLD_PATH):
	with open(BEST_THRESHOLD_PATH) as f:
	EVAL_THRESHOLD = float(f.read().strip())
	else:
	EVAL_THRESHOLD = EVAL_THRESHOLD_DEFAULT

	# -----------------------------------------------------------------------------
	# Load models
	# -----------------------------------------------------------------------------
	import joblib

	model_name = "facebook/wav2vec2-large-xlsr-53"
	feature_extractor = AutoFeatureExtractor.from_pretrained(model_name)
	wav2vec = Wav2Vec2Model.from_pretrained(model_name).to(device)
	wav2vec.eval()

	pca = joblib.load(PCA_PATH)

	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()


	def get_embedding(audio_np, embed_dim=None):
	"""Prosody embedding: mean+std+pitch (2051-d) or +MFCC (2090-d); L2 norm to match training."""
	if embed_dim is None:
	embed_dim = embedding_dim
	with torch.no_grad():
	input_values = feature_extractor(
	audio_np, return_tensors="pt", sampling_rate=16000
	).input_values.to(device)
	outputs = wav2vec(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) # (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)
	pitch = get_pitch_stats(audio_np, sr=16000)
	out = np.concatenate([emb_mean, emb_std, pitch]).astype(np.float32)
	if embed_dim == EMBED_DIM_WITH_MFCC:
	mfcc_39 = get_mfcc_from_audio(audio_np, sr=16000)
	out = np.concatenate([out, mfcc_39]).astype(np.float32)
	n = np.linalg.norm(out)
	if n > 1e-8:
	out = out / n
	out = out.reshape(1, -1)
	# Apply PCA to match training pipeline (2051-d -> 256-d)
	out = pca.transform(out).astype(np.float32)
	return torch.tensor(out, dtype=torch.float32, device=device)


	def process_file_and_predict(url):
	"""
	Download, chunk (10s), embed each chunk, then aggregate with confidence-weighted voting.
	Returns (label_str, confidence).
	"""
	try:
	response = requests.get(url, timeout=15)
	response.raise_for_status()
	raw = response.content
	if not raw or len(raw) < 100:
	return "Error", 0.0
	ct = (response.headers.get("Content-Type") or "").lower()
	if "text/html" in ct or raw.lstrip()[:50].startswith(b"<!") or raw.lstrip()[:20].startswith(b"<?"):
	return "Error", 0.0
	audio_bytes = io.BytesIO(raw)
	y, sr = librosa.load(audio_bytes, sr=16000, mono=True)
	except Exception:
	return "Error", 0.0

	chunk_size = 160000 # 10 sec at 16 kHz
	chunks = [
	y[i : i + chunk_size]
	for i in range(0, len(y), chunk_size)
	if len(y[i : i + chunk_size]) == chunk_size
	]
	if not chunks:
	return "Error", 0.0

	chunk_preds = [] # 0 = Native, 1 = Non-Native
	chunk_confs = [] # probability of predicted class

	with torch.no_grad():
	for chunk in chunks:
	emb = get_embedding(chunk, embed_dim=embedding_dim)
	logits = classifier(emb)
	probs = torch.softmax(logits, dim=1)
	p1 = probs[0, 1].item()
	pred = 1 if p1 >= EVAL_THRESHOLD else 0
	conf = p1 if pred == 1 else (1.0 - p1)
	chunk_preds.append(pred)
	chunk_confs.append(conf)

	# Confidence-weighted voting: weight each chunk by its confidence
	weights = np.array(chunk_confs, dtype=np.float64)
	weights = weights / (weights.sum() + 1e-10)
	weighted_score = np.sum(np.array(chunk_preds) * weights)
	final_label = "Non-Native" if weighted_score >= EVAL_THRESHOLD else "Native"
	# Final confidence = weighted average of confidences (or max)
	final_conf = float(np.average(chunk_confs, weights=weights))

	return final_label, final_conf


	def main():
	df = pd.read_csv(input_csv)
	results = []
	print(f"Predicting for {len(df)} files (device={device})...")
	for _, row in tqdm(df.iterrows(), total=len(df)):
	label, conf = process_file_and_predict(row["audio_url"])
	results.append({
	"dp_id": row["dp_id"],
	"prediction": label,
	"confidence_score": round(conf, 4),
	})
	out_df = pd.DataFrame(results)
	out_df.to_csv(output_csv, index=False)
	n_ok = (out_df["prediction"] != "Error").sum()
	n_fail = len(out_df) - n_ok
	print(f"Done. Saved to {output_csv} ({n_ok} predicted, {n_fail} failed to load)")

	# If ground-truth labels exist (nativity_status), print test metrics like validation
	if "nativity_status" in df.columns:
	mask = out_df["prediction"] != "Error"
	y_true_str = df.loc[mask, "nativity_status"].values
	y_pred_str = out_df.loc[mask, "prediction"].values
	map_fn = lambda s: 1 if s == "Non-Native" else 0
	y_true = np.array([map_fn(s) for s in y_true_str])
	y_pred = np.array([map_fn(s) for s in y_pred_str])

	acc = accuracy_score(y_true, y_pred)
	f1 = f1_score(y_true, y_pred, average="macro", zero_division=0)
	prec = precision_score(y_true, y_pred, average="macro", zero_division=0)
	rec = recall_score(y_true, y_pred, average="macro", zero_division=0)
	cm = confusion_matrix(y_true, y_pred)

	print("\nTest metrics (from test_data.csv, using nativity_status):")
	print(f" Accuracy: {acc:.4f}")
	print(f" F1 macro: {f1:.4f}")
	print(f" Precision: {prec:.4f}")
	print(f" Recall: {rec:.4f}")
	print(" Confusion matrix (rows=true, cols=pred):")
	print(f" {cm}")
	print(" (Native=0, Non-Native=1)")


	if __name__ == "__main__":
	main()