Release 0.11.0: audio robustness and benchmark improvements

.github/workflows/ci.yml

@@ -79,7 +79,7 @@ jobs:
                                  "--index-url", "https://download.pytorch.org/whl/cpu"])
           EOF
       - run: pip install '.[gpu]' pytest
-      - run: pytest tests/unit/test_benchmark_resume.py tests/unit/test_cached_sample_labels.py
+      - run: pytest tests/unit/test_benchmark_resume.py tests/unit/test_cached_sample_labels.py tests/unit/test_audio_augmentations.py tests/unit/test_video_prefetch.py
 
   typecheck:
     name: mypy (advisory)

README.md

@@ -60,6 +60,9 @@ gasbench run --video-model ./my_video_model/ --debug
 # Run audio benchmark
 gasbench run --audio-model ./my_audio_model/ --debug
 
+# Add an audio robustness pass on 100 samples per dataset
+gasbench run --audio-model ./my_audio_model/ --n-aug-per-dataset 100 --aug-weight 0.2
+
 # Full benchmark (all datasets)
 gasbench run --image-model ./my_model/ --full
 
@@ -77,6 +80,15 @@ Model directory must contain: `model_config.yaml`, `model.py`, `*.safetensors`
 
 Results are automatically saved to a timestamped JSON file.
 
+The optional audio robustness pass applies an 8 kHz downsample/upsample round
+trip, -6 dB gain, and white noise at 30 dB signal-to-noise ratio after normal
+mono/16 kHz/six-second preprocessing. It supports both raw audio and cached
+preprocessed tensors. The base pass stays unchanged; augmented predictions use
+the same scoring, sample pairing, and checkpoint path as image/video robustness
+passes. Use `--aug-cache-dir` to reuse versioned augmented arrays and
+`--aug-cache-readonly` to prevent writes (cache misses are computed in memory).
+Audio cache entries include the sample seed so different seeds use different noise.
+
 Every benchmark run checkpoints through `BenchmarkRunRecorder` after each inference
 batch. Checkpoints default to `<cache-dir>/runs/<run-id>/checkpoint`; use
 `--checkpoint-dir` to choose another location. To resume, rerun the same command

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "gasbench"
-version = "0.10.1"
+version = "0.11.0"
 description = "GASBench - ML model benchmark evaluation package"
 readme = "README.md"
 requires-python = ">=3.10"

src/gasbench/benchmark.py

@@ -333,6 +333,10 @@ async def execute_benchmark(
             multiclass_scoring=multiclass_scoring,
             checkpoint_dir=checkpoint_dir,
             checkpoint_persist=checkpoint_persist,
+            n_aug_per_dataset=n_aug_per_dataset,
+            aug_weight=aug_weight,
+            aug_cache_dir=aug_cache_dir,
+            aug_cache_readonly=aug_cache_readonly,
         )
         benchmark_score = benchmark_results.get("audio_results", {}).get("benchmark_score", 0.0)
     else:

src/gasbench/benchmarks/audio_bench.py

@@ -4,16 +4,19 @@
 
 from ..logger import get_logger
 from ..processing.media import process_audio_sample
+from ..processing.transforms import apply_audio_robustness_augmentations
 from ..dataset.iterator import load_audio_sample
 
 from ._checkpoint import CheckpointError
-from .recording import BenchmarkRunRecorder, log_dataset_summary
+from .recording import BenchmarkRunRecorder, log_dataset_summary, build_sample_id
+from .aug_cache import aud_aug_cache_path, write_aug_cache
 from .common import (
     BenchmarkRunConfig,
     build_plan,
     create_tracker,
     create_dataset_iterator,
     verify_sample,
+    use_augmentation_cache,
     finalize_run,
     run_batch_and_record,
 )
@@ -31,6 +34,7 @@ def process_batch(
     batch_metadata,
     tracker: BenchmarkRunRecorder,
     batch_id: int,
+    aug_pass: bool = False,
 ):
     """Push a batch of audio samples through the model and record rows in tracker."""
     if not batch_audio:
@@ -58,12 +62,13 @@ def process_batch(
                 sample_index=sample_index,
                 sample=sample,
                 error_message=f"stack-failed: {str(e)[:160]}",
+                aug_pass=aug_pass,
             )
         tracker.checkpoint()
         return
 
     run_batch_and_record(
-        session, input_specs, batch_array, batch_metadata, tracker, batch_id
+        session, input_specs, batch_array, batch_metadata, tracker, batch_id, aug_pass
     )
 
 
@@ -90,6 +95,10 @@ async def run_audio_benchmark(
     multiclass_scoring: bool = False,
     checkpoint_dir: Optional[str] = None,
     checkpoint_persist=None,
+    n_aug_per_dataset: int = 0,
+    aug_weight: float = 0.2,
+    aug_cache_dir: Optional[str] = None,
+    aug_cache_readonly: bool = False,
 ) -> pd.DataFrame:
     """Test model on benchmark audio datasets for AI-generated content detection.
 
@@ -109,7 +118,6 @@ async def run_audio_benchmark(
         else:
             logger.info("Loading benchmark audio datasets")
 
-        # Build run config - audio doesn't use augmentation or crop
         run_config = BenchmarkRunConfig(
             modality="audio",
             mode=mode,
@@ -131,6 +139,10 @@ async def run_audio_benchmark(
             content_category=content_category,
             score_composition=score_composition,
             multiclass_scoring=multiclass_scoring,
+            n_aug_per_dataset=n_aug_per_dataset,
+            aug_weight=aug_weight,
+            aug_cache_dir=aug_cache_dir,
+            aug_cache_readonly=aug_cache_readonly,
         )
 
         plan = build_plan(logger, run_config, input_specs)
@@ -154,116 +166,138 @@ async def run_audio_benchmark(
         logger.info(
             f"Sampling plan targets {plan.sampling_summary.actual_total_samples} samples across {plan.sampling_summary.num_datasets} datasets"
         )
-        for dataset_idx, dataset_cfg in enumerate(plan.available_datasets):
-            dataset_cap = plan.sampling_plan[dataset_cfg.name]
-            logger.info(
-                f"Processing dataset {dataset_idx + 1}/{len(plan.available_datasets)}: "
-                f"{dataset_cfg.name} ({dataset_cap} samples)"
-            )
-
-            try:
-                dataset_iterator = create_dataset_iterator(
-                    run_config, plan, dataset_cfg, aug_pass=False,
+        for aug_pass in ([False, True] if n_aug_per_dataset > 0 else [False]):
+            for dataset_idx, dataset_cfg in enumerate(plan.available_datasets):
+                dataset_cap = n_aug_per_dataset if aug_pass else plan.sampling_plan[dataset_cfg.name]
+                logger.info(
+                    f"{'Robustness pass' if aug_pass else 'Processing dataset'} "
+                    f"{dataset_idx + 1}/{len(plan.available_datasets)}: "
+                    f"{dataset_cfg.name} ({dataset_cap} samples)"
                 )
 
-                if skip_missing and dataset_iterator.get_total_cached_count() == 0:
-                    logger.warning(f"Skipping {dataset_cfg.name} (not cached, --skip-missing enabled)")
-                    continue
-
-                batch_audio = []
-                batch_metadata = []
-                batch_id = 0
-                sample_index = 0
+                try:
+                    dataset_iterator = create_dataset_iterator(
+                        run_config, plan, dataset_cfg, aug_pass=aug_pass,
+                    )
 
-                for sample in dataset_iterator:
-                    sample_index += 1
-                    if tracker.is_checkpointed(
-                        dataset_name=dataset_cfg.name, sample_index=sample_index, sample=sample,
-                    ):
+                    if skip_missing and dataset_iterator.get_total_cached_count() == 0:
+                        logger.warning(f"Skipping {dataset_cfg.name} (not cached, --skip-missing enabled)")
                         continue
-                    verify_sample(sample)
-                    try:
-                        audio_sample = load_audio_sample(sample)
-                        # Check if sample is already preprocessed
-                        if audio_sample.get("is_preprocessed", False):
-                            audio_array = audio_sample.get("preprocessed_waveform")
-                            label = audio_sample.get("label")
-
-                            if audio_array is None or label is None:
-                                continue
-
-                            if hasattr(audio_array, "numpy"):
-                                audio_array = audio_array.numpy()
-                        else:
-                            # Process raw audio bytes
-                            sample_seed_val = None if seed is None else (seed + sample_index)
-                            audio_array, label = process_audio_sample(
-                                audio_sample,
-                                target_sr=target_sr,
-                                seed=sample_seed_val,
-                            )
-
-                            if audio_array is None or label is None:
-                                continue
 
-                            if hasattr(audio_array, "numpy"):
-                                audio_array = audio_array.numpy()
-
-                        sample_seed_val = None if seed is None else (seed + sample_index)
-                        batch_audio.append(audio_array)
-                        batch_metadata.append(
-                            (label, sample, sample_index, dataset_cfg.name, sample_seed_val)
-                        )
+                    batch_audio = []
+                    batch_metadata = []
+                    batch_id = 0
+                    sample_index = 0
+
+                    for sample in dataset_iterator:
+                        sample_index += 1
+                        if tracker.is_checkpointed(
+                            dataset_name=dataset_cfg.name, sample_index=sample_index, sample=sample,
+                            aug_pass=aug_pass,
+                        ):
+                            continue
+                        verify_sample(sample)
+                        try:
+                            audio_sample = load_audio_sample(sample)
+                            # Check if sample is already preprocessed
+                            if audio_sample.get("is_preprocessed", False):
+                                audio_array = audio_sample.get("preprocessed_waveform")
+                                label = audio_sample.get("label")
+
+                                if audio_array is None or label is None:
+                                    continue
+
+                                if hasattr(audio_array, "numpy"):
+                                    audio_array = audio_array.numpy()
+                            else:
+                                # Process raw audio bytes
+                                sample_seed_val = None if seed is None else (seed + sample_index)
+                                audio_array, label = process_audio_sample(
+                                    audio_sample,
+                                    target_sr=target_sr,
+                                    seed=sample_seed_val,
+                                )
+
+                                if audio_array is None or label is None:
+                                    continue
+
+                                if hasattr(audio_array, "numpy"):
+                                    audio_array = audio_array.numpy()
 
-                        if len(batch_audio) >= batch_size:
-                            batch_id += 1
-                            process_batch(
-                                session,
-                                input_specs,
-                                batch_audio,
-                                batch_metadata,
-                                tracker,
-                                batch_id,
+                            sample_seed_val = None if seed is None else (seed + sample_index)
+                            if aug_pass:
+                                cache_path = (
+                                    aud_aug_cache_path(
+                                        aug_cache_dir, build_sample_id(sample), sample_seed_val
+                                    )
+                                    if aug_cache_dir else None
+                                )
+                                if cache_path and use_augmentation_cache(sample, cache_path):
+                                    audio_array = np.load(cache_path, allow_pickle=False)
+                                else:
+                                    audio_array, _, _, _ = apply_audio_robustness_augmentations(
+                                        np.asarray(audio_array).squeeze(),
+                                        target_sr=target_sr,
+                                        seed=sample_seed_val,
+                                    )
+                                    if cache_path and not aug_cache_readonly:
+                                        write_aug_cache(cache_path, audio_array)
+                            batch_audio.append(audio_array)
+                            batch_metadata.append(
+                                (label, sample, sample_index, dataset_cfg.name, sample_seed_val)
                             )
-                            batch_audio = []
-                            batch_metadata = []
 
-                            if tracker.count % 500 == 0:
-                                logger.info(f"Progress: {tracker.count} samples")
+                            if len(batch_audio) >= batch_size:
+                                batch_id += 1
+                                process_batch(
+                                    session,
+                                    input_specs,
+                                    batch_audio,
+                                    batch_metadata,
+                                    tracker,
+                                    batch_id,
+                                    aug_pass=aug_pass,
+                                )
+                                batch_audio = []
+                                batch_metadata = []
+
+                                if tracker.count % 500 == 0:
+                                    logger.info(f"Progress: {tracker.count} samples")
+
+                        except CheckpointError:
+                            raise
+                        except Exception as e:
+                            logger.warning(
+                                f"Failed to process audio sample from {dataset_cfg.name}: {e}"
+                            )
+                            benchmark_results["errors"].append(
+                                f"Audio processing error: {str(e)[:100]}"
+                            )
 
-                    except CheckpointError:
-                        raise
-                    except Exception as e:
-                        logger.warning(
-                            f"Failed to process audio sample from {dataset_cfg.name}: {e}"
-                        )
-                        benchmark_results["errors"].append(
-                            f"Audio processing error: {str(e)[:100]}"
+                    # Process remaining samples
+                    if batch_audio:
+                        batch_id += 1
+                        process_batch(
+                            session,
+                            input_specs,
+                            batch_audio,
+                            batch_metadata,
+                            tracker,
+                            batch_id,
+                            aug_pass=aug_pass,
                         )
 
-                # Process remaining samples
-                if batch_audio:
-                    batch_id += 1
-                    process_batch(
-                        session,
-                        input_specs,
-                        batch_audio,
-                        batch_metadata,
-                        tracker,
-                        batch_id,
+                    log_dataset_summary(
+                        logger, tracker, dataset_cfg.name, include_skipped=False
                     )
 
-                log_dataset_summary(
-                    logger, tracker, dataset_cfg.name, include_skipped=False
-                )
-
-            except CheckpointError:
-                raise
-            except Exception as e:
-                logger.error(f"Failed to process dataset {dataset_cfg.name}: {e}")
-                benchmark_results["errors"].append(
-                    f"Dataset error for {dataset_cfg.name}: {str(e)[:100]}"
-                )
+                except CheckpointError:
+                    raise
+                except Exception as e:
+                    logger.error(f"Failed to process dataset {dataset_cfg.name}: {e}")
+                    benchmark_results["errors"].append(
+                        f"Dataset error for {dataset_cfg.name}: {str(e)[:100]}"
+                    )
 
         df = finalize_run(
             config=run_config,

src/gasbench/benchmarks/aug_cache.py

@@ -8,6 +8,8 @@
 
 _IMG_AUG_VERSION = "img_v1"
 _VID_AUG_VERSION = "vid_v1"
+# Includes the fixed mono 16 kHz / six-second audio preprocessing contract.
+_AUD_AUG_VERSION = "aud_v1"
 
 
 def img_aug_cache_path(cache_dir: str, sample_id: str, target_size: Tuple[int, int]) -> str:
@@ -20,6 +22,11 @@ def vid_aug_cache_path(cache_dir: str, sample_id: str, target_size: Tuple[int, i
     return os.path.join(cache_dir, "vid", sample_id[:2], f"{sample_id}_{_VID_AUG_VERSION}_{H}x{W}.npy")
 
 
+def aud_aug_cache_path(cache_dir: str, sample_id: str, seed: int) -> str:
+    # Noise depends on the sample seed, so a different run seed cannot reuse it.
+    return os.path.join(cache_dir, "aud", sample_id[:2], f"{sample_id}_{_AUD_AUG_VERSION}_seed{seed}.npy")
+
+
 def write_aug_cache(path: str, array: np.ndarray) -> None:
     """Atomically write an augmented array to the cache via temp-file + os.replace."""
     os.makedirs(os.path.dirname(path), exist_ok=True)