"""Collector for interesting sample records from the dataset.""" import json import random from pathlib import Path from typing import List, Dict, Any, Optional from collections import defaultdict import logging from .base_analyzer import BaseAnalyzer logger = logging.getLogger(__name__) # ============================================================================ # SAMPLING TASK TOGGLES - Enable/disable specific sampling tasks # ============================================================================ ENABLE_TAG_LENGTH_SAMPLES = False # Max/longest tag samples ENABLE_TAG_COUNT_SAMPLES = False # Most tags sample ENABLE_TEXT_LENGTH_SAMPLES = False # Shortest/longest text samples ENABLE_LANGUAGE_SAMPLES = False # One sample per language ENABLE_STEM_SAMPLES = False # One sample per stem type ENABLE_BRACKET_SAMPLES = False # Samples with different bracket types ENABLE_WEIGHT_SAMPLES = False # Min/max weight samples ENABLE_DURATION_SAMPLES = False # Min/max duration samples ENABLE_SPECIAL_SAMPLES = False # Multi-artist, multi-cover, etc. ENABLE_DIRECTORY_SAMPLES = False # Samples from different directories ENABLE_S3_PATTERN_SAMPLES = False # Samples from different S3 patterns # S3 category sampling with tag length filter (MAIN FOCUS) ENABLE_S3_CATEGORY_SAMPLING = True # Reservoir sampling by s3_category S3_CATEGORY_MIN_TAG_LENGTH = 0 # Sample all records (0 = no tag length filter) # ============================================================================ class SampleCollector(BaseAnalyzer): """Collect interesting sample records for analysis.""" def __init__(self, output_dir, samples_per_s3_category=10, random_seed=42): super().__init__("samples", output_dir) self.samples_dir = self.output_dir / "samples" self.samples_dir.mkdir(exist_ok=True) self.samples_per_s3_category = samples_per_s3_category self.random_seed = random_seed random.seed(random_seed) # Reservoir sampling state for s3_categories self.s3_category_reservoirs = defaultdict(list) self.s3_category_counts = defaultdict(int) def process_chunk(self, records: List[Dict[str, Any]]) -> Dict[str, Any]: """Process a chunk to collect interesting samples. Args: records: List of JSON records Returns: Dictionary with collected samples """ samples = { "max_total_tag_length": {"length": 0, "record": None}, "longest_single_tag": {"length": 0, "tag": "", "record": None}, "most_tags": {"count": 0, "record": None}, "shortest_text": {"length": float("inf"), "record": None}, "longest_text": {"length": 0, "record": None}, "by_language": {}, "by_stem_type": {}, "by_bracket_type": {}, "by_directory": {}, "by_s3_pattern": {}, "extreme_weights": { "min_weight": {"value": float("inf"), "record": None}, "max_weight": {"value": float("-inf"), "record": None}, }, "extreme_durations": { "min_duration": {"value": float("inf"), "record": None}, "max_duration": {"value": float("-inf"), "record": None}, }, "multi_artist": None, "multi_cover": None, "has_vocal_pitch_range": None, "has_text_aligned": None, "all_stems_types": None, # Record with the most diverse stem types } bracket_patterns = { "square": r"\[", "curly": r"\{", "angle": r"<", "parentheses": r"\(", "chinese_angle": "【", "chinese_book": "《", "japanese_corner": "「", "japanese_double": "『", } for record in records: # Total tag length tags = record.get("tags", []) if ENABLE_TAG_LENGTH_SAMPLES and tags: total_length = sum(len(tag) for tag in tags) if total_length > samples["max_total_tag_length"]["length"]: samples["max_total_tag_length"] = {"length": total_length, "record": record} # Longest single tag for tag in tags: if len(tag) > samples["longest_single_tag"]["length"]: samples["longest_single_tag"] = { "length": len(tag), "tag": tag, "record": record, } # Most tags if ENABLE_TAG_COUNT_SAMPLES and tags: if len(tags) > samples["most_tags"]["count"]: samples["most_tags"] = {"count": len(tags), "record": record} # Text length if ENABLE_TEXT_LENGTH_SAMPLES: text = record.get("text", "") if text: text_length = len(text) if text_length < samples["shortest_text"]["length"]: samples["shortest_text"] = {"length": text_length, "record": record} if text_length > samples["longest_text"]["length"]: samples["longest_text"] = {"length": text_length, "record": record} # Bracket types if ENABLE_BRACKET_SAMPLES: text = record.get("text", "") if text: for bracket_type, pattern in bracket_patterns.items(): if pattern in text and bracket_type not in samples["by_bracket_type"]: samples["by_bracket_type"][bracket_type] = record # Language samples if ENABLE_LANGUAGE_SAMPLES: lang = record.get("lang") if lang and lang not in samples["by_language"]: samples["by_language"][lang] = record # Stem samples if ENABLE_STEM_SAMPLES: stems = record.get("stems", {}) for stem_name in stems.keys(): if stem_name not in samples["by_stem_type"]: samples["by_stem_type"][stem_name] = record # All stem types - find record with most diverse stems if stems and len(stems) > 0: current_all_stems = samples["all_stems_types"] if current_all_stems is None or len(stems) > len(current_all_stems.get("stems", {})): samples["all_stems_types"] = record # Directory samples if ENABLE_DIRECTORY_SAMPLES: local_path = record.get("local_filepath", "") if local_path: dir_path = "/".join(local_path.split("/")[:-1]) if dir_path and dir_path not in samples["by_directory"]: samples["by_directory"][dir_path] = record # S3 pattern samples s3_path = record.get("s3_filepath", "") if s3_path and s3_path.startswith("s3://"): parts = s3_path[5:].split("/") # Original S3 pattern sampling if ENABLE_S3_PATTERN_SAMPLES and len(parts) > 3: pattern = "/".join(parts[:4]) if pattern not in samples["by_s3_pattern"]: samples["by_s3_pattern"][pattern] = record # Reservoir sampling for s3_categories with tag length filter if ENABLE_S3_CATEGORY_SAMPLING and len(parts) >= 4: category = "/".join(parts[1:4]) # Check if record has a tag with length > S3_CATEGORY_MIN_TAG_LENGTH tags = record.get("tags", []) has_long_tag = any(len(tag) > S3_CATEGORY_MIN_TAG_LENGTH for tag in tags) if has_long_tag: self.s3_category_counts[category] += 1 n = self.s3_category_counts[category] if len(self.s3_category_reservoirs[category]) < self.samples_per_s3_category: self.s3_category_reservoirs[category].append(record) else: # Replace with probability k/n replace_idx = random.randint(0, n - 1) if replace_idx < self.samples_per_s3_category: self.s3_category_reservoirs[category][replace_idx] = record # Weight extremes if ENABLE_WEIGHT_SAMPLES: weight = record.get("weight") if weight is not None: if weight < samples["extreme_weights"]["min_weight"]["value"]: samples["extreme_weights"]["min_weight"] = {"value": weight, "record": record} if weight > samples["extreme_weights"]["max_weight"]["value"]: samples["extreme_weights"]["max_weight"] = {"value": weight, "record": record} # Duration extremes if ENABLE_DURATION_SAMPLES: duration = record.get("duration_s") if duration is not None: if duration < samples["extreme_durations"]["min_duration"]["value"]: samples["extreme_durations"]["min_duration"] = { "value": duration, "record": record, } if duration > samples["extreme_durations"]["max_duration"]["value"]: samples["extreme_durations"]["max_duration"] = { "value": duration, "record": record, } # Multi-artist/cover samples and special fields if ENABLE_SPECIAL_SAMPLES: artist_ids = record.get("artist_ids", []) if artist_ids and len(artist_ids) > 1 and samples["multi_artist"] is None: samples["multi_artist"] = record cover_ids = record.get("cover_ids", []) if cover_ids and len(cover_ids) > 1 and samples["multi_cover"] is None: samples["multi_cover"] = record # Special fields if "vocal_pitch_range" in record and samples["has_vocal_pitch_range"] is None: samples["has_vocal_pitch_range"] = record if "text_aligned" in record and samples["has_text_aligned"] is None: samples["has_text_aligned"] = record # Add reservoir samples to the return samples["s3_category_reservoirs"] = dict(self.s3_category_reservoirs) samples["s3_category_counts"] = dict(self.s3_category_counts) return samples def aggregate(self, chunk_results: List[Dict[str, Any]]) -> Dict[str, Any]: """Aggregate samples from all chunks. Args: chunk_results: List of sample results from each chunk Returns: Final collection of interesting samples """ final_samples = { "max_total_tag_length": {"length": 0, "record": None}, "longest_single_tag": {"length": 0, "tag": "", "record": None}, "most_tags": {"count": 0, "record": None}, "shortest_text": {"length": float("inf"), "record": None}, "longest_text": {"length": 0, "record": None}, "by_language": {}, "by_stem_type": {}, "by_bracket_type": {}, "by_directory": {}, "by_s3_pattern": {}, "extreme_weights": { "min_weight": {"value": float("inf"), "record": None}, "max_weight": {"value": float("-inf"), "record": None}, }, "extreme_durations": { "min_duration": {"value": float("inf"), "record": None}, "max_duration": {"value": float("-inf"), "record": None}, }, "multi_artist": None, "multi_cover": None, "has_vocal_pitch_range": None, "has_text_aligned": None, "all_stems_types": None, } for chunk in chunk_results: # Max total tag length if chunk["max_total_tag_length"]["length"] > final_samples["max_total_tag_length"]["length"]: final_samples["max_total_tag_length"] = chunk["max_total_tag_length"] # Longest single tag if chunk["longest_single_tag"]["length"] > final_samples["longest_single_tag"]["length"]: final_samples["longest_single_tag"] = chunk["longest_single_tag"] # Most tags if chunk["most_tags"]["count"] > final_samples["most_tags"]["count"]: final_samples["most_tags"] = chunk["most_tags"] # Text extremes if ( chunk["shortest_text"]["record"] and chunk["shortest_text"]["length"] < final_samples["shortest_text"]["length"] ): final_samples["shortest_text"] = chunk["shortest_text"] if chunk["longest_text"]["length"] > final_samples["longest_text"]["length"]: final_samples["longest_text"] = chunk["longest_text"] # Merge language samples (keep first found) for lang, record in chunk["by_language"].items(): if lang not in final_samples["by_language"]: final_samples["by_language"][lang] = record # Merge stem samples (limit to top 50) for stem, record in chunk["by_stem_type"].items(): if len(final_samples["by_stem_type"]) < 50 and stem not in final_samples["by_stem_type"]: final_samples["by_stem_type"][stem] = record # Merge bracket samples for bracket, record in chunk["by_bracket_type"].items(): if bracket not in final_samples["by_bracket_type"]: final_samples["by_bracket_type"][bracket] = record # Merge directory samples (limit to 10) for dir_path, record in chunk["by_directory"].items(): if ( len(final_samples["by_directory"]) < 10 and dir_path not in final_samples["by_directory"] ): final_samples["by_directory"][dir_path] = record # Merge S3 patterns (limit to 10) for pattern, record in chunk["by_s3_pattern"].items(): if ( len(final_samples["by_s3_pattern"]) < 10 and pattern not in final_samples["by_s3_pattern"] ): final_samples["by_s3_pattern"][pattern] = record # Weight extremes if chunk["extreme_weights"]["min_weight"]["record"]: if ( chunk["extreme_weights"]["min_weight"]["value"] < final_samples["extreme_weights"]["min_weight"]["value"] ): final_samples["extreme_weights"]["min_weight"] = chunk["extreme_weights"][ "min_weight" ] if chunk["extreme_weights"]["max_weight"]["record"]: if ( chunk["extreme_weights"]["max_weight"]["value"] > final_samples["extreme_weights"]["max_weight"]["value"] ): final_samples["extreme_weights"]["max_weight"] = chunk["extreme_weights"][ "max_weight" ] # Duration extremes if chunk["extreme_durations"]["min_duration"]["record"]: if ( chunk["extreme_durations"]["min_duration"]["value"] < final_samples["extreme_durations"]["min_duration"]["value"] ): final_samples["extreme_durations"]["min_duration"] = chunk["extreme_durations"][ "min_duration" ] if chunk["extreme_durations"]["max_duration"]["record"]: if ( chunk["extreme_durations"]["max_duration"]["value"] > final_samples["extreme_durations"]["max_duration"]["value"] ): final_samples["extreme_durations"]["max_duration"] = chunk["extreme_durations"][ "max_duration" ] # Other special samples if chunk["multi_artist"] and not final_samples["multi_artist"]: final_samples["multi_artist"] = chunk["multi_artist"] if chunk["multi_cover"] and not final_samples["multi_cover"]: final_samples["multi_cover"] = chunk["multi_cover"] if chunk["has_vocal_pitch_range"] and not final_samples["has_vocal_pitch_range"]: final_samples["has_vocal_pitch_range"] = chunk["has_vocal_pitch_range"] if chunk["has_text_aligned"] and not final_samples["has_text_aligned"]: final_samples["has_text_aligned"] = chunk["has_text_aligned"] if chunk["all_stems_types"]: current = final_samples["all_stems_types"] if current is None or len(chunk["all_stems_types"].get("stems", {})) > len( current.get("stems", {}) ): final_samples["all_stems_types"] = chunk["all_stems_types"] # Merge s3_category reservoir samples from all chunks if "s3_category_reservoirs" in chunk: for category, samples_list in chunk["s3_category_reservoirs"].items(): if category not in self.s3_category_reservoirs: self.s3_category_reservoirs[category] = [] self.s3_category_reservoirs[category].extend(samples_list) # Merge s3_category counts if "s3_category_counts" in chunk: for category, count in chunk["s3_category_counts"].items(): self.s3_category_counts[category] += count # Trim reservoirs to desired size (in case we got more from merging chunks) for category in self.s3_category_reservoirs: if len(self.s3_category_reservoirs[category]) > self.samples_per_s3_category: self.s3_category_reservoirs[category] = random.sample( self.s3_category_reservoirs[category], self.samples_per_s3_category ) # Save individual sample files self._save_samples(final_samples) # Return summary return { "samples_collected": { "max_total_tag_length": final_samples["max_total_tag_length"]["length"], "longest_single_tag": final_samples["longest_single_tag"]["length"], "most_tags_count": final_samples["most_tags"]["count"], "shortest_text_length": final_samples["shortest_text"]["length"] if final_samples["shortest_text"]["record"] else None, "longest_text_length": final_samples["longest_text"]["length"], "languages_sampled": len(final_samples["by_language"]), "stem_types_sampled": len(final_samples["by_stem_type"]), "bracket_types_sampled": len(final_samples["by_bracket_type"]), "directories_sampled": len(final_samples["by_directory"]), "s3_patterns_sampled": len(final_samples["by_s3_pattern"]), "min_weight": final_samples["extreme_weights"]["min_weight"]["value"] if final_samples["extreme_weights"]["min_weight"]["record"] else None, "max_weight": final_samples["extreme_weights"]["max_weight"]["value"] if final_samples["extreme_weights"]["max_weight"]["record"] else None, "min_duration": final_samples["extreme_durations"]["min_duration"]["value"] if final_samples["extreme_durations"]["min_duration"]["record"] else None, "max_duration": final_samples["extreme_durations"]["max_duration"]["value"] if final_samples["extreme_durations"]["max_duration"]["record"] else None, } } def _save_samples(self, samples: Dict[str, Any]) -> None: """Save individual sample files with proper formatting. Args: samples: Dictionary of collected samples """ # Save tag length extreme samples if ENABLE_TAG_LENGTH_SAMPLES: if samples["max_total_tag_length"]["record"]: self._save_sample("max_total_tag_length.json", samples["max_total_tag_length"]["record"]) if samples["longest_single_tag"]["record"]: self._save_sample("longest_single_tag.json", samples["longest_single_tag"]["record"]) # Save tag count samples if ENABLE_TAG_COUNT_SAMPLES and samples["most_tags"]["record"]: self._save_sample("most_tags.json", samples["most_tags"]["record"]) # Save text length samples if ENABLE_TEXT_LENGTH_SAMPLES: if samples["shortest_text"]["record"]: self._save_sample("shortest_text.json", samples["shortest_text"]["record"]) if samples["longest_text"]["record"]: self._save_sample("longest_text.json", samples["longest_text"]["record"]) # Save language samples if ENABLE_LANGUAGE_SAMPLES: for lang, record in samples["by_language"].items(): safe_lang = lang.replace("/", "_").replace(" ", "_") self._save_sample(f"language_{safe_lang}.json", record) # Save stem samples if ENABLE_STEM_SAMPLES: for stem, record in samples["by_stem_type"].items(): safe_stem = stem.replace("/", "_").replace(" ", "_") self._save_sample(f"stem_{safe_stem}.json", record) # Save bracket samples if ENABLE_BRACKET_SAMPLES: for bracket, record in samples["by_bracket_type"].items(): self._save_sample(f"bracket_{bracket}.json", record) # Save weight/duration extremes if ENABLE_WEIGHT_SAMPLES: if samples["extreme_weights"]["min_weight"]["record"]: self._save_sample("min_weight.json", samples["extreme_weights"]["min_weight"]["record"]) if samples["extreme_weights"]["max_weight"]["record"]: self._save_sample("max_weight.json", samples["extreme_weights"]["max_weight"]["record"]) if ENABLE_DURATION_SAMPLES: if samples["extreme_durations"]["min_duration"]["record"]: self._save_sample( "min_duration.json", samples["extreme_durations"]["min_duration"]["record"] ) if samples["extreme_durations"]["max_duration"]["record"]: self._save_sample( "max_duration.json", samples["extreme_durations"]["max_duration"]["record"] ) # Save special samples if ENABLE_SPECIAL_SAMPLES: if samples["multi_artist"]: self._save_sample("multi_artist.json", samples["multi_artist"]) if samples["multi_cover"]: self._save_sample("multi_cover.json", samples["multi_cover"]) if samples["has_vocal_pitch_range"]: self._save_sample("has_vocal_pitch_range.json", samples["has_vocal_pitch_range"]) if samples["has_text_aligned"]: self._save_sample("has_text_aligned.json", samples["has_text_aligned"]) if samples["all_stems_types"]: self._save_sample("all_stems_types.json", samples["all_stems_types"]) # Save directory samples if ENABLE_DIRECTORY_SAMPLES: for dir_path, record in samples["by_directory"].items(): safe_dir = dir_path.replace("/", "_")[-50:] # Limit length self._save_sample(f"dir_{safe_dir}.json", record) # Save S3 pattern samples if ENABLE_S3_PATTERN_SAMPLES: for pattern, record in samples["by_s3_pattern"].items(): safe_pattern = pattern.replace("/", "_") self._save_sample(f"s3pattern_{safe_pattern}.json", record) # Save s3_category reservoir samples for category, category_samples in self.s3_category_reservoirs.items(): safe_category = category.replace("/", "_") category_dir = self.samples_dir / f"s3_category_{safe_category}" category_dir.mkdir(exist_ok=True) for idx, record in enumerate(category_samples, 1): filename = category_dir / f"sample_{idx:02d}.json" with open(filename, "w") as f: json.dump(record, f, indent=2, ensure_ascii=False) logger.info(f"Saved {len(category_samples)} samples for category: {category}") def _save_sample(self, filename: str, record: Dict[str, Any]) -> None: """Save a single sample record to file. Args: filename: Name of the file record: Record to save """ filepath = self.samples_dir / filename with open(filepath, "w") as f: json.dump(record, f, indent=2, ensure_ascii=False) logger.debug(f"Saved sample: {filename}")