#!/usr/bin/env python3 """ Analyze and count all keywords from voice_description_keywords in stems_captions. This script processes a JSONL file containing captioned vocal stems and extracts all keywords from the voice_description_keywords prompt type, counting their frequencies and generating comprehensive statistics. """ import argparse import json import logging from collections import Counter, defaultdict from pathlib import Path from typing import Dict, List, Set from datetime import datetime from tqdm import tqdm def setup_logging(output_dir: Path) -> logging.Logger: """Set up logging configuration.""" log_file = output_dir / f"keyword_analysis_{datetime.now().strftime('%Y%m%d_%H%M%S')}.log" logging.basicConfig( level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s", handlers=[logging.FileHandler(log_file), logging.StreamHandler()], ) return logging.getLogger(__name__) class VoiceKeywordAnalyzer: """Analyzer for voice description keywords in captioned stems.""" def __init__(self, input_jsonl: str, output_dir: str): """Initialize the analyzer. Args: input_jsonl: Path to input JSONL file output_dir: Directory for output files """ self.input_jsonl = Path(input_jsonl) self.output_dir = Path(output_dir) self.output_dir.mkdir(parents=True, exist_ok=True) self.logger = setup_logging(self.output_dir) # Statistics tracking self.keyword_counter = Counter() self.stem_type_keywords = defaultdict(Counter) self.records_with_keywords = 0 self.total_records = 0 self.total_captions = 0 self.empty_captions = 0 self.error_captions = 0 def extract_keywords_from_caption(self, caption: str) -> List[str]: """Extract individual keywords from a caption string. Args: caption: Comma-separated keyword string Returns: List of individual keywords (normalized) """ if not caption: return [] # Split by comma and clean each keyword keywords = [] for keyword in caption.split(","): cleaned = keyword.strip().lower() # Remove trailing periods that sometimes appear if cleaned.endswith("."): cleaned = cleaned[:-1].strip() if cleaned: keywords.append(cleaned) return keywords def process_record(self, record: Dict) -> Dict[str, List[str]]: """Process a single record to extract keywords. Args: record: Record dictionary Returns: Dictionary mapping stem names to keyword lists """ stem_keywords = {} stems_captions = record.get("stems_captions", {}) for stem_name, caption_list in stems_captions.items(): keywords = [] for caption_entry in caption_list: if caption_entry.get("prompt_type") == "voice_description_keywords": caption = caption_entry.get("caption", "") error = caption_entry.get("error") if error: self.error_captions += 1 self.logger.debug(f"Caption error for {record.get('id')}: {error}") elif caption: extracted = self.extract_keywords_from_caption(caption) keywords.extend(extracted) self.total_captions += 1 else: self.empty_captions += 1 if keywords: stem_keywords[stem_name] = keywords return stem_keywords def analyze_file(self) -> None: """Analyze the entire JSONL file for keywords.""" self.logger.info(f"Starting analysis of {self.input_jsonl}") # Count total lines first for progress bar with open(self.input_jsonl, "r") as f: total_lines = sum(1 for line in f if line.strip()) self.logger.info(f"Total records to process: {total_lines:,}") # Process records with open(self.input_jsonl, "r") as f: for line_num, line in enumerate(tqdm(f, total=total_lines, desc="Analyzing keywords"), 1): if not line.strip(): continue self.total_records += 1 try: record = json.loads(line.strip()) stem_keywords = self.process_record(record) if stem_keywords: self.records_with_keywords += 1 # Update counters for stem_name, keywords in stem_keywords.items(): for keyword in keywords: self.keyword_counter[keyword] += 1 self.stem_type_keywords[stem_name][keyword] += 1 # Log progress every 10,000 records if line_num % 10000 == 0: self.logger.info( f"Processed {line_num:,} records - " f"Found {len(self.keyword_counter):,} unique keywords" ) except json.JSONDecodeError as e: self.logger.warning(f"JSON decode error at line {line_num}: {e}") except Exception as e: self.logger.error(f"Error processing line {line_num}: {e}") self.logger.info("Analysis complete!") self.logger.info(f"Total records processed: {self.total_records:,}") self.logger.info(f"Records with keywords: {self.records_with_keywords:,}") self.logger.info(f"Unique keywords found: {len(self.keyword_counter):,}") self.logger.info(f"Total captions processed: {self.total_captions:,}") self.logger.info(f"Empty captions: {self.empty_captions:,}") self.logger.info(f"Error captions: {self.error_captions:,}") def generate_reports(self) -> None: """Generate output reports in various formats.""" # 1. Complete keyword frequency (all keywords) self.logger.info("Generating complete keyword frequency report...") all_keywords_file = self.output_dir / "all_keywords_frequency.json" all_keywords_data = { "metadata": { "total_records": self.total_records, "records_with_keywords": self.records_with_keywords, "total_unique_keywords": len(self.keyword_counter), "total_keyword_occurrences": sum(self.keyword_counter.values()), "total_captions": self.total_captions, "empty_captions": self.empty_captions, "error_captions": self.error_captions, "input_file": str(self.input_jsonl), "analysis_timestamp": datetime.now().isoformat(), }, "keyword_frequencies": dict(self.keyword_counter.most_common()), } with open(all_keywords_file, "w") as f: json.dump(all_keywords_data, f, indent=2) self.logger.info(f"Saved complete keyword frequencies to {all_keywords_file}") # 2. Top keywords report (top 1000) self.logger.info("Generating top keywords report...") top_keywords_file = self.output_dir / "top_1000_keywords.json" top_keywords_data = { "metadata": all_keywords_data["metadata"], "top_1000_keywords": dict(self.keyword_counter.most_common(1000)), } with open(top_keywords_file, "w") as f: json.dump(top_keywords_data, f, indent=2) self.logger.info(f"Saved top 1000 keywords to {top_keywords_file}") # 3. Keywords by stem type self.logger.info("Generating stem-specific keyword report...") stem_keywords_file = self.output_dir / "keywords_by_stem_type.json" stem_keywords_data = { "metadata": { "total_stem_types": len(self.stem_type_keywords), "stem_types": list(self.stem_type_keywords.keys()), }, "stem_keywords": {}, } for stem_name, keyword_counter in self.stem_type_keywords.items(): stem_keywords_data["stem_keywords"][stem_name] = { "total_unique_keywords": len(keyword_counter), "total_occurrences": sum(keyword_counter.values()), "top_100_keywords": dict(keyword_counter.most_common(100)), } with open(stem_keywords_file, "w") as f: json.dump(stem_keywords_data, f, indent=2) self.logger.info(f"Saved stem-specific keywords to {stem_keywords_file}") # 4. Keyword categories (grouped by type) self.logger.info("Generating keyword categories report...") categories_file = self.output_dir / "keyword_categories.json" categories = self.categorize_keywords() with open(categories_file, "w") as f: json.dump(categories, f, indent=2) self.logger.info(f"Saved keyword categories to {categories_file}") # 5. Summary statistics self.logger.info("Generating summary statistics...") summary_file = self.output_dir / "analysis_summary.json" # Calculate percentiles keyword_counts = list(self.keyword_counter.values()) keyword_counts.sort() summary_data = { "total_statistics": { "total_records_analyzed": self.total_records, "records_with_voice_keywords": self.records_with_keywords, "percentage_with_keywords": round( 100 * self.records_with_keywords / self.total_records, 2 ), "total_unique_keywords": len(self.keyword_counter), "total_keyword_occurrences": sum(self.keyword_counter.values()), "average_keywords_per_record": round( sum(self.keyword_counter.values()) / self.records_with_keywords, 2 ) if self.records_with_keywords else 0, }, "keyword_distribution": { "most_common_keyword": self.keyword_counter.most_common(1)[0] if self.keyword_counter else None, "keywords_appearing_once": sum( 1 for count in self.keyword_counter.values() if count == 1 ), "keywords_appearing_10+_times": sum( 1 for count in self.keyword_counter.values() if count >= 10 ), "keywords_appearing_100+_times": sum( 1 for count in self.keyword_counter.values() if count >= 100 ), "keywords_appearing_1000+_times": sum( 1 for count in self.keyword_counter.values() if count >= 1000 ), }, "top_20_keywords": dict(self.keyword_counter.most_common(20)), "stem_type_summary": { stem: { "unique_keywords": len(counter), "total_occurrences": sum(counter.values()), "top_5": dict(counter.most_common(5)), } for stem, counter in self.stem_type_keywords.items() }, } with open(summary_file, "w") as f: json.dump(summary_data, f, indent=2) self.logger.info(f"Saved analysis summary to {summary_file}") def categorize_keywords(self) -> Dict: """Categorize keywords into semantic groups.""" categories = { "gender": {"keywords": [], "total_count": 0}, "age": {"keywords": [], "total_count": 0}, "accent": {"keywords": [], "total_count": 0}, "emotion": {"keywords": [], "total_count": 0}, "vocal_quality": {"keywords": [], "total_count": 0}, "genre": {"keywords": [], "total_count": 0}, "effects": {"keywords": [], "total_count": 0}, "other": {"keywords": [], "total_count": 0}, } # Define category patterns gender_keywords = ["male", "female", "masculine", "feminine", "androgynous"] age_keywords = ["young", "adult", "child", "teen", "elderly", "mature", "youthful"] accent_keywords = [ "accent", "american", "british", "german", "french", "spanish", "italian", "chinese", "japanese", "korean", "indian", "russian", "scottish", "irish", "australian", ] emotion_keywords = [ "happy", "sad", "angry", "emotional", "passionate", "melancholic", "joyful", "soulful", "expressive", "intense", "aggressive", ] quality_keywords = [ "clear", "raspy", "smooth", "rough", "deep", "high", "soft", "loud", "bright", "dark", "warm", "breathy", "powerful", "gentle", "resonant", ] genre_keywords = [ "pop", "rock", "jazz", "blues", "soul", "r&b", "hip-hop", "country", "folk", "metal", "electronic", "classical", "opera", ] effect_keywords = [ "reverb", "echo", "distortion", "auto-tune", "vocoder", "delay", "chorus", "compressed", "filtered", ] for keyword, count in self.keyword_counter.items(): categorized = False # Check each category if any(g in keyword for g in gender_keywords): categories["gender"]["keywords"].append((keyword, count)) categories["gender"]["total_count"] += count categorized = True elif any(a in keyword for a in age_keywords): categories["age"]["keywords"].append((keyword, count)) categories["age"]["total_count"] += count categorized = True elif any(acc in keyword for acc in accent_keywords): categories["accent"]["keywords"].append((keyword, count)) categories["accent"]["total_count"] += count categorized = True elif any(e in keyword for e in emotion_keywords): categories["emotion"]["keywords"].append((keyword, count)) categories["emotion"]["total_count"] += count categorized = True elif any(q in keyword for q in quality_keywords): categories["vocal_quality"]["keywords"].append((keyword, count)) categories["vocal_quality"]["total_count"] += count categorized = True elif any(g in keyword for g in genre_keywords): categories["genre"]["keywords"].append((keyword, count)) categories["genre"]["total_count"] += count categorized = True elif any(fx in keyword for fx in effect_keywords): categories["effects"]["keywords"].append((keyword, count)) categories["effects"]["total_count"] += count categorized = True if not categorized: categories["other"]["keywords"].append((keyword, count)) categories["other"]["total_count"] += count # Sort keywords within each category and limit to top 50 for category in categories.values(): category["keywords"] = sorted(category["keywords"], key=lambda x: x[1], reverse=True)[:50] category["keywords"] = dict(category["keywords"]) return categories def run(self) -> None: """Run the complete analysis pipeline.""" self.logger.info("=" * 60) self.logger.info("Voice Keyword Analysis Starting") self.logger.info("=" * 60) # Analyze the file self.analyze_file() # Generate reports self.generate_reports() self.logger.info("=" * 60) self.logger.info("Analysis Complete!") self.logger.info(f"Output files saved to: {self.output_dir}") self.logger.info("=" * 60) def main(): """Main execution function.""" parser = argparse.ArgumentParser( description="Analyze voice description keywords from captioned stems" ) parser.add_argument( "--input-jsonl", default="/home/vibert/data/voice_designer/metas_v6_tr_vocal_stems_captioned_w30.jsonl", help="Path to input JSONL file with captioned stems", ) parser.add_argument( "--output-dir", default="/home/vibert/tmp/voice_keywords_analysis", help="Output directory for analysis results", ) args = parser.parse_args() # Run analyzer analyzer = VoiceKeywordAnalyzer(args.input_jsonl, args.output_dir) analyzer.run() if __name__ == "__main__": main()