""" Wav2Vec2 emotion classifier for the capstone audio sentiment module. This script fine-tunes Wav2Vec2 for 6-class speech emotion classification. Default training uses CREMA-D metadata. For Model V2, pass the combined CREMA-D + RAVDESS metadata file. Run quick smoke test from ml-services: python -m src.models.test_wav2vec2_setup Run CREMA-D training: python -m src.models.wav2vec2_emotion_model \ --model-checkpoint Dpngtm/wav2vec2-emotion-recognition \ --run-name model_v1_cremad \ --num-epochs 3 \ --batch-size 2 \ --learning-rate 1e-5 Run combined CREMA-D + RAVDESS training: python -m src.models.wav2vec2_emotion_model \ --metadata-path data/processed/combined_emotion_metadata.csv \ --model-checkpoint Dpngtm/wav2vec2-emotion-recognition \ --run-name model_v2_cremad_ravdess \ --num-epochs 3 \ --batch-size 2 \ --learning-rate 1e-5 Outputs: outputs/wav2vec2//best_model outputs/reports/_report.json outputs/reports/_confusion_matrix.csv """ import argparse import json import mlflow from dataclasses import dataclass from pathlib import Path from typing import Dict, List, Literal, Optional import numpy as np import pandas as pd import torch from sklearn.metrics import accuracy_score, classification_report, confusion_matrix, f1_score from torch.utils.data import Dataset from transformers import ( AutoConfig, Trainer, TrainingArguments, Wav2Vec2ForSequenceClassification, Wav2Vec2Processor, set_seed, ) from src.data.audio_dataset import ( DEFAULT_SAMPLE_RATE, EMOTION_LABELS, build_label_encoding, get_split_dataframe, load_audio_file, load_metadata, resolve_audio_path, ) PROJECT_ROOT = Path(__file__).resolve().parents[3] ML_SERVICES_ROOT = PROJECT_ROOT / "ml-services" DEFAULT_METADATA_PATH = ML_SERVICES_ROOT / "data" / "processed" / "cremad_metadata.csv" DEFAULT_OUTPUT_ROOT = ML_SERVICES_ROOT / "outputs" / "wav2vec2" DEFAULT_REPORTS_DIR = ML_SERVICES_ROOT / "outputs" / "reports" DEFAULT_MLFLOW_EXPERIMENT_NAME = "audio_sentiment_emotion_classification" DEFAULT_MLFLOW_TRACKING_URI = f"sqlite:///{ML_SERVICES_ROOT / 'mlflow.db'}" DEFAULT_MODEL_CHECKPOINT = "facebook/wav2vec2-base" TASK_NAME = "6-class speech emotion classification" RANDOM_SEED = 42 def add_background_noise( waveform: np.ndarray, noise_factor: float = 0.003, ) -> np.ndarray: """ Add light Gaussian noise to audio. """ noise = np.random.normal(0, noise_factor, waveform.shape) augmented = waveform + noise return augmented.astype(np.float32) def apply_volume_gain( waveform: np.ndarray, min_gain: float = 0.85, max_gain: float = 1.15, ) -> np.ndarray: """ Randomly increase or decrease audio volume slightly. """ gain = np.random.uniform(min_gain, max_gain) augmented = waveform * gain return np.clip(augmented, -1.0, 1.0).astype(np.float32) def apply_speed_change( waveform: np.ndarray, min_rate: float = 0.97, max_rate: float = 1.03, ) -> np.ndarray: """ Apply a very small speed change using interpolation. This is intentionally light because strong speed changes can distort emotion labels. """ rate = np.random.uniform(min_rate, max_rate) original_indices = np.arange(len(waveform)) new_length = max(1, int(len(waveform) / rate)) new_indices = np.linspace(0, len(waveform) - 1, new_length) augmented = np.interp(new_indices, original_indices, waveform) return augmented.astype(np.float32) def apply_light_audio_augmentation( waveform: np.ndarray, probability: float = 0.5, ) -> np.ndarray: """ Apply light training-only audio augmentation. Each augmentation is optional and mild. """ augmented = waveform.astype(np.float32) if np.random.random() < probability: augmented = add_background_noise(augmented) if np.random.random() < probability: augmented = apply_volume_gain(augmented) if np.random.random() < probability: augmented = apply_speed_change(augmented) return augmented.astype(np.float32) class Wav2Vec2EmotionDataset(Dataset): """ PyTorch Dataset for Wav2Vec2 emotion classification. Each item returns raw audio values and a numeric label. Padding is handled later by Wav2Vec2DataCollator. """ def __init__( self, metadata: pd.DataFrame, label_to_id: Dict[str, int], sample_rate: int = DEFAULT_SAMPLE_RATE, max_duration_seconds: Optional[float] = 6.0, enable_augmentation: bool = False, augmentation_probability: float = 0.5, ) -> None: self.metadata = metadata.reset_index(drop=True) self.label_to_id = label_to_id self.sample_rate = sample_rate self.max_duration_seconds = max_duration_seconds self.enable_augmentation = enable_augmentation self.augmentation_probability = augmentation_probability required_columns = {"file_path", "emotion_label", "filename"} missing_columns = required_columns - set(self.metadata.columns) if missing_columns: raise ValueError(f"Metadata is missing columns: {sorted(missing_columns)}") def __len__(self) -> int: return len(self.metadata) def __getitem__(self, index: int) -> Dict: row = self.metadata.iloc[index] audio_path = resolve_audio_path(row["file_path"]) waveform, sample_rate = load_audio_file( audio_path=audio_path, target_sample_rate=self.sample_rate, max_duration_seconds=self.max_duration_seconds, ) if self.enable_augmentation: waveform = apply_light_audio_augmentation( waveform, probability=self.augmentation_probability, ) label_name = row["emotion_label"] label_id = self.label_to_id[label_name] return { "input_values": waveform, "labels": label_id, "filename": row["filename"], } @dataclass class Wav2Vec2DataCollator: """ Pads variable-length audio inputs for Wav2Vec2 training. Audio clips do not all have the same length, so this collator pads each batch dynamically using the Wav2Vec2 processor. """ processor: Wav2Vec2Processor sampling_rate: int = DEFAULT_SAMPLE_RATE def __call__(self, features: List[Dict]) -> Dict[str, torch.Tensor]: """ Convert a list of dataset samples into one padded training batch. The processor expects a list of raw waveform arrays, not a list of dictionaries. Labels are added separately after padding. """ input_values = [feature["input_values"] for feature in features] batch = self.processor( input_values, sampling_rate=self.sampling_rate, padding=True, return_attention_mask=True, return_tensors="pt", ) batch["labels"] = torch.tensor( [feature["labels"] for feature in features], dtype=torch.long, ) return batch def limit_metadata_per_split( metadata: pd.DataFrame, limit_per_split: Optional[int], ) -> pd.DataFrame: """ Limit records per split for quick testing. This is useful to verify training works before running full fine-tuning. """ if limit_per_split is None: return metadata limited_parts = [] for split_name in ["train", "validation", "test"]: split_df = metadata[metadata["split"] == split_name].head(limit_per_split) limited_parts.append(split_df) limited_metadata = pd.concat(limited_parts, ignore_index=True) print(f"Using limit_per_split={limit_per_split}") print(f"Limited split distribution: {limited_metadata['split'].value_counts().to_dict()}") return limited_metadata def infer_dataset_source(metadata: pd.DataFrame) -> str: """ Infer dataset source from metadata. If the metadata has a dataset column, return the joined dataset names. Otherwise, assume CREMA-D for backwards compatibility. """ if "dataset" not in metadata.columns: return "CREMA-D" datasets = sorted(metadata["dataset"].dropna().unique().tolist()) if not datasets: return "Unknown" return " + ".join(datasets) def build_wav2vec2_datasets( metadata_path: Path = DEFAULT_METADATA_PATH, limit_per_split: Optional[int] = None, max_duration_seconds: Optional[float] = 6.0, enable_augmentation: bool = False, augmentation_probability: float = 0.5, ) -> Dict[str, Wav2Vec2EmotionDataset]: """ Build train, validation, and test datasets for Wav2Vec2. """ metadata = load_metadata(metadata_path) metadata = limit_metadata_per_split(metadata, limit_per_split) label_encoding = build_label_encoding(task="emotion") train_df = get_split_dataframe(metadata, "train") validation_df = get_split_dataframe(metadata, "validation") test_df = get_split_dataframe(metadata, "test") return { "train": Wav2Vec2EmotionDataset( metadata=train_df, label_to_id=label_encoding.label_to_id, max_duration_seconds=max_duration_seconds, enable_augmentation=enable_augmentation, augmentation_probability=augmentation_probability, ), "validation": Wav2Vec2EmotionDataset( metadata=validation_df, label_to_id=label_encoding.label_to_id, max_duration_seconds=max_duration_seconds, enable_augmentation=False, ), "test": Wav2Vec2EmotionDataset( metadata=test_df, label_to_id=label_encoding.label_to_id, max_duration_seconds=max_duration_seconds, enable_augmentation=False, ), } def normalize_checkpoint_label(label: str) -> str: """ Normalize checkpoint labels to match the project emotion schema. Checkpoint labels: angry, calm, disgust, fearful, happy, sad, surprised Project labels: anger, neutral, disgust, fear, happy, sadness """ label = str(label).strip().lower() aliases = { "angry": "anger", "fearful": "fear", "sad": "sadness", "calm": "neutral", } return aliases.get(label, label) def transfer_matching_classifier_weights( model: Wav2Vec2ForSequenceClassification, model_checkpoint: str, target_label_to_id: Dict[str, int], ) -> bool: """ Transfer classifier weights from the pretrained 7-label checkpoint into our 6-label project classifier. The checkpoint has: angry, calm, disgust, fearful, happy, sad, surprised Our project has: anger, disgust, fear, happy, neutral, sadness We reuse: angry -> anger calm -> neutral disgust -> disgust fearful -> fear happy -> happy sad -> sadness We ignore: surprised """ source_config = AutoConfig.from_pretrained(model_checkpoint) source_id_to_project_label = { int(label_id): normalize_checkpoint_label(label_name) for label_id, label_name in source_config.id2label.items() } source_project_label_to_id = { label_name: label_id for label_id, label_name in source_id_to_project_label.items() } missing_labels = [ label for label in target_label_to_id.keys() if label.lower() not in source_project_label_to_id ] if missing_labels: print( "Classifier weight transfer skipped. " f"Missing labels in checkpoint after alias mapping: {missing_labels}" ) print(f"Checkpoint labels after normalization: {source_id_to_project_label}") print(f"Target labels: {target_label_to_id}") return False source_model = Wav2Vec2ForSequenceClassification.from_pretrained(model_checkpoint) with torch.no_grad(): for target_label, target_id in target_label_to_id.items(): normalized_target_label = target_label.lower() source_id = source_project_label_to_id[normalized_target_label] model.classifier.weight[target_id].copy_( source_model.classifier.weight[source_id] ) model.classifier.bias[target_id].copy_( source_model.classifier.bias[source_id] ) print("Transferred matching classifier weights from checkpoint.") print(f"Checkpoint labels after normalization: {source_id_to_project_label}") print(f"Target labels: {target_label_to_id}") return True def build_model_and_processor( model_checkpoint: str = DEFAULT_MODEL_CHECKPOINT, freeze_feature_encoder: bool = True, freeze_transformer_layers: int = 0, transfer_classifier_weights: bool = False, ) -> tuple[Wav2Vec2ForSequenceClassification, Wav2Vec2Processor, bool]: """ Load Wav2Vec2 model and processor for 6-class classification. """ label_encoding = build_label_encoding(task="emotion") processor = Wav2Vec2Processor.from_pretrained(model_checkpoint) model = Wav2Vec2ForSequenceClassification.from_pretrained( model_checkpoint, num_labels=len(EMOTION_LABELS), label2id=label_encoding.label_to_id, id2label=label_encoding.id_to_label, ignore_mismatched_sizes=True, ) classifier_weights_transferred = False if transfer_classifier_weights: classifier_weights_transferred = transfer_matching_classifier_weights( model=model, model_checkpoint=model_checkpoint, target_label_to_id=label_encoding.label_to_id, ) # This is safer for laptops and speeds up training. # Later we can unfreeze for stronger fine-tuning. if freeze_feature_encoder: model.freeze_feature_encoder() if freeze_transformer_layers > 0: encoder_layers = model.wav2vec2.encoder.layers total_layers = len(encoder_layers) if freeze_transformer_layers > total_layers: raise ValueError( f"freeze_transformer_layers={freeze_transformer_layers} is larger " f"than total transformer layers={total_layers}" ) for layer_index in range(freeze_transformer_layers): for parameter in encoder_layers[layer_index].parameters(): parameter.requires_grad = False return model, processor, classifier_weights_transferred def compute_metrics(eval_prediction) -> Dict[str, float]: """ Compute evaluation metrics during training. """ logits, labels = eval_prediction predictions = np.argmax(logits, axis=-1) return { "accuracy": float(accuracy_score(labels, predictions)), "macro_f1": float(f1_score(labels, predictions, average="macro")), "weighted_f1": float(f1_score(labels, predictions, average="weighted")), } def get_device_note() -> str: """ Return a readable note about available acceleration. """ if torch.cuda.is_available(): return "CUDA GPU available" if hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): return "Apple Silicon MPS available" return "CPU only" def count_trainable_parameters(model: torch.nn.Module) -> Dict[str, int]: """ Count trainable and total model parameters. """ total_parameters = sum(parameter.numel() for parameter in model.parameters()) trainable_parameters = sum( parameter.numel() for parameter in model.parameters() if parameter.requires_grad ) return { "total_parameters": int(total_parameters), "trainable_parameters": int(trainable_parameters), "frozen_parameters": int(total_parameters - trainable_parameters), } def create_training_arguments( output_dir: Path, num_epochs: int, batch_size: int, learning_rate: float, weight_decay: float, warmup_ratio: float, ) -> TrainingArguments: """ Create Hugging Face training arguments. Uses eval_strategy because newer Transformers versions replaced the older evaluation_strategy argument. """ return TrainingArguments( output_dir=str(output_dir), eval_strategy="epoch", save_strategy="epoch", logging_strategy="steps", logging_steps=25, learning_rate=learning_rate, per_device_train_batch_size=batch_size, per_device_eval_batch_size=batch_size, num_train_epochs=num_epochs, weight_decay=weight_decay, warmup_ratio=warmup_ratio, load_best_model_at_end=True, metric_for_best_model="macro_f1", greater_is_better=True, save_total_limit=2, report_to=[], seed=RANDOM_SEED, fp16=torch.cuda.is_available(), dataloader_num_workers=0, ) def evaluate_on_test_set( trainer: Trainer, test_dataset: Dataset, id_to_label: Dict[int, str], confusion_matrix_path: Path, ) -> Dict: """ Evaluate the final model on the held-out test set. """ predictions_output = trainer.predict(test_dataset) logits = predictions_output.predictions labels = predictions_output.label_ids predictions = np.argmax(logits, axis=-1) target_names = [id_to_label[index] for index in sorted(id_to_label.keys())] label_ids = sorted(id_to_label.keys()) report = { "accuracy": float(accuracy_score(labels, predictions)), "macro_f1": float( f1_score( labels, predictions, labels=label_ids, average="macro", zero_division=0, ) ), "weighted_f1": float( f1_score( labels, predictions, labels=label_ids, average="weighted", zero_division=0, ) ), "classification_report": classification_report( labels, predictions, labels=label_ids, target_names=target_names, output_dict=True, zero_division=0, ), } matrix = confusion_matrix(labels, predictions, labels=label_ids) matrix_df = pd.DataFrame( matrix, index=[f"actual_{label}" for label in target_names], columns=[f"predicted_{label}" for label in target_names], ) confusion_matrix_path.parent.mkdir(parents=True, exist_ok=True) matrix_df.to_csv(confusion_matrix_path) return report def log_training_run_to_mlflow( full_report: Dict, report_path: Path, confusion_matrix_path: Path, run_name: str, mlflow_experiment_name: str, ) -> str: """ Log Wav2Vec2 training results to MLflow. This tracks model configuration, dataset information, validation/test metrics, and report artifacts for experiment comparison. """ mlflow.set_tracking_uri(DEFAULT_MLFLOW_TRACKING_URI) mlflow.set_experiment(mlflow_experiment_name) with mlflow.start_run(run_name=run_name): # Tags mlflow.set_tag("module", "audio_sentiment_analysis") mlflow.set_tag("model_family", "Wav2Vec2") mlflow.set_tag("training_framework", "Hugging Face Transformers") mlflow.set_tag("dataset_source", full_report.get("dataset_source", "Unknown")) mlflow.set_tag("purpose", "model_training") # Parameters mlflow.log_param("run_name", full_report.get("run_name")) mlflow.log_param("model_name", full_report.get("model_name")) mlflow.log_param("base_checkpoint", full_report.get("base_checkpoint")) mlflow.log_param("task", full_report.get("task")) mlflow.log_param("device", full_report.get("device")) mlflow.log_param("dataset_source", full_report.get("dataset_source")) mlflow.log_param("metadata_path", full_report.get("metadata_path")) mlflow.log_param("output_dir", full_report.get("output_dir")) mlflow.log_param("best_model_dir", full_report.get("best_model_dir")) mlflow.log_param("train_samples", full_report.get("train_samples")) mlflow.log_param("validation_samples", full_report.get("validation_samples")) mlflow.log_param("test_samples", full_report.get("test_samples")) mlflow.log_param("num_epochs", full_report.get("num_epochs")) mlflow.log_param("batch_size", full_report.get("batch_size")) mlflow.log_param("learning_rate", full_report.get("learning_rate")) mlflow.log_param("enable_augmentation", full_report.get("enable_augmentation")) mlflow.log_param( "transfer_classifier_weights", full_report.get("transfer_classifier_weights"), ) mlflow.log_param( "classifier_weights_transferred", full_report.get("classifier_weights_transferred"), ) mlflow.log_param("augmentation_probability", full_report.get("augmentation_probability")) mlflow.log_param("weight_decay", full_report.get("weight_decay")) mlflow.log_param("warmup_ratio", full_report.get("warmup_ratio")) mlflow.log_param( "freeze_feature_encoder", full_report.get("freeze_feature_encoder"), ) mlflow.log_param( "freeze_transformer_layers", full_report.get("freeze_transformer_layers"), ) mlflow.log_param("total_parameters", full_report.get("total_parameters")) mlflow.log_param("trainable_parameters", full_report.get("trainable_parameters")) mlflow.log_param("frozen_parameters", full_report.get("frozen_parameters")) # Validation metrics validation_metrics = full_report.get("validation", {}) mlflow.log_metric( "validation_accuracy", float(validation_metrics.get("eval_accuracy", 0.0)), ) mlflow.log_metric( "validation_macro_f1", float(validation_metrics.get("eval_macro_f1", 0.0)), ) mlflow.log_metric( "validation_weighted_f1", float(validation_metrics.get("eval_weighted_f1", 0.0)), ) mlflow.log_metric( "validation_loss", float(validation_metrics.get("eval_loss", 0.0)), ) # Test metrics test_metrics = full_report.get("test", {}) mlflow.log_metric("test_accuracy", float(test_metrics.get("accuracy", 0.0))) mlflow.log_metric("test_macro_f1", float(test_metrics.get("macro_f1", 0.0))) mlflow.log_metric( "test_weighted_f1", float(test_metrics.get("weighted_f1", 0.0)), ) # Artifacts if report_path.exists(): mlflow.log_artifact(str(report_path), artifact_path="reports") if confusion_matrix_path.exists(): mlflow.log_artifact(str(confusion_matrix_path), artifact_path="reports") run_id = mlflow.active_run().info.run_id return run_id def train_wav2vec2_emotion_model( model_checkpoint: str = DEFAULT_MODEL_CHECKPOINT, metadata_path: Path = DEFAULT_METADATA_PATH, run_name: str = "wav2vec2_emotion", output_root: Path = DEFAULT_OUTPUT_ROOT, limit_per_split: Optional[int] = None, num_epochs: int = 5, batch_size: int = 4, learning_rate: float = 3e-5, warmup_ratio: float = 0.1, weight_decay: float = 0.01, freeze_feature_encoder: bool = True, freeze_transformer_layers: int = 0, transfer_classifier_weights: bool = False, max_duration_seconds: Optional[float] = 6.0, enable_augmentation: bool = False, augmentation_probability: float = 0.5, enable_mlflow: bool = False, mlflow_experiment_name: str = DEFAULT_MLFLOW_EXPERIMENT_NAME, ) -> Dict: """ Fine-tune Wav2Vec2 for emotion classification. """ set_seed(RANDOM_SEED) metadata_path = Path(metadata_path) output_root = Path(output_root) output_dir = output_root / run_name best_model_dir = output_dir / "best_model" report_path = DEFAULT_REPORTS_DIR / f"{run_name}_report.json" confusion_matrix_path = DEFAULT_REPORTS_DIR / f"{run_name}_confusion_matrix.csv" metadata = load_metadata(metadata_path) dataset_source = infer_dataset_source(metadata) print("\nStarting Wav2Vec2 emotion fine-tuning") print("-" * 70) print(f"Run name: {run_name}") print(f"Device: {get_device_note()}") print(f"Dataset source: {dataset_source}") print(f"Metadata path: {metadata_path}") print(f"Model checkpoint: {model_checkpoint}") print(f"Epochs: {num_epochs}") print(f"Batch size: {batch_size}") print(f"Learning rate: {learning_rate}") print(f"Warmup ratio: {warmup_ratio}") print(f"Augmentation enabled: {enable_augmentation}") print(f"Augmentation probability: {augmentation_probability}") print(f"Transfer classifier weights: {transfer_classifier_weights}") print(f"Output directory: {output_dir}") print(f"MLflow enabled: {enable_mlflow}") if enable_mlflow: print(f"MLflow experiment: {mlflow_experiment_name}") print("-" * 70) datasets = build_wav2vec2_datasets( metadata_path=metadata_path, limit_per_split=limit_per_split, max_duration_seconds=max_duration_seconds, enable_augmentation=enable_augmentation, augmentation_probability=augmentation_probability, ) label_encoding = build_label_encoding(task="emotion") model, processor, classifier_weights_transferred = build_model_and_processor( model_checkpoint=model_checkpoint, freeze_feature_encoder=freeze_feature_encoder, freeze_transformer_layers=freeze_transformer_layers, transfer_classifier_weights=transfer_classifier_weights, ) parameter_counts = count_trainable_parameters(model) print(f"Classifier weights transferred: {classifier_weights_transferred}") print("Parameter counts:") print(f" Total parameters: {parameter_counts['total_parameters']}") print(f" Trainable parameters: {parameter_counts['trainable_parameters']}") print(f" Frozen parameters: {parameter_counts['frozen_parameters']}") data_collator = Wav2Vec2DataCollator(processor=processor) training_args = create_training_arguments( output_dir=output_dir, num_epochs=num_epochs, batch_size=batch_size, learning_rate=learning_rate, weight_decay=weight_decay, warmup_ratio=warmup_ratio, ) trainer = Trainer( model=model, args=training_args, train_dataset=datasets["train"], eval_dataset=datasets["validation"], processing_class=processor, data_collator=data_collator, compute_metrics=compute_metrics, ) trainer.train() validation_metrics = trainer.evaluate(datasets["validation"]) test_report = evaluate_on_test_set( trainer=trainer, test_dataset=datasets["test"], id_to_label=label_encoding.id_to_label, confusion_matrix_path=confusion_matrix_path, ) best_model_dir.mkdir(parents=True, exist_ok=True) trainer.save_model(str(best_model_dir)) processor.save_pretrained(str(best_model_dir)) full_report = { "run_name": run_name, "model_name": "Wav2Vec2 emotion classifier", "base_checkpoint": model_checkpoint, "task": TASK_NAME, "device": get_device_note(), "dataset_source": dataset_source, "metadata_path": str(metadata_path), "output_dir": str(output_dir), "best_model_dir": str(best_model_dir), "mlflow_enabled": enable_mlflow, "mlflow_experiment_name": mlflow_experiment_name if enable_mlflow else None, "labels": label_encoding.id_to_label, "train_samples": len(datasets["train"]), "validation_samples": len(datasets["validation"]), "test_samples": len(datasets["test"]), "num_epochs": num_epochs, "batch_size": batch_size, "learning_rate": learning_rate, "warmup_ratio": warmup_ratio, "freeze_feature_encoder": freeze_feature_encoder, "freeze_transformer_layers": freeze_transformer_layers, "total_parameters": parameter_counts["total_parameters"], "trainable_parameters": parameter_counts["trainable_parameters"], "frozen_parameters": parameter_counts["frozen_parameters"], "enable_augmentation": enable_augmentation, "augmentation_probability": augmentation_probability, "transfer_classifier_weights": transfer_classifier_weights, "classifier_weights_transferred": classifier_weights_transferred, "weight_decay": weight_decay, "validation": validation_metrics, "test": test_report, } report_path.parent.mkdir(parents=True, exist_ok=True) with report_path.open("w", encoding="utf-8") as file: json.dump(full_report, file, indent=2) mlflow_run_id = None if enable_mlflow: mlflow_run_id = log_training_run_to_mlflow( full_report=full_report, report_path=report_path, confusion_matrix_path=confusion_matrix_path, run_name=run_name, mlflow_experiment_name=mlflow_experiment_name, ) full_report["mlflow_run_id"] = mlflow_run_id with report_path.open("w", encoding="utf-8") as file: json.dump(full_report, file, indent=2) print("\nWav2Vec2 Results") print("-" * 70) print(f"Freeze feature encoder: {freeze_feature_encoder}") print(f"Freeze transformer layers: {freeze_transformer_layers}") print(f"Validation accuracy: {validation_metrics.get('eval_accuracy'):.4f}") print(f"Validation macro F1: {validation_metrics.get('eval_macro_f1'):.4f}") print(f"Test accuracy: {test_report['accuracy']:.4f}") print(f"Test macro F1: {test_report['macro_f1']:.4f}") print(f"Classifier weights transferred: {classifier_weights_transferred}") print("-" * 70) print(f"Saved model to: {best_model_dir}") print(f"Saved report to: {report_path}") print(f"Saved confusion matrix to: {confusion_matrix_path}") if enable_mlflow: print(f"Logged MLflow run ID: {mlflow_run_id}") return full_report def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser( description="Fine-tune Wav2Vec2 for speech emotion classification." ) parser.add_argument( "--model-checkpoint", type=str, default=DEFAULT_MODEL_CHECKPOINT, help="Hugging Face Wav2Vec2 checkpoint.", ) parser.add_argument( "--metadata-path", type=Path, default=DEFAULT_METADATA_PATH, help="Path to metadata CSV. Defaults to CREMA-D metadata.", ) parser.add_argument( "--run-name", type=str, default="wav2vec2_emotion", help="Run name used for model folder and report filenames.", ) parser.add_argument( "--enable-mlflow", action="store_true", help="Enable MLflow logging for this training run.", ) parser.add_argument( "--mlflow-experiment-name", type=str, default=DEFAULT_MLFLOW_EXPERIMENT_NAME, help="MLflow experiment name.", ) parser.add_argument( "--output-root", type=Path, default=DEFAULT_OUTPUT_ROOT, help="Root directory for saved Wav2Vec2 models.", ) parser.add_argument( "--limit-per-split", type=int, default=None, help="Optional sample limit per split for quick testing.", ) parser.add_argument( "--num-epochs", type=int, default=5, help="Number of training epochs.", ) parser.add_argument( "--batch-size", type=int, default=4, help="Per-device batch size.", ) parser.add_argument( "--learning-rate", type=float, default=3e-5, help="Learning rate.", ) parser.add_argument( "--weight-decay", type=float, default=0.01, help="Weight decay.", ) parser.add_argument( "--warmup-ratio", type=float, default=0.1, help="Warmup ratio for learning rate scheduler.", ) parser.add_argument( "--freeze-feature-encoder", action=argparse.BooleanOptionalAction, default=True, help="Freeze Wav2Vec2 CNN feature encoder. Use --no-freeze-feature-encoder to unfreeze.", ) parser.add_argument( "--freeze-transformer-layers", type=int, default=0, help="Number of lower Wav2Vec2 transformer layers to freeze.", ) parser.add_argument( "--enable-augmentation", action="store_true", help="Enable light audio augmentation for training samples only.", ) parser.add_argument( "--augmentation-probability", type=float, default=0.5, help="Probability of applying each light augmentation to a training sample.", ) parser.add_argument( "--max-duration-seconds", type=float, default=6.0, help="Maximum audio duration per sample.", ) parser.add_argument( "--transfer-classifier-weights", action="store_true", help="Reuse matching classifier weights from a larger-label checkpoint.", ) return parser.parse_args() if __name__ == "__main__": args = parse_args() train_wav2vec2_emotion_model( model_checkpoint=args.model_checkpoint, metadata_path=args.metadata_path, run_name=args.run_name, output_root=args.output_root, limit_per_split=args.limit_per_split, num_epochs=args.num_epochs, batch_size=args.batch_size, learning_rate=args.learning_rate, weight_decay=args.weight_decay, warmup_ratio=args.warmup_ratio, freeze_feature_encoder=args.freeze_feature_encoder, freeze_transformer_layers=args.freeze_transformer_layers, transfer_classifier_weights=args.transfer_classifier_weights, enable_augmentation=args.enable_augmentation, augmentation_probability=args.augmentation_probability, max_duration_seconds=args.max_duration_seconds, enable_mlflow=args.enable_mlflow, mlflow_experiment_name=args.mlflow_experiment_name, )