#!/usr/bin/env python3 """ Extract samples for specific keywords from voice descriptions. Allows targeting specific keywords or phrases for sample extraction. """ import argparse import json import logging import os import shutil import subprocess from datetime import datetime from pathlib import Path from typing import Dict, List, Optional from tqdm import tqdm def setup_logging(output_dir: Path, log_name: str = "keyword_extraction") -> logging.Logger: """Set up logging configuration.""" log_file = output_dir / f"{log_name}_{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 SpecificKeywordExtractor: """Extract audio samples for specific keywords.""" def __init__(self, dataset_file: str, output_dir: str, samples_per_keyword: int = 5): """ Initialize the extractor. Args: dataset_file: Path to captioned dataset JSONL output_dir: Base output directory samples_per_keyword: Number of records to find per keyword """ self.dataset_file = Path(dataset_file) self.output_dir = Path(output_dir) self.output_dir.mkdir(parents=True, exist_ok=True) self.samples_per_keyword = samples_per_keyword self.logger = setup_logging(self.output_dir, "specific_keywords") # Track statistics self.stats = { "keywords_processed": 0, "total_records_found": 0, "total_audio_converted": 0, "keywords_with_samples": [], "keywords_without_samples": [], "errors": [], } def find_records_with_keyword(self, keyword: str, max_records: int = None) -> List[Dict]: """ Find records containing a specific keyword (partial match allowed). Args: keyword: Keyword to search for (will match if keyword is contained in caption keywords) max_records: Maximum records to return Returns: List of matching records """ if max_records is None: max_records = self.samples_per_keyword found_records = [] keyword_lower = keyword.lower() self.logger.info(f"Searching for keyword: '{keyword}'") records_scanned = 0 with open(self.dataset_file, "r") as f: for line_num, line in enumerate(f, 1): if len(found_records) >= max_records: break if not line.strip(): continue records_scanned += 1 # Progress update if records_scanned % 50000 == 0: self.logger.debug( f" Scanned {records_scanned:,} records, found {len(found_records)} matches..." ) try: record = json.loads(line.strip()) # Check stems_captions for keyword stems_captions = record.get("stems_captions", {}) for stem_name, caption_list in stems_captions.items(): if len(found_records) >= max_records: break for caption_entry in caption_list: if caption_entry.get("prompt_type") == "voice_description_keywords": caption = caption_entry.get("caption", "").lower() # Check if keyword appears in caption (partial match) # Split by comma and check each caption keyword caption_keywords = [k.strip() for k in caption.split(",")] # Check for partial match for cap_keyword in caption_keywords: if keyword_lower in cap_keyword: # Add metadata for processing record["_matched_stem"] = stem_name record["_matched_keyword"] = keyword record["_matched_caption_keyword"] = cap_keyword found_records.append(record) self.logger.debug( f" Found match: '{cap_keyword}' in record {record.get('id')}" ) break if len(found_records) > len(found_records) - 1: # Found a match break except json.JSONDecodeError: continue except Exception as e: self.logger.warning(f"Error processing line {line_num}: {e}") self.logger.info( f" Found {len(found_records)} records for keyword '{keyword}' (scanned {records_scanned:,} records)" ) if found_records: self.stats["keywords_with_samples"].append(keyword) else: self.stats["keywords_without_samples"].append(keyword) return found_records def convert_audio_to_mp3( self, input_path: str, output_path: str, skip_if_exists: bool = True ) -> bool: """ Convert audio file to MP3. Args: input_path: Input audio file path output_path: Output MP3 file path skip_if_exists: Skip conversion if output exists Returns: True if successful, False otherwise """ output_path = Path(output_path) if skip_if_exists and output_path.exists(): self.logger.debug(f" MP3 already exists: {output_path.name}") return True if not os.path.exists(input_path): self.logger.warning(f" Input file not found: {input_path}") return False try: cmd = [ "ffmpeg", "-i", input_path, "-acodec", "libmp3lame", "-b:a", "192k", str(output_path), "-y", "-loglevel", "error", ] subprocess.run(cmd, check=True, capture_output=True) self.stats["total_audio_converted"] += 1 return True except subprocess.CalledProcessError as e: self.logger.error(f" Failed to convert {input_path}: {e}") self.stats["errors"].append(f"Audio conversion failed: {input_path}") return False def process_keyword(self, keyword: str) -> Dict: """ Process a single keyword to extract samples. Args: keyword: Keyword to process Returns: Dictionary with processing results """ self.logger.info(f"\n{'='*60}") self.logger.info(f"Processing keyword: '{keyword}'") self.logger.info(f"{'='*60}") # Create keyword directory keyword_safe = keyword.replace(" ", "_").replace("/", "_").replace("-", "_") keyword_dir = self.output_dir / keyword_safe keyword_dir.mkdir(exist_ok=True) # Find records records = self.find_records_with_keyword(keyword) keyword_results = { "keyword": keyword, "records_found": len(records), "samples": [], "audio_files_created": 0, } if not records: self.logger.warning(f"No records found for keyword: '{keyword}'") return keyword_results # Process each record for idx, record in enumerate(records, 1): record_id = record.get("id", f"unknown_{idx}") matched_stem = record.get("_matched_stem", "Vocals") matched_caption = record.get("_matched_caption_keyword", keyword) self.logger.info(f" [{idx}/{len(records)}] Processing: {record_id}") self.logger.info(f" Matched: '{matched_caption}' in {matched_stem}") # Create record directory record_dir = keyword_dir / f"{idx:02d}_{record_id}" record_dir.mkdir(exist_ok=True) sample_info = { "id": record_id, "matched_stem": matched_stem, "matched_caption": matched_caption, "audio_files": [], } # Save metadata metadata_file = record_dir / "metadata.json" with open(metadata_file, "w") as f: json.dump( { "id": record_id, "keyword_searched": keyword, "keyword_matched": matched_caption, "matched_stem": matched_stem, "text": record.get("text", ""), "duration_s": record.get("duration_s", 0), "tags": record.get("tags", [])[:20], # Limit tags to save space "stems_captions": record.get("stems_captions", {}), }, f, indent=2, ) # Convert full song local_filepath = record.get("local_filepath") if local_filepath: full_mp3 = record_dir / f"full_{record_id}.mp3" if self.convert_audio_to_mp3(local_filepath, full_mp3): sample_info["audio_files"].append("full_song") keyword_results["audio_files_created"] += 1 self.logger.info(f" ✓ Converted full song") # Convert matched stem stems = record.get("stems", {}) if matched_stem in stems: stem_path = stems[matched_stem] stem_mp3 = record_dir / f"{matched_stem.lower().replace(' ', '_')}_{record_id}.mp3" if self.convert_audio_to_mp3(stem_path, stem_mp3): sample_info["audio_files"].append(matched_stem) keyword_results["audio_files_created"] += 1 self.logger.info(f" ✓ Converted {matched_stem} stem") # Also get main Vocals if different if matched_stem != "Vocals" and "Vocals" in stems: vocals_path = stems["Vocals"] vocals_mp3 = record_dir / f"vocals_{record_id}.mp3" if self.convert_audio_to_mp3(vocals_path, vocals_mp3): sample_info["audio_files"].append("Vocals") keyword_results["audio_files_created"] += 1 self.logger.info(f" ✓ Converted Vocals stem") keyword_results["samples"].append(sample_info) self.stats["keywords_processed"] += 1 self.stats["total_records_found"] += len(records) return keyword_results def process_keywords(self, keywords: List[str]) -> List[Dict]: """ Process multiple keywords. Args: keywords: List of keywords to process Returns: List of results for each keyword """ self.logger.info(f"\nProcessing {len(keywords)} keywords") self.logger.info(f"Target samples per keyword: {self.samples_per_keyword}") results = [] for keyword in tqdm(keywords, desc="Processing keywords"): keyword_results = self.process_keyword(keyword) results.append(keyword_results) return results def generate_summary(self, results: List[Dict]) -> Dict: """Generate extraction summary.""" summary = { "timestamp": datetime.now().isoformat(), "configuration": { "dataset_file": str(self.dataset_file), "output_dir": str(self.output_dir), "samples_per_keyword": self.samples_per_keyword, }, "statistics": { "keywords_processed": self.stats["keywords_processed"], "keywords_with_samples": len(self.stats["keywords_with_samples"]), "keywords_without_samples": len(self.stats["keywords_without_samples"]), "total_records_found": self.stats["total_records_found"], "total_audio_converted": self.stats["total_audio_converted"], "errors": len(self.stats["errors"]), }, "keywords_with_samples": self.stats["keywords_with_samples"], "keywords_without_samples": self.stats["keywords_without_samples"], "detailed_results": results, } return summary def run(self, keywords: List[str]): """Run extraction for specified keywords.""" self.logger.info("=" * 80) self.logger.info("SPECIFIC KEYWORD EXTRACTION") self.logger.info("=" * 80) # Process keywords results = self.process_keywords(keywords) # Generate summary summary = self.generate_summary(results) # Save summary summary_file = self.output_dir / "extraction_summary.json" with open(summary_file, "w") as f: json.dump(summary, f, indent=2) # Create README self.create_readme(summary) # Final report self.logger.info("\n" + "=" * 80) self.logger.info("EXTRACTION COMPLETE") self.logger.info("=" * 80) self.logger.info(f"Keywords processed: {self.stats['keywords_processed']}") self.logger.info(f"Keywords with samples: {len(self.stats['keywords_with_samples'])}") self.logger.info(f"Keywords without samples: {len(self.stats['keywords_without_samples'])}") self.logger.info(f"Total records found: {self.stats['total_records_found']}") self.logger.info(f"Audio files converted: {self.stats['total_audio_converted']}") self.logger.info(f"\nOutput directory: {self.output_dir}") self.logger.info(f"Summary saved to: {summary_file}") if self.stats["keywords_without_samples"]: self.logger.warning( f"\nNo samples found for: {', '.join(self.stats['keywords_without_samples'])}" ) def create_readme(self, summary: Dict): """Create README file.""" readme_file = self.output_dir / "README.md" with open(readme_file, "w") as f: f.write("# Specific Keyword Extraction Results\n\n") f.write(f"Generated: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n\n") f.write("## Summary\n\n") stats = summary["statistics"] f.write(f"- Keywords processed: {stats['keywords_processed']}\n") f.write(f"- Keywords with samples: {stats['keywords_with_samples']}\n") f.write(f"- Total records found: {stats['total_records_found']}\n") f.write(f"- Audio files created: {stats['total_audio_converted']}\n\n") f.write("## Keywords with Samples\n\n") for result in summary["detailed_results"]: if result["records_found"] > 0: f.write(f"### {result['keyword']}\n") f.write(f"- Records found: {result['records_found']}\n") f.write(f"- Audio files: {result['audio_files_created']}\n") f.write(f"- Directory: `{result['keyword'].replace(' ', '_')}/`\n\n") if summary["keywords_without_samples"]: f.write("## Keywords without Samples\n\n") for keyword in summary["keywords_without_samples"]: f.write(f"- {keyword}\n") f.write("\n") f.write("## Directory Structure\n\n") f.write("```\n") f.write("[keyword]/\n") f.write(" ├── 01_[record_id]/\n") f.write(" │ ├── metadata.json\n") f.write(" │ ├── full_*.mp3\n") f.write(" │ └── [stem]_*.mp3\n") f.write(" └── ...\n") f.write("```\n") self.logger.info(f"README created: {readme_file}") def main(): """Main execution.""" parser = argparse.ArgumentParser(description="Extract audio samples for specific keywords") parser.add_argument( "--dataset-file", default="/home/vibert/data/voice_designer/metas_v6_tr_vocal_stems_captioned_w30.jsonl", help="Path to captioned dataset JSONL", ) parser.add_argument( "--output-dir", default="/home/vibert/tmp/specific_keyword_samples", help="Output directory for samples", ) parser.add_argument( "--keywords", nargs="+", required=True, help="Keywords to search for (space-separated)" ) parser.add_argument( "--samples-per-keyword", type=int, default=5, help="Number of samples per keyword (default: 5)" ) args = parser.parse_args() # Create extractor extractor = SpecificKeywordExtractor( dataset_file=args.dataset_file, output_dir=args.output_dir, samples_per_keyword=args.samples_per_keyword, ) # Run extraction extractor.run(args.keywords) if __name__ == "__main__": main()