fix(audio): harden AudioBench and MMAU evaluation

Signed-off-by: Piotr Żelasko <pzelasko@nvidia.com>

Author: pzelasko

nemo_skills/dataset/audiobench/nonjudge/__init__.py

@@ -26,5 +26,7 @@
 # Evaluation settings
 EVAL_ARGS = "++eval_type=audio ++eval_config.normalization_mode=audiobench"
 
-# Generation settings - OpenAI format for audio-language models
-GENERATION_ARGS = "++prompt_format=openai ++enable_audio=true"
+# Generation settings - OpenAI format for audio-language models.
+# Soft-fail keeps one over-limit or malformed audio request from aborting an
+# entire chunk; audio metrics count the empty generation as incorrect.
+GENERATION_ARGS = "++prompt_format=openai ++enable_audio=true ++server.enable_soft_fail=true"

nemo_skills/dataset/audiobench/prepare.py

@@ -25,6 +25,7 @@
 """
 
 import argparse
+import copy
 import json
 import os
 import shutil
@@ -155,12 +156,21 @@ def create_manifest_entry(
         Manifest entry dict with proper format for nemo-skills
     """
     instruction = sample.get("instruction", sample.get("text", "Process the audio"))
+    if isinstance(instruction, dict):
+        try:
+            instruction = instruction["text"]
+        except KeyError as exc:
+            raise ValueError(
+                f"Instruction dict missing required 'text' field for {dataset_name} sample {sample_id}: {instruction}"
+            ) from exc
+    if not isinstance(instruction, str):
+        raise ValueError(f"Instruction must be a string for {dataset_name} sample {sample_id}: {instruction!r}")
     reference = sample.get("reference", sample.get("answer", ""))
     task_type = sample.get("task_type", "unknown")
 
-    # Create absolute audio path with /data/ prefix for cluster deployment
-    # Format: /data/audiobench/{category}/audio/{dataset_name}/{filename}
-    audio_rel_path = f"/data/audiobench/{category}/audio/{dataset_name}/{audio_filename}"
+    # Paths are resolved relative to the manifest directory by the inference code.
+    # The combined category manifest lives at audiobench/{category}/test.jsonl.
+    audio_rel_path = f"audio/{dataset_name}/{audio_filename}"
 
     # Create audio metadata (both singular and plural forms for compatibility)
     audio_metadata = {"path": audio_rel_path, "duration": duration}
@@ -203,6 +213,28 @@ def create_manifest_entry(
     return entry
 
 
+def make_dataset_manifest_entry(entry: Dict) -> Dict:
+    """Adjust category-relative audio paths for per-dataset manifests."""
+    entry = copy.deepcopy(entry)
+
+    def rewrite(path: str) -> str:
+        return f"../{path}" if path.startswith("audio/") else path
+
+    if isinstance(entry.get("audio_path"), list):
+        entry["audio_path"] = [rewrite(path) for path in entry["audio_path"]]
+    elif isinstance(entry.get("audio_path"), str):
+        entry["audio_path"] = rewrite(entry["audio_path"])
+
+    for message in entry.get("messages", []):
+        if "audio" in message and "path" in message["audio"]:
+            message["audio"]["path"] = rewrite(message["audio"]["path"])
+        for audio in message.get("audios", []):
+            if "path" in audio:
+                audio["path"] = rewrite(audio["path"])
+
+    return entry
+
+
 def process_dataset(
     dataset_name: str,
     output_dir: Path,
@@ -473,7 +505,7 @@ def process_dataset(
     manifest_path = dataset_dir / f"{split}.jsonl"
     with open(manifest_path, "w", encoding="utf-8") as f:
         for entry in manifest_entries:
-            f.write(json.dumps(entry, ensure_ascii=False) + "\n")
+            f.write(json.dumps(make_dataset_manifest_entry(entry), ensure_ascii=False) + "\n")
 
     print(f"✓ Saved {successful} samples to {manifest_path}")
     if failed > 0:
@@ -566,6 +598,7 @@ def main():
 
     total_samples = 0
     total_datasets = 0
+    combined_entries = {"judge": [], "nonjudge": []}
 
     for name in target_datasets:
         # Normalize dataset name: allow passing without _test suffix
@@ -575,23 +608,33 @@ def main():
             if f"{dataset_name}_test" in JUDGE_DATASETS or f"{dataset_name}_test" in NONJUDGE_DATASETS:
                 dataset_name = f"{dataset_name}_test"
 
-        # Determine category for logging
-        category = "judge" if name in JUDGE_DATASETS else "nonjudge"
+        if dataset_name in JUDGE_DATASETS:
+            category = "judge"
+        elif dataset_name in NONJUDGE_DATASETS:
+            category = "nonjudge"
+        else:
+            raise ValueError(f"Unsupported AudioBench dataset name: {name}")
+
+        num_samples, manifest_entries = process_dataset(
+            dataset_name=dataset_name,
+            output_dir=output_dir,
+            save_audio=args.save_audio,
+            split=args.split,
+            max_samples=args.max_samples,
+        )
+        total_samples += num_samples
+        total_datasets += 1
+        combined_entries[category].extend(manifest_entries)
+        print(f"✓ Completed {dataset_name}: {num_samples} samples")
 
-        try:
-            num_samples, _ = process_dataset(
-                dataset_name=dataset_name,
-                output_dir=output_dir,
-                save_audio=args.save_audio,
-                split=args.split,
-                max_samples=args.max_samples,
-            )
-            total_samples += num_samples
-            total_datasets += 1
-            print(f"✓ Completed {dataset_name}: {num_samples} samples")
-        except Exception as e:
-            print(f"✗ Failed {dataset_name}: {e}")
+    for category, entries in combined_entries.items():
+        if not entries:
             continue
+        combined_manifest = output_dir / category / f"{args.split}.jsonl"
+        with open(combined_manifest, "w", encoding="utf-8") as f:
+            for entry in entries:
+                f.write(json.dumps(entry, ensure_ascii=False) + "\n")
+        print(f"✓ Saved combined {category} manifest with {len(entries)} samples to {combined_manifest}")
 
     print("\n" + "=" * 60)
     print("AudioBench Preparation Summary")

nemo_skills/dataset/mmau-pro/closed_form/__init__.py

@@ -13,8 +13,7 @@
 # limitations under the License.
 METRICS_TYPE = "mmau_pro_closed_form"
 SCORE_MODULE = "nemo_skills.evaluation.metrics.mmau_pro_metrics"
-GENERATION_ARGS = "++prompt_format=openai ++enable_audio=true"
-EVAL_ARGS = "++eval_type=mmau-pro"
+GENERATION_ARGS = "++prompt_format=openai ++enable_audio=true ++server.enable_soft_fail=true"
 
 # NVEmbed judge configuration for closed-form evaluation
 JUDGE_PIPELINE_ARGS = {

nemo_skills/dataset/mmau-pro/prepare.py

@@ -86,9 +86,11 @@ def format_entry(entry, with_audio=False):
         audio_path = entry["audio_path"]
 
         if isinstance(audio_path, list) and audio_path:
-            user_message["audios"] = [{"path": path, "duration": 10.0} for path in audio_path]
+            formatted_entry["audio_path"] = [f"../{path}" for path in audio_path]
+            user_message["audios"] = [{"path": path, "duration": 10.0} for path in formatted_entry["audio_path"]]
         elif isinstance(audio_path, str):
-            user_message["audio"] = {"path": audio_path, "duration": 10.0}
+            formatted_entry["audio_path"] = f"../{audio_path}"
+            user_message["audio"] = {"path": formatted_entry["audio_path"], "duration": 10.0}
 
     formatted_entry["messages"] = [user_message]
     return formatted_entry

nemo_skills/evaluation/evaluator/audio.py

@@ -735,12 +735,28 @@ def eval_audio(cfg):
     asyncio.run(evaluator.eval_full())
 
 
+def coerce_text_field(value: Any) -> str:
+    """Convert loose AudioBench text/reference fields to plain text."""
+    if value is None:
+        return ""
+    if isinstance(value, str):
+        return value.strip()
+    if isinstance(value, dict):
+        for key in ("text", "expected_answer", "answer", "transcript", "reference"):
+            if key in value and value[key] is not None:
+                return coerce_text_field(value[key])
+        return " ".join(coerce_text_field(item) for item in value.values() if item is not None).strip()
+    if isinstance(value, (list, tuple)):
+        return " ".join(coerce_text_field(item) for item in value if item is not None).strip()
+    return str(value).strip()
+
+
 def evaluate_sample(sample: dict[str, Any], config: AudioEvaluatorConfig) -> dict[str, Any]:
     """Evaluate single sample based on task_type. Returns dict of updates to merge."""
     updates = {}
     task_type = sample.get("task_type", "unknown")
-    generation = sample["generation"].strip()
-    expected_answer = sample.get("expected_answer", "").strip()
+    generation = coerce_text_field(sample.get("generation", ""))
+    expected_answer = coerce_text_field(sample.get("expected_answer", ""))
 
     # Extract ASR text from generation
     # E.g Qwen ASR uses <asr_text> tags to indicate the ASR text

nemo_skills/evaluation/evaluator/mmau_pro.py

@@ -56,6 +56,10 @@ def eval_mmau_pro(cfg):
 
 def evaluate_instruction_following_sample(sample: dict[str, Any]) -> dict[str, Any]:
     """Evaluate a single instruction following sample."""
+    if sample is None:
+        LOG.warning("Encountered null MMAU-Pro sample; marking it incorrect")
+        return {"is_correct": False, "error": "null_sample"}
+
     sample = sample.copy()
     generation = sample.get("generation", "").strip()
 

nemo_skills/evaluation/evaluator/nvembed_judge.py

@@ -42,7 +42,18 @@ def install_packages():
     """Install required packages for NVEmbed evaluation."""
     LOG.info("Installing required packages...")
     subprocess.run(
-        ["pip", "install", "-q", "datasets", "einops", "transformers==4.42.4", "tqdm"],
+        [
+            sys.executable,
+            "-m",
+            "pip",
+            "install",
+            "-q",
+            "numpy<2",
+            "datasets",
+            "einops",
+            "transformers==4.42.4",
+            "tqdm",
+        ],
         check=True,
         capture_output=True,
         text=True,
@@ -120,15 +131,23 @@ def evaluate_sample_with_nvembed(sample: dict[str, Any], model_name: str = "nvid
     if "nvembed_confidence" in sample:
         return sample
 
-    generation = sample.get("generation", "").strip()
+    generation_value = sample.get("generation", "")
+    generation = generation_value.strip() if isinstance(generation_value, str) else ""
     choices = sample.get("choices", [])
     expected_answer = sample.get("expected_answer", "")
 
-    # Fail fast if data is malformed - this indicates a pipeline error
+    # Empty model outputs are valid failed predictions. Keep malformed-data
+    # checks strict, but do not let one blank generation abort the whole eval.
     if not generation:
-        raise ValueError(
-            f"Sample missing generation field or has empty generation. Sample ID: {sample.get('id', 'unknown')}"
+        sample.update(
+            {
+                "nvembed_matched_choice": "",
+                "nvembed_confidence": 0.0,
+                "is_correct": False,
+                "nvembed_error": "empty_generation",
+            }
         )
+        return sample
 
     if not choices:
         raise ValueError(

nemo_skills/pipeline/judges/nvembed_judge.py

@@ -15,6 +15,7 @@
 """NVEmbed judge implementation for embedding-based similarity matching."""
 
 import logging
+import shlex
 
 from nemo_skills.pipeline.utils import add_task
 from nemo_skills.pipeline.utils.generation import get_remaining_jobs
@@ -35,6 +36,8 @@ def create_judge_tasks(
     judge_server_gpus,
     judge_server_nodes,
     partition,
+    account,
+    judge_container,
     run_after,
     reuse_code_exp,
     reuse_code,
@@ -59,6 +62,8 @@ def create_judge_tasks(
         judge_server_gpus: Number of GPUs for judge
         judge_server_nodes: Number of nodes for judge
         partition: SLURM partition
+        account: SLURM account
+        judge_container: Container to use for the judge
         run_after: Dependencies to run after
         reuse_code_exp: Experiment to reuse code from
         reuse_code: Whether to reuse code
@@ -95,32 +100,50 @@ def create_judge_tasks(
         return []
 
     # Build command to run NVEmbed judge script
-    script_args = [f"--output-dir {output_dir_path}"]
+    script_args = ["--output-dir", str(output_dir_path)]
 
     if input_file is None:
-        input_dir = judge_pipeline_args.get("input_dir")
-        script_args.append(f"--input-dir {input_dir}")
-        script_args.append(f"--num-seeds {num_seeds}")
+        try:
+            input_dir = judge_pipeline_args["input_dir"]
+        except KeyError as exc:
+            raise ValueError(
+                "NVEmbed judge requires judge_pipeline_args['input_dir'] when input_file is unset"
+            ) from exc
+        if input_dir is None:
+            raise ValueError("NVEmbed judge requires a non-null input_dir when input_file is unset")
+        script_args.extend(["--input-dir", str(input_dir)])
+        script_args.extend(["--num-seeds", str(num_seeds)])
     else:
-        script_args.append(f"--input-file {input_file}")
+        script_args.extend(["--input-file", str(input_file)])
 
     # Add skip-existing flag unless rerun_done is set
     if not rerun_done:
         script_args.append("--skip-existing")
-
-    run_cmd = f"python3 -I /nemo_run/code/nemo_skills/evaluation/evaluator/nvembed_judge.py {' '.join(script_args)}"
+    script_args.append("--skip-install")
+
+    run_cmd = (
+        "NVEMBED_DEPS_DIR=/tmp/nvembed_deps_${SLURM_JOB_ID:-$$} && "
+        'mkdir -p "$NVEMBED_DEPS_DIR" && '
+        'python3 -m pip install -q --upgrade --target "$NVEMBED_DEPS_DIR" '
+        "'numpy==1.26.4' datasets einops transformers==4.42.4 tqdm && "
+        'export PYTHONPATH="$NVEMBED_DEPS_DIR:${PYTHONPATH:-}" && '
+        "export HF_HUB_OFFLINE=0 TRANSFORMERS_OFFLINE=0 HF_DATASETS_OFFLINE=0 && "
+        "python3 /nemo_run/code/nemo_skills/evaluation/evaluator/nvembed_judge.py "
+        f"{' '.join(shlex.quote(arg) for arg in script_args)}"
+    )
 
     # Create task with GPU support for NVEmbed
     judge_task = add_task(
         exp,
         cmd=run_cmd,
         task_name=f"{expname}-{benchmark}-nvembed-judge",
         log_dir=log_dir + "/judge",
-        container=cluster_config["containers"]["vllm"],
+        container=judge_container or cluster_config["containers"]["vllm"],
         cluster_config=cluster_config,
         num_gpus=judge_server_gpus or 1,
         num_nodes=judge_server_nodes or 1,
         partition=partition,
+        account=account,
         run_after=run_after,
         reuse_code_exp=reuse_code_exp,
         reuse_code=reuse_code,