import json import os import tempfile import itertools from abc import ABC, abstractmethod from enum import Enum from typing import Dict, List, Optional, Tuple, Any from dataclasses import dataclass import boto3 import numpy as np import pandas as pd from tqdm import tqdm from pydantic import BaseModel from suno_utils.utils.s3 import download_s3_files from suno_utils.worker.modal_image_utils.eval_util import EvalTestType @dataclass class EvaluationConfig: """Global evaluation configuration.""" timestamp: str output_folder: str s3_bucket: str = "suno-data-uploads" s3_ditto_path: str = "s3://suno-data-uploads/tasks/feature_eval/cover_persona/" prompt_audio_timestamp: str = "2025_07_08-01_39_45" @dataclass class AuxiliaryMetric: """Definition of an auxiliary metric.""" name: str # e.g., "drift", "intro" description: str # Human-readable description file_suffix: str = None # Optional suffix for separate file saving include_in_summary: bool = True # Whether to generate summary stats @dataclass class EvaluationResult: """Standard result container.""" scores: Dict[str, Any] metadata: Dict[str, Any] = None class ScoreSummary(BaseModel): """Statistical summary.""" mean: float min: float max: float std: float count: int median: Optional[float] = None # Utilities def cosine_similarity(a: np.ndarray, b: np.ndarray) -> float: """Calculate cosine similarity.""" dot_product = np.dot(a, b) mag_a, mag_b = np.linalg.norm(a), np.linalg.norm(b) return 0.0 if mag_a == 0 or mag_b == 0 else dot_product / (mag_a * mag_b) def load_embedding(file_path: str) -> Optional[np.ndarray]: """Load embedding from npz file.""" if not os.path.isfile(file_path): return None try: data = np.load(file_path, allow_pickle=True) embedding = data.get("embedding") if embedding is not None and len(embedding.shape) > 0: return embedding return None except Exception as e: print(f"Error loading {file_path}: {e}") return None def calculate_stats(scores: List[float]) -> ScoreSummary: """Calculate summary statistics.""" if not scores: return ScoreSummary(mean=0.0, min=0.0, max=0.0, std=0.0, count=0, median=0.0) scores = [s for s in scores if not np.isnan(s)] return ScoreSummary( mean=float(np.mean(scores)), min=float(np.min(scores)), max=float(np.max(scores)), std=float(np.std(scores)), count=len(scores), median=float(np.median(scores)), ) def save_result( result: EvaluationResult, config: EvaluationConfig, extra_tags: str, test_type: str, eval_task: str, model_name: str, suffix: str = None, ) -> str: """Save result with standard filename format.""" # For auxiliary metrics (with suffix), skip eval_task to match baseline format # e.g., *lyrics_eval_feat_eval_cer* instead of *lyrics_eval_feat_eval_whisper_cer* if suffix: parts = [extra_tags, test_type, "feat_eval", suffix] else: parts = [extra_tags, test_type, "feat_eval", eval_task] parts.extend([model_name, config.timestamp]) filename = "_".join(parts) + ".npz" filepath = os.path.join(config.output_folder, filename) np.savez(filepath, **result.scores, metadata=result.metadata or {}) print(f"Saved to {filepath}") return filepath def clear_aws_env(): """Clear AWS environment variables.""" for env_var in ["AWS_PROFILE", "AWS_DEFAULT_PROFILE"]: os.environ.pop(env_var, None) boto3.DEFAULT_SESSION = None # Base Evaluator class BaseEvaluator(ABC): """Base evaluator class.""" def __init__(self, config: EvaluationConfig): self.config = config self.s3_client = boto3.client("s3") clear_aws_env() def get_auxiliary_metrics(self) -> List[AuxiliaryMetric]: """Return list of auxiliary metrics this evaluator provides. Example: return [ AuxiliaryMetric( name="quality", description="Audio quality scores", file_suffix="quality", include_in_summary=True ) ] """ return [] # Default: no auxiliary metrics @abstractmethod def evaluate(self, generations: List[Dict]) -> EvaluationResult: """Run evaluation on list of generations from standardized format.""" pass def get_task_name(self) -> str: """Get the task name for this evaluator (defaults to test_type).""" return self.__class__.__name__.lower().replace("evaluator", "") # Similarity-based Evaluators class SimilarityEvaluator(BaseEvaluator): """Base for embedding similarity evaluations.""" def __init__(self, config: EvaluationConfig, task: str): super().__init__(config) self.task = task # e.g., "self_sim", "artist_vox_sim" def evaluate(self, generations: List[Dict]) -> EvaluationResult: scores = {} # Group generations by source for similarity comparison grouped_generations = group_generations_by_source(generations) with tempfile.TemporaryDirectory() as tmp_dir: # Download files s3_paths, local_paths = self._get_file_paths(grouped_generations, tmp_dir) download_s3_files(s3_paths, local_paths) # Calculate similarities for source_id, source_generations in tqdm( grouped_generations.items(), desc=f"Processing {self.task}" ): source_data = source_generations[0].get("source_data", {}) source_file_id = ( source_data.get("source_id") or source_data.get("s3_id") or source_data.get("id") ) if not source_file_id: continue source_embedding = load_embedding(os.path.join(tmp_dir, f"{source_file_id}.npz")) if source_embedding is None: continue source_scores = [] for gen in source_generations: generation_id = gen["generation_id"] gen_embedding = load_embedding(os.path.join(tmp_dir, f"{generation_id}.npz")) if gen_embedding is not None: similarity = cosine_similarity(source_embedding, gen_embedding) source_scores.append(similarity) if source_scores: scores[source_id] = source_scores return EvaluationResult(scores) def _get_file_paths( self, grouped_generations: Dict[str, List[Dict]], tmp_dir: str ) -> Tuple[List[str], List[str]]: s3_paths, local_paths = [], [] for source_id, source_generations in grouped_generations.items(): # Source file (from first generation's source_data) source_data = source_generations[0].get("source_data", {}) source_file_id = ( source_data.get("source_id") or source_data.get("s3_id") or source_data.get("id") ) if source_file_id: s3_paths.append( f"{self.config.s3_ditto_path}{self.config.prompt_audio_timestamp}/{source_file_id}_{self.task}_ditto.npz" ) local_paths.append(os.path.join(tmp_dir, f"{source_file_id}.npz")) # Generation files for gen in source_generations: generation_id = gen["generation_id"] s3_paths.append( f"{self.config.s3_ditto_path}{self.config.timestamp}/{generation_id}_{self.task}_ditto.npz" ) local_paths.append(os.path.join(tmp_dir, f"{generation_id}.npz")) return s3_paths, local_paths class CoverEvaluator(SimilarityEvaluator): def __init__(self, config: EvaluationConfig): super().__init__(config, "self_sim") class ArtistEvaluator(SimilarityEvaluator): def __init__(self, config: EvaluationConfig, task: str = "self_sim"): super().__init__(config, task) # Specialized Evaluators class SubgenreEvaluator(BaseEvaluator): """Subgenre similarity evaluation.""" def evaluate(self, generations: List[Dict]) -> EvaluationResult: baseline_embeddings = np.load("/app2/suno/data/ditto_evals/baselines/subgenre_ditto_map.npz") with tempfile.TemporaryDirectory() as tmp_dir: # Download and process files s3_paths, local_paths, rows = [], [], [] for gen in generations: generation_id = gen["generation_id"] source_data = gen.get("source_data", {}) s3_paths.append( f"{self.config.s3_ditto_path}{self.config.timestamp}/{generation_id}_genre_sim_ditto.npz" ) local_paths.append(os.path.join(tmp_dir, f"{generation_id}.npz")) # Create row data from source_data row_data = source_data.copy() row_data["s3_id"] = generation_id rows.append(row_data) download_s3_files(s3_paths, local_paths) # Group by parent-child relationships parent_to_children = {} for row, path in zip(rows, local_paths): if not os.path.isfile(path): continue data = np.load(path, allow_pickle=True) current_genre = row.get("name", None) parent_genre = row.get("parent_genre", None) if parent_genre: parent_to_children.setdefault(parent_genre, {}).setdefault(current_genre, []) parent_to_children[parent_genre][current_genre].append(data["embedding"]) # Calculate similarity pairs result = {"within_genre_pairs": {}, "between_genre_pairs": {}} all_within_pairs = [] all_between_pairs = [] for parent_genre, child_genres in parent_to_children.items(): within_pairs, between_pairs = [], [] # Within-genre pairs for child_genre, embeddings in child_genres.items(): within_pairs.extend( cosine_similarity(e1, e2) for e1, e2 in itertools.combinations(embeddings, 2) ) # Between-genre pairs genre_items = list(child_genres.items()) for (g1, emb1), (g2, emb2) in itertools.combinations(genre_items, 2): between_pairs.extend( cosine_similarity(e1, e2) for e1, e2 in itertools.product(emb1, emb2) ) result["within_genre_pairs"][parent_genre] = within_pairs result["between_genre_pairs"][parent_genre] = between_pairs # Collect all pairs for summary all_within_pairs.extend(within_pairs) all_between_pairs.extend(between_pairs) # Return flattened scores for summary, keep raw results for file saving return EvaluationResult( scores={"within": all_within_pairs, "between": all_between_pairs}, metadata={"raw_results": result}, ) class GenreEvaluator(BaseEvaluator): """Genre diversity evaluation with auxiliary metrics (intro, duration).""" def get_auxiliary_metrics(self) -> List[AuxiliaryMetric]: return [ AuxiliaryMetric( name="intro", description="Introduction lengths for lyrical content", file_suffix="intro", include_in_summary=True, ), AuxiliaryMetric( name="duration", description="Audio duration from ditto embeddings", file_suffix="duration", include_in_summary=True, ), ] def evaluate(self, generations: List[Dict]) -> EvaluationResult: # Prepare data from generations item_rows = [] for gen in generations: generation_id = gen["generation_id"] source_data = gen.get("source_data", {}) # Create row data from source_data and generation info row_data = source_data.copy() row_data["s3_id"] = generation_id row_data["genre"] = source_data.get("group", "unknown") row_data["instrumental"] = len(source_data.get("lyrics", "")) < 15 item_rows.append(row_data) with tempfile.TemporaryDirectory() as tmp_dir: # Download embeddings s3_paths = [ f"{self.config.s3_ditto_path}{self.config.timestamp}/{row['s3_id']}_self_sim_ditto.npz" for row in item_rows ] local_paths = [os.path.join(tmp_dir, f"{row['s3_id']}.npz") for row in item_rows] download_s3_files(s3_paths, local_paths) # Process embeddings all_embeddings = {} for row, path in zip(item_rows, local_paths): try: data = np.load(path) row["ditto_embed"] = data["embedding"] row["duration"] = data.get("duration", None) # Extract duration from NPZ all_embeddings[row["s3_id"]] = data["chunk_embedding"] except Exception as e: print(f"Skipping {row['s3_id']}: {e}") row["ditto_embed"] = None row["duration"] = None # Calculate similarities df = pd.DataFrame([r for r in item_rows if r.get("ditto_embed") is not None]) # Check if we have any valid data to process if df.empty or "tags" not in df.columns: print(f"No valid embeddings found for processing. Skipping similarity calculations.") return EvaluationResult({}) df["s_tags"] = df["tags"].apply( lambda x: ", ".join(sorted(x.split(", "))) if isinstance(x, str) else "" ) similarity_pairs = [] for tag, group in df.groupby("s_tags"): pairs = list(itertools.combinations(group.itertuples(index=False), 2)) for row_a, row_b in pairs: similarity = cosine_similarity(row_a.ditto_embed, row_b.ditto_embed) similarity_pairs.append( { "tags": tag, "similarity": similarity, "instrumental": row_a.instrumental, "genre": row_a.genre, } ) # Calculate auxiliary metrics intro_scores = self._calculate_intro_lengths(df[~df["instrumental"]]["s3_id"].tolist()) duration_scores = self._calculate_duration(df) # Include all metrics in the scores dictionary scores = { "main": similarity_pairs, # Main diversity metric "intro": intro_scores, # Auxiliary metric 1 "duration": duration_scores, # Auxiliary metric 2 } return EvaluationResult(scores) def _calculate_intro_lengths(self, lyric_ids: List[str]) -> List[float]: """Calculate intro lengths from hoot files.""" s3_paths = [f"s3://suno-data-uploads/studio/uploads/{s3_id}_hoot.json" for s3_id in lyric_ids] with tempfile.TemporaryDirectory() as tmp_dir: local_paths = [os.path.join(tmp_dir, f"{s3_id}_hoot.json") for s3_id in lyric_ids] download_s3_files(s3_paths, local_paths) intro_lengths = [] for s3_id, path in zip(lyric_ids, local_paths): try: if os.path.isfile(path): data = pd.read_json(path) intro_length = data.iloc[0]["start_s"] if len(data) > 1 else np.nan if not np.isnan(intro_length): intro_lengths.append(intro_length) # Skip NaN values to maintain consistent list format except Exception: # Skip failed loads to maintain consistent list format pass return intro_lengths def _calculate_duration(self, df: pd.DataFrame) -> List[float]: """Extract duration values from already-loaded ditto embedding data.""" duration_scores = [] for _, row in df.iterrows(): duration = row.get("duration") # Handle numpy arrays (from NPZ files) and regular values if duration is not None: if hasattr(duration, "item"): # numpy scalar duration = duration.item() if not pd.isna(duration) and isinstance(duration, (int, float)): duration_scores.append(float(duration)) return duration_scores class WEREvaluator(BaseEvaluator): """Generic WER evaluation for lyrics and infill tasks.""" def __init__(self, config: EvaluationConfig, file_suffix: str, task_name: str): """ Args: config: Evaluation configuration file_suffix: S3 file suffix (e.g., "infill_wer", "lyrics_wer") task_name: Display name for progress bar (e.g., "infill", "lyrics") """ super().__init__(config) self.file_suffix = file_suffix self.task_name = task_name def get_auxiliary_metrics(self) -> List[AuxiliaryMetric]: return [ AuxiliaryMetric( name="cer", description="Character Error Rate", file_suffix="cer", include_in_summary=True, ), ] def evaluate(self, generations: List[Dict]) -> EvaluationResult: all_wer_data = [] grouped_results = {} grouped_cer_results = {} # Group generations by the 'group' field in source_data (e.g., "simple") grouped_generations = group_generations_by_field(generations, "group") for group_id, group_generations in tqdm( grouped_generations.items(), desc=f"Processing {self.task_name}" ): group_wers = [] group_wer_data = [] for gen in group_generations: generation_id = gen["generation_id"] wer_data = self._load_wer_data(generation_id) if wer_data and "wer" in wer_data: wer_data["s3_id"] = generation_id all_wer_data.append(wer_data) group_wer_data.append(wer_data) group_wers.append(wer_data["wer"]) if group_wers: # Store raw WER scores for plotting, keep statistics for metadata grouped_results[group_id] = { "wers": group_wers, # Raw WER scores for plotting "mean_wer": np.mean(group_wers), "count": len(group_wers), "min_wer": min(group_wers), "max_wer": max(group_wers), "std_wer": np.std(group_wers), } # Calculate CER for this group group_cers = self._calculate_cer(group_wer_data) if group_cers: grouped_cer_results[group_id] = group_cers # Overall results if all_wer_data: all_wers = [d["wer"] for d in all_wer_data] grouped_results["all"] = { "wers": all_wers, # Raw WER scores for plotting "mean_wer": np.mean(all_wers), "count": len(all_wers), "min_wer": min(all_wers), "max_wer": max(all_wers), "std_wer": np.std(all_wers), } # Calculate overall CER all_cers = self._calculate_cer(all_wer_data) if all_cers: grouped_cer_results["all"] = all_cers # Include all metrics in the scores dictionary scores = { "main": grouped_results, # Main WER metric "cer": grouped_cer_results, # Auxiliary metric with grouping } return EvaluationResult(scores) def _load_wer_data(self, file_id: str) -> Optional[Dict]: s3_key = ( f"tasks/feature_eval/cover_persona/{self.config.timestamp}/{file_id}_{self.file_suffix}.json" ) try: response = self.s3_client.get_object(Bucket=self.config.s3_bucket, Key=s3_key) return json.loads(response["Body"].read().decode("utf-8")) except Exception: return None def _calculate_cer(self, wer_data_list: List[Dict]) -> List[float]: """Extract CER values from loaded WER data files.""" cer_scores = [] for wer_data in wer_data_list: if "cer" in wer_data: cer_scores.append(wer_data["cer"]) return cer_scores class InfillEvaluator(WEREvaluator): """Infill WER evaluation.""" def __init__(self, config: EvaluationConfig): super().__init__(config, file_suffix="infill_wer", task_name="infill") class LyricsEvaluator(WEREvaluator): """Lyrics WER evaluation with auxiliary CER metrics.""" def __init__(self, config: EvaluationConfig): super().__init__(config, file_suffix="lyrics_wer", task_name="lyrics") def get_auxiliary_metrics(self) -> List[AuxiliaryMetric]: # Get CER from parent class base_metrics = super().get_auxiliary_metrics() # Add hoot_cer for lyrics base_metrics.append( AuxiliaryMetric( name="hoot_cer", description="Character Error Rate from hoot alignment", file_suffix="hoot_cer", include_in_summary=True, ) ) return base_metrics def evaluate(self, generations: List[Dict]) -> EvaluationResult: # Call parent evaluate to get WER results and CER (both grouped) result = super().evaluate(generations) # Calculate hoot_cer per group to match WER and CER grouping grouped_hoot_cer_results = {} grouped_generations = group_generations_by_field(generations, "group") for group_id, group_generations in grouped_generations.items(): # Extract generation IDs for this group lyric_ids = [gen["generation_id"] for gen in group_generations] # Calculate hoot_cer for this group group_hoot_cer = self._calculate_hoot_cer(lyric_ids) if group_hoot_cer: grouped_hoot_cer_results[group_id] = group_hoot_cer # Calculate overall hoot_cer all_lyric_ids = [gen["generation_id"] for gen in generations] all_hoot_cer = self._calculate_hoot_cer(all_lyric_ids) if all_hoot_cer: grouped_hoot_cer_results["all"] = all_hoot_cer # Add to scores if grouped_hoot_cer_results: result.scores["hoot_cer"] = grouped_hoot_cer_results return result def _calculate_hoot_cer(self, lyric_ids: List[str]) -> List[float]: """Calculate Character Error Rate from hoot alignment files.""" s3_paths = [f"s3://suno-data-uploads/studio/uploads/{s3_id}_hoot.json" for s3_id in lyric_ids] with tempfile.TemporaryDirectory() as tmp_dir: local_paths = [os.path.join(tmp_dir, f"{s3_id}_hoot.json") for s3_id in lyric_ids] download_s3_files(s3_paths, local_paths) cer_scores = [] for s3_id, path in zip(lyric_ids, local_paths): try: if os.path.isfile(path): # Load JSON as raw data first with open(path, "r") as f: raw_data = json.load(f) # Search through the list of dictionaries for hoot_cer for item in raw_data: if isinstance(item, dict) and "hoot_cer" in item: cer_value = item["hoot_cer"] if cer_value is not None and isinstance(cer_value, (int, float)): cer_scores.append(float(cer_value)) break # Found it, move to next file # Skip NaN values to maintain consistent list format except Exception as e: # Skip failed loads to maintain consistent list format print(f"Warning: Failed to extract hoot_cer for {s3_id}: {e}") pass return cer_scores class InstrumentalDriftEvaluator(BaseEvaluator): """Evaluates instrumental drift by comparing embeddings from early and late slices of source audio.""" def __init__(self, config: EvaluationConfig): super().__init__(config) self.task = "artist_sim" # Uses artist_sim embeddings from DittoInstrumentalBatchProcessor def evaluate(self, generations: List[Dict]) -> EvaluationResult: """Compare early vs late instrumental embeddings for each generated track. For each generation, embeddings are saved with filenames: {gen_id}_{slice_type}_artist_sim_ditto.npz We read slice_configs from source_data to know what slices to load. """ with tempfile.TemporaryDirectory() as tmp_dir: s3_paths = [] local_paths = [] slice_info = [] # Track which generation and slice type each file belongs to for gen in generations: gen_id = gen["generation_id"] source_data = gen.get("source_data", {}) slice_configs = source_data.get("slice_configs", []) group = source_data.get("group", "unknown") # For each slice configuration, construct the embedding file path for slice_config in slice_configs: slice_type = slice_config["type"] filename_id = f"{gen_id}_{slice_type}" s3_path = f"{self.config.s3_ditto_path}{self.config.timestamp}/{filename_id}_{self.task}_ditto.npz" local_path = os.path.join(tmp_dir, f"{filename_id}.npz") s3_paths.append(s3_path) local_paths.append(local_path) slice_info.append( { "generation_id": gen_id, "slice_type": slice_type, "group": group, } ) # Download all embeddings download_s3_files(s3_paths, local_paths) # Group slices by generation_id source_groups = {} for path, info in zip(local_paths, slice_info): try: data = np.load(path) embedding = data["embedding"] gen_id = info["generation_id"] slice_type = info["slice_type"] group = info["group"] if gen_id not in source_groups: source_groups[gen_id] = {"early": None, "late": None, "group": group} source_groups[gen_id][slice_type] = embedding source_groups[gen_id]["group"] = group except Exception as e: print(f"Failed to load {path}: {e}") continue # Calculate drift scores drift_scores_by_group = {} all_drift_scores = [] for gen_id, slices in tqdm(source_groups.items(), desc="Calculating instrumental drift"): if slices["early"] is None or slices["late"] is None: print(f"Warning: Missing slices for generation {gen_id}") continue # Calculate similarity (higher = less drift) similarity = cosine_similarity(slices["early"], slices["late"]) all_drift_scores.append(similarity) # Group by category if available group = slices.get("group", "all") if group not in drift_scores_by_group: drift_scores_by_group[group] = [] drift_scores_by_group[group].append(similarity) # Prepare results - include both grouped and overall scores scores = {} if drift_scores_by_group: scores = drift_scores_by_group # Always include overall "all" category if all_drift_scores: scores["all"] = all_drift_scores return EvaluationResult( scores, metadata={ "description": "Instrumental drift: cosine similarity between early and late instrumental embeddings. Higher = less drift." }, ) # Factory class EvaluatorFactory: """Factory for creating evaluators.""" _REGISTRY = { EvalTestType.COVER: CoverEvaluator, EvalTestType.ARTIST: ArtistEvaluator, # Handles both self_sim and artist_vox_sim EvalTestType.GENRE_EVAL: SubgenreEvaluator, EvalTestType.BASE_GENERATION: GenreEvaluator, EvalTestType.INFILL_EVAL: InfillEvaluator, EvalTestType.LYRICS_EVAL: LyricsEvaluator, EvalTestType.INST_DRIFT: InstrumentalDriftEvaluator, } @classmethod def create(cls, eval_type, config: EvaluationConfig, eval_task: str = None) -> BaseEvaluator: """Create evaluator instance.""" if eval_type not in cls._REGISTRY: raise ValueError(f"Unknown evaluation type: {eval_type}") evaluator_class = cls._REGISTRY[eval_type] if evaluator_class is None: raise ValueError(f"No evaluator implemented for: {eval_type}") # Handle special cases where we need to pass the task if evaluator_class == ArtistEvaluator and eval_task: return evaluator_class(config, eval_task) else: return evaluator_class(config) @classmethod def register(cls, eval_type, evaluator_class: type): """Register new evaluator type.""" cls._REGISTRY[eval_type] = evaluator_class # Processing Pipeline def run_evaluation( eval_type, mapping_data: Dict, metadata: Dict, config: EvaluationConfig, summary: Dict ): """Run single evaluation and save results.""" eval_task = metadata.get("eval_task") evaluator = EvaluatorFactory.create(eval_type, config, eval_task) # Extract generations from the new standardized format generations = extract_generations_data(mapping_data) result = evaluator.evaluate(generations) # Extract metadata fields with fallbacks extra_tags = metadata.get("extra_tags", "unknown") eval_type_value = eval_type.value if hasattr(eval_type, "value") else str(eval_type).lower() test_type = metadata.get("test_type", eval_type_value) eval_task = metadata.get("eval_task", eval_type_value) model_name = metadata.get("model_name", "unknown") eval_type_name = eval_type.value if hasattr(eval_type, "value") else str(eval_type) print(f"Processing {eval_type_name}: {extra_tags}, {test_type}, {eval_task}, {model_name}") # Handle special cases if eval_type == EvalTestType.GENRE_EVAL: # Save within/between files separately using raw results raw_results = result.metadata.get("raw_results", {}) within_result = EvaluationResult( {"within_genre_pairs": raw_results.get("within_genre_pairs", {})} ) between_result = EvaluationResult( {"between_genre_pairs": raw_results.get("between_genre_pairs", {})} ) save_result(within_result, config, extra_tags, test_type, eval_task, model_name, "within_parent") save_result(between_result, config, extra_tags, test_type, eval_task, model_name, "with_parent") # Add to summary using flattened scores if result.scores: for subtask, scores_list in result.scores.items(): if isinstance(scores_list, list) and scores_list: score_summary = calculate_stats(scores_list).model_dump() summary_key = f"{eval_type.value}_{subtask}_{extra_tags}_{model_name}" summary[summary_key] = score_summary elif eval_type == EvalTestType.BASE_GENERATION: # Handle auxiliary metrics first auxiliary_metrics = evaluator.get_auxiliary_metrics() for aux_metric in auxiliary_metrics: if aux_metric.name in result.scores and result.scores[aux_metric.name]: # Save auxiliary metric to separate file aux_result = EvaluationResult({aux_metric.name: result.scores[aux_metric.name]}) suffix = aux_metric.file_suffix or aux_metric.name save_result(aux_result, config, extra_tags, "genre", eval_task, model_name, suffix) # Add to summary if requested if aux_metric.include_in_summary: # Handle grouped auxiliary metrics (dict) vs flat list aux_data = result.scores[aux_metric.name] if isinstance(aux_data, dict) and "all" in aux_data: # Grouped format: use the "all" group for overall summary aux_summary = calculate_stats(aux_data["all"]).model_dump() elif isinstance(aux_data, list): # Flat list format aux_summary = calculate_stats(aux_data).model_dump() else: # Skip if format is unexpected continue summary[f"{aux_metric.name}_{extra_tags}_{model_name}"] = aux_summary # Remove auxiliary metrics from main result before saving main_result = EvaluationResult( scores={ k: v for k, v in result.scores.items() if k not in [m.name for m in auxiliary_metrics] }, metadata=result.metadata, ) save_result(main_result, config, extra_tags, test_type, eval_task, model_name) # Add diversity summary (main metric) if "main" in result.scores: df = pd.DataFrame(result.scores["main"]) diversity_summary = {} for group_name, group_data in [ ("instrumental", df[df["instrumental"]]), ("lyrical", df[~df["instrumental"]]), ("all", df), ]: if len(group_data) > 0: diversity_summary[group_name] = calculate_stats( group_data["similarity"].tolist() ).model_dump() else: diversity_summary[group_name] = calculate_stats([]).model_dump() summary[f"diversity_{extra_tags}_{model_name}"] = diversity_summary else: # Standard evaluations # Handle auxiliary metrics first auxiliary_metrics = evaluator.get_auxiliary_metrics() for aux_metric in auxiliary_metrics: if aux_metric.name in result.scores and result.scores[aux_metric.name]: # Save auxiliary metric to separate file aux_result = EvaluationResult({aux_metric.name: result.scores[aux_metric.name]}) suffix = aux_metric.file_suffix or aux_metric.name save_result(aux_result, config, extra_tags, test_type, eval_task, model_name, suffix) # Add to summary if requested if aux_metric.include_in_summary: # Handle grouped auxiliary metrics (dict) vs flat list aux_data = result.scores[aux_metric.name] if isinstance(aux_data, dict) and "all" in aux_data: # Grouped format: use the "all" group for overall summary aux_summary = calculate_stats(aux_data["all"]).model_dump() elif isinstance(aux_data, list): # Flat list format aux_summary = calculate_stats(aux_data).model_dump() else: # Skip if format is unexpected continue summary[f"{aux_metric.name}_{extra_tags}_{model_name}"] = aux_summary # Remove auxiliary metrics from main result before saving main_result = EvaluationResult( scores={ k: v for k, v in result.scores.items() if k not in [m.name for m in auxiliary_metrics] }, metadata=result.metadata, ) save_result(main_result, config, extra_tags, test_type, eval_task, model_name) # Add to summary if eval_type == EvalTestType.INFILL_EVAL: if "main" in result.scores and "all" in result.scores["main"]: stats = result.scores["main"]["all"] summary[f"infill_{extra_tags}_{model_name}"] = { "mean": stats["mean_wer"], "count": stats["count"], } else: print(f"Warning: No 'all' stats found in infill results for {extra_tags}") summary[f"infill_{extra_tags}_{model_name}"] = {"mean": 0.0, "count": 0} elif eval_type == EvalTestType.LYRICS_EVAL: if "main" in result.scores and "all" in result.scores["main"]: stats = result.scores["main"]["all"] summary[f"lyrics_{extra_tags}_{model_name}"] = { "mean": stats["mean_wer"], "count": stats["count"], } else: print(f"Warning: No 'all' stats found in lyrics results for {extra_tags}") summary[f"lyrics_{extra_tags}_{model_name}"] = {"mean": 0.0, "count": 0} else: # Similarity evaluations - flatten scores if result.scores: all_scores = [score for scores in result.scores.values() for score in scores] score_summary = calculate_stats(all_scores).model_dump() summary[f"{eval_type.value}_{extra_tags}_{model_name}"] = score_summary else: print(f"Warning: No scores found for {eval_type.value} {extra_tags}") summary[f"{eval_type.value}_{extra_tags}_{model_name}"] = calculate_stats( [] ).model_dump() def extract_generations_data(mapping_data: Dict) -> List[Dict]: """Extract generations list from standardized mapping format.""" if "generations" not in mapping_data: raise ValueError("Expected new standardized mapping format with 'generations' key") return mapping_data["generations"] def group_generations_by_source(generations: List[Dict]) -> Dict[str, List[Dict]]: """Group generations by their source identifier for similarity evaluations.""" grouped = {} for gen in generations: source_data = gen.get("source_data", {}) # Extract source identifier - try different possible keys source_id = source_data.get("source_id") or source_data.get("s3_id") or source_data.get("id") if not source_id: # Fallback: use a default group source_id = "all_generations" if source_id not in grouped: grouped[source_id] = [] grouped[source_id].append(gen) return grouped def group_generations_by_field(generations: List[Dict], field_name: str) -> Dict[str, List[Dict]]: """Group generations by a specific field in source_data.""" grouped = {} for gen in generations: source_data = gen.get("source_data", {}) group_key = source_data.get(field_name) if not group_key: # Fallback: use a default group group_key = "unknown" if group_key not in grouped: grouped[group_key] = [] grouped[group_key].append(gen) return grouped def get_eval_type_from_filename(filename: str): """Extract evaluation type from mapping filename.""" # Parse the filename to extract test_type and eval_task # New format: {name}_{test_type}_{eval_task}_{timestamp}.json parts = filename.replace(".json", "").split("_") if len(parts) < 3: return None # Extract test_type (second to last part before timestamp) test_type = parts[-3] # Map test_type to EvalTestType for eval_type in EvalTestType: if eval_type.value == test_type: return eval_type return None def main( file: Optional[str] = None, timestamp: str = "2025_06_20-11_17_20", data_folder: str = "/app2/suno/data/ditto_evals", output_folder: str = "/app2/suno/data/ditto_evals", ): """Main evaluation function.""" config = EvaluationConfig(timestamp=timestamp, output_folder=os.path.join(output_folder, timestamp)) os.makedirs(config.output_folder, exist_ok=True) if file: # Single file processing eval_type = get_eval_type_from_filename(os.path.basename(file)) if eval_type is None: raise ValueError(f"Cannot identify evaluation type from: {file}") eval_type_value = eval_type.value if hasattr(eval_type, "value") else str(eval_type).lower() name = os.path.basename(file).split(eval_type_value)[0].rstrip("_") metadata = { "extra_tags": name, "test_type": eval_type_value, "eval_task": eval_type_value, "model_name": "unknown", } with open(file) as f: mapping_data = json.load(f) summary = {} run_evaluation(eval_type, mapping_data, metadata, config, summary) with open( os.path.join(config.output_folder, f"summary_{eval_type.value}_{name}_{timestamp}.json"), "w" ) as f: json.dump(summary, f, indent=4) else: # Batch processing data_dir = os.path.join(data_folder, timestamp) with open(os.path.join(data_dir, "metadata.json")) as f: all_metadata = json.load(f) all_metadata.pop("timestamp", None) all_metadata.pop("checkpoint", None) all_summaries = {} print(all_metadata) for name, metadata in all_metadata.items(): print(name, metadata) if not isinstance(metadata, dict) or "model_name" not in metadata: continue print(f"Processing: {name}") print(f"Metadata: {metadata}") # Debug metadata summary = {} # Map evaluation types to the combinations we need to look for eval_mappings = { EvalTestType.COVER: [("cover", "self_sim")], EvalTestType.ARTIST: [ ("artist", "self_sim"), ("artist", "artist_vox_sim"), ], # Includes VOX EvalTestType.GENRE_EVAL: [("subgenre", "genre_sim")], EvalTestType.BASE_GENERATION: [("base_gen", "self_sim")], EvalTestType.INFILL_EVAL: [("infill", "whisper")], EvalTestType.LYRICS_EVAL: [("lyrics_eval", "whisper")], # EvalTestType.DURATION: Removed - now handled as auxiliary metric in GenreEvaluator EvalTestType.INST_DRIFT: [("inst_drift", "artist_sim")], } # Find and process all mapping files for this experiment for eval_type, type_task_pairs in eval_mappings.items(): for test_type_val, eval_task_val in type_task_pairs: mapping_file = os.path.join( data_dir, f"{name}_{test_type_val}_{eval_task_val}_{timestamp}.json" ) if os.path.exists(mapping_file): print(f"Found mapping file: {mapping_file}") with open(mapping_file) as f: mapping_data = json.load(f) run_evaluation(eval_type, mapping_data, metadata, config, summary) if summary: all_summaries[name] = summary # Print summary table summary_table = {} for name, evaluations in all_summaries.items(): for eval_name, summary in evaluations.items(): if "diversity" in eval_name: for group, group_summary in summary.items(): key = f"diversity_{group}" summary_table[key] = [group_summary.get("mean"), group_summary.get("count", 0)] else: # Clean key extraction - handle multi-task evaluations parts = eval_name.split("_") if len(parts) >= 2 and parts[0] == "subgenre": # For subgenre, keep both eval_type and subtask key = f"{parts[0]}_{parts[1]}" # e.g., "subgenre_within", "subgenre_between" else: # For other evaluations, just use eval type key = parts[0] # Just take the first part (evaluation type) summary_table[key] = [summary["mean"], summary["count"]] print("\nSummary Table:") if summary_table: print(pd.DataFrame(summary_table, index=["mean", "count"]).round(2).T) with open(os.path.join(config.output_folder, f"score_summary_{timestamp}.json"), "w") as f: json.dump(all_summaries, f, indent=4) return all_summaries if not file else summary if __name__ == "__main__": import argparse parser = argparse.ArgumentParser(description="Process music evaluation mappings") parser.add_argument("--file", type=str, help="Single mapping file") parser.add_argument("--timestamp", type=str, default="2025_06_20-11_17_20", help="Timestamp") parser.add_argument( "--data_folder", type=str, default="/app2/suno/data/ditto_evals", help="Input folder" ) parser.add_argument( "--output_folder", type=str, default="/app2/suno/data/ditto_evals", help="Output folder" ) args = parser.parse_args() main(args.file, args.timestamp, args.data_folder, args.output_folder)