#!/usr/bin/env python3 """Quick validation script to inspect preference dataset structure. This script provides utilities to: 1. Inspect dataset files (mmap, metadata, info) 2. Validate dataset integrity 3. Display dataset statistics """ import os import sys from typing import Optional import json from collections import Counter from merge_preference_datasets import ( read_jsonl, read_json, load_sample_from_mmap, N_TOKENS_AUDIO, SEMANTIC_N_CODEBOOKS, SEMANTIC_PAD_TOKEN, ) import numpy as np def inspect_dataset( dataset_dir: str, is_val: bool = False, t_data_memmap: int = N_TOKENS_AUDIO, show_samples: int = 5, ) -> None: """Inspect a preference dataset and display its structure. Args: dataset_dir: Path to the dataset directory is_val: Whether this is a validation set t_data_memmap: Number of tokens per sample show_samples: Number of sample metadata to display """ dset_type = "val" if is_val else "tr" mmap_path = os.path.join(dataset_dir, f"data_{dset_type}.bin") meta_path = os.path.join(dataset_dir, f"meta_{dset_type}.jsonl") info_path = os.path.join(dataset_dir, f"info_{dset_type}.json") print("=" * 70) print(f"Dataset Inspection: {dataset_dir}") print(f"Type: {'Validation' if is_val else 'Training'}") print("=" * 70) # Check file existence print("\n[File Existence Check]") for name, path in [("Mmap", mmap_path), ("Metadata", meta_path), ("Info", info_path)]: exists = os.path.exists(path) size_str = "" if exists: size = os.path.getsize(path) size_str = f" ({size / 1024**3:.2f} GB)" if size > 1024**3 else f" ({size / 1024**2:.2f} MB)" print(f" {'✓' if exists else '✗'} {name}: {path}{size_str}") if not exists: print(f"\n❌ Missing file: {path}") return # Load and inspect metadata print("\n[Metadata Statistics]") metadata = read_jsonl(meta_path) n_samples = len(metadata) print(f" Total samples: {n_samples:,}") # Count by dataset type dataset_counter = Counter(meta["dataset"] for meta in metadata) print(f"\n Dataset distribution:") for dataset_name, count in sorted(dataset_counter.items()): percentage = (count / n_samples * 100) if n_samples > 0 else 0 print(f" {dataset_name}: {count:,} ({percentage:.1f}%)") # Count by task task_counter = Counter(meta.get("task", "gen") for meta in metadata) print(f"\n Task distribution:") for task, count in sorted(task_counter.items()): percentage = (count / n_samples * 100) if n_samples > 0 else 0 print(f" {task or 'gen'}: {count:,} ({percentage:.1f}%)") # Count by gender gender_counter = Counter(meta.get("gender", "unspecified") for meta in metadata) print(f"\n Gender distribution:") for gender, count in sorted(gender_counter.items()): percentage = (count / n_samples * 100) if n_samples > 0 else 0 print(f" {gender}: {count:,} ({percentage:.1f}%)") # Load and inspect info print("\n[Info File Statistics]") info = read_json(info_path) total_indices = sum(len(v.get("idx_list", [])) for v in info.values()) print(f" Total indices in info: {total_indices:,}") print(f" Dataset categories: {len(info)}") for dataset_name, dataset_info in sorted(info.items()): idx_list = dataset_info.get("idx_list", []) print(f" {dataset_name}: {len(idx_list):,} indices") # Validate mmap size print("\n[Mmap Validation]") sample_size = t_data_memmap * SEMANTIC_N_CODEBOOKS expected_size = n_samples * sample_size mmap = np.memmap(mmap_path, dtype=np.uint16, mode='r') actual_size = len(mmap) del mmap size_match = actual_size == expected_size print(f" Expected size: {expected_size:,} elements ({expected_size * 2 / 1024**3:.2f} GB)") print(f" Actual size: {actual_size:,} elements ({actual_size * 2 / 1024**3:.2f} GB)") print(f" {'✓' if size_match else '✗'} Size match: {size_match}") if not size_match: print(f"\n⚠️ Warning: Mmap size mismatch! Expected {expected_size:,}, got {actual_size:,}") return # Sample inspection print(f"\n[Sample Metadata (first {show_samples})]") for i, meta in enumerate(metadata[:show_samples]): print(f"\n Sample {i}:") print(f" ID: {meta.get('id', 'N/A')}") print(f" Dataset: {meta.get('dataset', 'N/A')}") print(f" Task: {meta.get('task', 'gen')}") print(f" Text: {meta.get('text', '')[:60]}...") print(f" Tags: {meta.get('tags', [])}") print(f" Gender: {meta.get('gender', 'N/A')}") print(f" Gen Start Index: {meta.get('generated_start_index', 0)}") # Load and inspect actual sample data print(f"\n[Sample Data Inspection (first 3)]") for i in range(min(3, n_samples)): sample = load_sample_from_mmap(mmap_path, i, t_data_memmap) non_pad = np.sum(sample != SEMANTIC_PAD_TOKEN) pad_count = np.sum(sample == SEMANTIC_PAD_TOKEN) min_val = np.min(sample) max_val = np.max(sample) print(f"\n Sample {i}:") print(f" Shape: {sample.shape}") print(f" Non-pad tokens: {non_pad:,}") print(f" Pad tokens: {pad_count:,}") print(f" Value range: [{min_val}, {max_val}]") print(f" Duration: ~{non_pad / 25:.1f}s") print("\n" + "=" * 70) print("✅ Dataset inspection complete!") print("=" * 70) def validate_dataset_indices( dataset_dir: str, is_val: bool = False, ) -> bool: """Validate that all info indices point to valid metadata entries. Args: dataset_dir: Path to the dataset directory is_val: Whether this is a validation set Returns: True if validation passes, False otherwise """ dset_type = "val" if is_val else "tr" meta_path = os.path.join(dataset_dir, f"meta_{dset_type}.jsonl") info_path = os.path.join(dataset_dir, f"info_{dset_type}.json") print("\n[Index Validation]") metadata = read_jsonl(meta_path) info = read_json(info_path) n_samples = len(metadata) print(f" Total samples: {n_samples:,}") errors = [] for dataset_name, dataset_info in info.items(): idx_list = dataset_info.get("idx_list", []) print(f"\n Checking {dataset_name}: {len(idx_list):,} indices") for idx in idx_list: # Check index in range if idx < 0 or idx >= n_samples: errors.append(f" ✗ Index {idx} out of range [0, {n_samples})") continue # Check dataset name matches meta = metadata[idx] if meta["dataset"] != dataset_name: errors.append( f" ✗ Index {idx} has dataset '{meta['dataset']}' but expected '{dataset_name}'" ) if errors: print("\n❌ Validation failed with errors:") for error in errors[:10]: # Show first 10 errors print(error) if len(errors) > 10: print(f" ... and {len(errors) - 10} more errors") return False else: print("\n✓ All indices are valid!") return True if __name__ == "__main__": import argparse parser = argparse.ArgumentParser(description="Inspect and validate preference dataset") parser.add_argument( "dataset_dir", type=str, help="Path to the dataset directory to inspect" ) parser.add_argument( "--is_val", action="store_true", help="Whether this is validation set (default: False for train set)" ) parser.add_argument( "--t_data_memmap", type=int, default=N_TOKENS_AUDIO, help=f"Number of tokens per sample (default: {N_TOKENS_AUDIO})" ) parser.add_argument( "--show_samples", type=int, default=5, help="Number of sample metadata to display (default: 5)" ) parser.add_argument( "--validate_indices", action="store_true", help="Perform detailed index validation" ) args = parser.parse_args() inspect_dataset( dataset_dir=args.dataset_dir, is_val=args.is_val, t_data_memmap=args.t_data_memmap, show_samples=args.show_samples, ) if args.validate_indices: validate_dataset_indices( dataset_dir=args.dataset_dir, is_val=args.is_val, )