#!/usr/bin/env python3 """ Validate vocal stem captions and pitch range data in processed metadata files. This script validates that vocal stem captioning and pitch range data have been properly added to the metadata, providing comprehensive statistics and visualizations. """ import argparse import json import os from pathlib import Path from typing import Dict, List, Any, Optional from collections import defaultdict, Counter import matplotlib.pyplot as plt import numpy as np from tqdm import tqdm from datetime import datetime def analyze_vocal_captions(record: Dict) -> Dict[str, Any]: """ Analyze vocal stem captions in a record. Args: record: Single record from JSONL Returns: Dictionary with caption analysis results """ analysis = { "has_stems_captions": False, "has_voice_keywords": False, "caption_count": 0, "stem_types_with_captions": [], "keywords_per_stem": {}, } if "stems_captions" in record: analysis["has_stems_captions"] = True for stem_name, caption_list in record["stems_captions"].items(): if isinstance(caption_list, list): for caption_data in caption_list: if isinstance(caption_data, dict): # Check for voice_description_keywords prompt type if caption_data.get("prompt_type") == "voice_description_keywords": analysis["has_voice_keywords"] = True analysis["caption_count"] += 1 analysis["stem_types_with_captions"].append(stem_name) # Count keywords in caption caption_text = caption_data.get("caption", "") if caption_text: keywords = [k.strip() for k in caption_text.split(",")] analysis["keywords_per_stem"][stem_name] = len(keywords) return analysis def analyze_pitch_range(record: Dict) -> Dict[str, Any]: """ Analyze vocal pitch range data in a record. Args: record: Single record from JSONL Returns: Dictionary with pitch range analysis """ analysis = {"has_pitch_range": False, "algorithms": {}, "pitch_data": {}} if "vocal_pitch_range" in record: analysis["has_pitch_range"] = True pitch_data = record["vocal_pitch_range"] # Handle both dict and list formats if isinstance(pitch_data, dict): for algo_name, algo_data in pitch_data.items(): if isinstance(algo_data, dict) and "min" in algo_data and "max" in algo_data: analysis["algorithms"][algo_name] = True analysis["pitch_data"][algo_name] = { "min": algo_data["min"], "max": algo_data["max"], } elif isinstance(pitch_data, list): # Handle list format with model and strategy fields for item in pitch_data: if isinstance(item, dict): # Create algorithm name from model and strategy model = item.get("model", "unknown") strategy = item.get("strategy", "unknown") algo_name = f"{model}_{strategy}" # Extract min/max frequencies min_f = item.get("min_f") max_f = item.get("max_f") if min_f is not None and max_f is not None: analysis["algorithms"][algo_name] = True analysis["pitch_data"][algo_name] = { "min": min_f, "max": max_f, "min_note": item.get("min_note"), "max_note": item.get("max_note"), "range_semitones": item.get("range_semitones"), } return analysis def plot_pitch_distributions(pitch_stats: Dict, output_dir: Path): """ Create visualizations for pitch range distributions. Args: pitch_stats: Statistics about pitch ranges output_dir: Directory to save plots """ output_dir.mkdir(parents=True, exist_ok=True) # Prepare data for plotting algorithms = list(pitch_stats["algorithm_distributions"].keys()) if not algorithms: print("No pitch data to plot") return # Create figure with subplots for each algorithm n_algos = len(algorithms) fig, axes = plt.subplots(n_algos, 2, figsize=(15, 5 * n_algos)) if n_algos == 1: axes = axes.reshape(1, -1) for idx, algo in enumerate(algorithms): dist = pitch_stats["algorithm_distributions"][algo] # Plot min pitch distribution ax_min = axes[idx, 0] if dist["min_values"]: ax_min.hist(dist["min_values"], bins=50, edgecolor="black", alpha=0.7, color="blue") ax_min.set_title(f"{algo} - Min Pitch Distribution") ax_min.set_xlabel("Min Pitch (Hz)") ax_min.set_ylabel("Count") ax_min.axvline( np.mean(dist["min_values"]), color="red", linestyle="--", label=f'Mean: {np.mean(dist["min_values"]):.1f} Hz', ) ax_min.legend() ax_min.grid(True, alpha=0.3) # Plot max pitch distribution ax_max = axes[idx, 1] if dist["max_values"]: ax_max.hist(dist["max_values"], bins=50, edgecolor="black", alpha=0.7, color="green") ax_max.set_title(f"{algo} - Max Pitch Distribution") ax_max.set_xlabel("Max Pitch (Hz)") ax_max.set_ylabel("Count") ax_max.axvline( np.mean(dist["max_values"]), color="red", linestyle="--", label=f'Mean: {np.mean(dist["max_values"]):.1f} Hz', ) ax_max.legend() ax_max.grid(True, alpha=0.3) plt.tight_layout() plot_path = output_dir / "pitch_distributions.png" plt.savefig(plot_path, dpi=150, bbox_inches="tight") print(f"Saved pitch distribution plots to {plot_path}") plt.close() # Create summary plot comparing all algorithms if len(algorithms) > 1: fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 6)) # Box plot for min values min_data = [pitch_stats["algorithm_distributions"][algo]["min_values"] for algo in algorithms] ax1.boxplot(min_data, labels=algorithms) ax1.set_title("Min Pitch Comparison Across Algorithms") ax1.set_ylabel("Min Pitch (Hz)") ax1.set_xlabel("Algorithm") ax1.grid(True, alpha=0.3) # Box plot for max values max_data = [pitch_stats["algorithm_distributions"][algo]["max_values"] for algo in algorithms] ax2.boxplot(max_data, labels=algorithms) ax2.set_title("Max Pitch Comparison Across Algorithms") ax2.set_ylabel("Max Pitch (Hz)") ax2.set_xlabel("Algorithm") ax2.grid(True, alpha=0.3) plt.tight_layout() comparison_path = output_dir / "pitch_comparison.png" plt.savefig(comparison_path, dpi=150, bbox_inches="tight") print(f"Saved pitch comparison plot to {comparison_path}") plt.close() def validate_dataset( input_path: str, original_path: str = "/app2/suno/data/auk_v0/metas_v6_tr.jsonl", sample_size: int = 10000, output_dir: Optional[str] = None, ) -> Dict[str, Any]: """ Validate the processed dataset for captions and pitch data. Args: input_path: Path to processed JSONL file original_path: Path to original JSONL file for line count comparison sample_size: Number of records to sample for detailed analysis output_dir: Directory for output plots and reports Returns: Dictionary with validation results """ input_file = Path(input_path) output_path = Path(output_dir) if output_dir else input_file.parent / "validation_output" output_path.mkdir(parents=True, exist_ok=True) # Initialize statistics stats = { "total_records": 0, "sampled_records": 0, "caption_stats": { "records_with_stems_captions": 0, "records_with_voice_keywords": 0, "total_captions": 0, "stem_type_counts": defaultdict(int), "keyword_counts": [], }, "pitch_stats": { "records_with_pitch": 0, "algorithm_counts": defaultdict(int), "algorithm_distributions": defaultdict(lambda: {"min_values": [], "max_values": []}), }, "sample_analysis": [], } print(f"Validating dataset: {input_path}") print(f"Sampling first {sample_size:,} records for detailed analysis...") # Process the dataset with open(input_file, "r") as f: for line_num, line in enumerate(tqdm(f, desc="Processing records"), 1): stats["total_records"] += 1 if not line.strip(): continue try: record = json.loads(line.strip()) # Detailed analysis for sampled records if line_num <= sample_size: stats["sampled_records"] += 1 # Analyze captions caption_analysis = analyze_vocal_captions(record) if caption_analysis["has_stems_captions"]: stats["caption_stats"]["records_with_stems_captions"] += 1 if caption_analysis["has_voice_keywords"]: stats["caption_stats"]["records_with_voice_keywords"] += 1 stats["caption_stats"]["total_captions"] += caption_analysis["caption_count"] for stem_type in caption_analysis["stem_types_with_captions"]: stats["caption_stats"]["stem_type_counts"][stem_type] += 1 for stem, count in caption_analysis["keywords_per_stem"].items(): stats["caption_stats"]["keyword_counts"].append(count) # Analyze pitch pitch_analysis = analyze_pitch_range(record) if pitch_analysis["has_pitch_range"]: stats["pitch_stats"]["records_with_pitch"] += 1 for algo in pitch_analysis["algorithms"]: stats["pitch_stats"]["algorithm_counts"][algo] += 1 for algo, data in pitch_analysis["pitch_data"].items(): if data["min"] is not None: stats["pitch_stats"]["algorithm_distributions"][algo][ "min_values" ].append(data["min"]) if data["max"] is not None: stats["pitch_stats"]["algorithm_distributions"][algo][ "max_values" ].append(data["max"]) # Store sample record info if line_num <= 10: # Store first 10 for inspection stats["sample_analysis"].append( { "id": record.get("id", f"record_{line_num}"), "line": line_num, "has_captions": caption_analysis["has_voice_keywords"], "caption_count": caption_analysis["caption_count"], "has_pitch": pitch_analysis["has_pitch_range"], "algorithms": list(pitch_analysis["algorithms"].keys()), } ) except json.JSONDecodeError as e: print(f"Error parsing line {line_num}: {e}") continue # Check line count against original print(f"\nChecking line count against original file...") original_line_count = 0 if Path(original_path).exists(): with open(original_path, "r") as f: for _ in f: original_line_count += 1 stats["line_count_validation"] = { "processed_lines": stats["total_records"], "original_lines": original_line_count, "match": stats["total_records"] == original_line_count, "difference": stats["total_records"] - original_line_count, } else: print(f"Original file not found: {original_path}") stats["line_count_validation"] = { "processed_lines": stats["total_records"], "original_lines": "Not found", "match": "Unknown", } # Generate visualizations if stats["pitch_stats"]["algorithm_distributions"]: plot_pitch_distributions(stats["pitch_stats"], output_path) # Generate summary report report_path = output_path / "validation_report.json" with open(report_path, "w") as f: json.dump(stats, f, indent=2, default=str) # Print summary print_validation_summary(stats) return stats def print_validation_summary(stats: Dict): """Print a formatted summary of validation results.""" print("\n" + "=" * 80) print("VALIDATION SUMMARY") print("=" * 80) # Dataset size print(f"\nDataset Size:") print(f" Total records: {stats['total_records']:,}") print(f" Sampled for analysis: {stats['sampled_records']:,}") # Line count validation if "line_count_validation" in stats: lc = stats["line_count_validation"] print(f"\nLine Count Validation:") print(f" Processed file: {lc['processed_lines']:,} lines") print(f" Original file: {lc.get('original_lines', 'N/A')}") if isinstance(lc.get("original_lines"), int): if lc["match"]: print(f" ✓ Line counts match!") else: print(f" ✗ Line count difference: {lc['difference']:+,}") # Caption statistics caption_stats = stats["caption_stats"] print(f"\nVocal Stem Caption Statistics (first {stats['sampled_records']:,} records):") print(f" Records with stems_captions field: {caption_stats['records_with_stems_captions']:,}") print(f" Records with voice_description_keywords: {caption_stats['records_with_voice_keywords']:,}") print(f" Total voice keyword captions: {caption_stats['total_captions']:,}") if caption_stats["records_with_voice_keywords"] > 0: avg_captions = caption_stats["total_captions"] / caption_stats["records_with_voice_keywords"] print(f" Average captions per record: {avg_captions:.2f}") if caption_stats["keyword_counts"]: avg_keywords = np.mean(caption_stats["keyword_counts"]) print(f" Average keywords per caption: {avg_keywords:.1f}") if caption_stats["stem_type_counts"]: print(f"\n Stem types with captions:") for stem_type, count in sorted( caption_stats["stem_type_counts"].items(), key=lambda x: x[1], reverse=True ): print(f" - {stem_type}: {count:,}") # Pitch statistics pitch_stats = stats["pitch_stats"] print(f"\nVocal Pitch Range Statistics (first {stats['sampled_records']:,} records):") print(f" Records with vocal_pitch_range: {pitch_stats['records_with_pitch']:,}") if pitch_stats["algorithm_counts"]: print(f"\n Algorithm coverage:") for algo, count in sorted(pitch_stats["algorithm_counts"].items()): percentage = (count / stats["sampled_records"]) * 100 print(f" - {algo}: {count:,} records ({percentage:.1f}%)") if pitch_stats["algorithm_distributions"]: print(f"\n Pitch range statistics by algorithm:") for algo, dist in pitch_stats["algorithm_distributions"].items(): if dist["min_values"] and dist["max_values"]: min_mean = np.mean(dist["min_values"]) min_std = np.std(dist["min_values"]) max_mean = np.mean(dist["max_values"]) max_std = np.std(dist["max_values"]) print(f" {algo}:") print(f" Min pitch: {min_mean:.1f} ± {min_std:.1f} Hz") print(f" Max pitch: {max_mean:.1f} ± {max_std:.1f} Hz") print(f" Range: {max_mean - min_mean:.1f} Hz") # Sample records if stats.get("sample_analysis"): print(f"\nFirst {len(stats['sample_analysis'])} records inspection:") print(f" {'ID':<20} {'Line':<8} {'Captions':<10} {'Pitch':<10} {'Algorithms'}") print(" " + "-" * 70) for sample in stats["sample_analysis"]: id_str = sample["id"][:20] if len(sample["id"]) > 20 else sample["id"] caption_str = "✓" if sample["has_captions"] else "✗" pitch_str = "✓" if sample["has_pitch"] else "✗" algo_str = ", ".join(sample["algorithms"]) if sample["algorithms"] else "None" print(f" {id_str:<20} {sample['line']:<8} {caption_str:<10} {pitch_str:<10} {algo_str}") print("\n" + "=" * 80) def main(): """Main execution function.""" parser = argparse.ArgumentParser( description="Validate vocal stem captions and pitch range data", formatter_class=argparse.RawDescriptionHelpFormatter, ) parser.add_argument( "--input-path", type=str, default="/home/vibert/data/voice_designer/metas_v6_pitch_latest.jsonl", help="Path to processed JSONL file with captions and pitch data", ) parser.add_argument( "--original-path", type=str, default="/app2/suno/data/auk_v0/metas_v6_tr.jsonl", help="Path to original JSONL file for line count comparison", ) parser.add_argument( "--sample-size", type=int, default=10000, help="Number of records to sample for detailed analysis", ) parser.add_argument("--output-dir", type=str, help="Directory for output plots and reports") args = parser.parse_args() # Run validation stats = validate_dataset( input_path=args.input_path, original_path=args.original_path, sample_size=args.sample_size, output_dir=args.output_dir, ) print(f"\nValidation complete. Report saved to output directory.") if __name__ == "__main__": main()