#!/usr/bin/env python3 """ Extract samples for multiple keywords with AND logic from voice descriptions. Records must contain ALL specified keywords to match. """ 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, Set 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 AndKeywordExtractor: """Extract audio samples that match ALL specified keywords.""" def __init__(self, dataset_file: str, output_dir: str, samples_per_query: int = 5): """ Initialize the extractor. Args: dataset_file: Path to captioned dataset JSONL output_dir: Base output directory samples_per_query: Number of records to find """ self.dataset_file = Path(dataset_file) self.output_dir = Path(output_dir) self.output_dir.mkdir(parents=True, exist_ok=True) self.samples_per_query = samples_per_query self.logger = setup_logging(self.output_dir, "and_keywords") # Track statistics self.stats = { "keywords_processed": [], "total_records_found": 0, "total_audio_converted": 0, "query_successful": False, "errors": [], } def find_records_with_all_keywords(self, keywords: List[str], max_records: int = None) -> List[Dict]: """ Find records containing ALL specified keywords. Args: keywords: List of keywords that must ALL be present max_records: Maximum records to return Returns: List of matching records """ if max_records is None: max_records = self.samples_per_query found_records = [] keywords_lower = [k.lower() for k in keywords] self.logger.info(f"Searching for records with ALL keywords: {keywords}") 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 ALL keywords 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() # Split by comma to get individual caption keywords caption_keywords = [k.strip() for k in caption.split(",")] # Check if ALL required keywords are found (partial match allowed) matched_keywords = {} for required_keyword in keywords_lower: for cap_keyword in caption_keywords: if required_keyword in cap_keyword: matched_keywords[required_keyword] = cap_keyword break # If all keywords were matched if len(matched_keywords) == len(keywords_lower): # Add metadata for processing record["_matched_stem"] = stem_name record["_matched_keywords"] = keywords record["_matched_caption_keywords"] = matched_keywords record["_all_caption_keywords"] = caption_keywords found_records.append(record) self.logger.debug( f" Found match in record {record.get('id')}: {matched_keywords}" ) 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 with ALL keywords (scanned {records_scanned:,} records)" ) if found_records: self.stats["query_successful"] = True 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_keywords(self, keywords: List[str]) -> Dict: """ Process keywords with AND logic to extract samples. Args: keywords: List of keywords that must ALL be present Returns: Dictionary with processing results """ self.logger.info(f"\n{'='*60}") self.logger.info(f"Processing AND query: {' AND '.join(keywords)}") self.logger.info(f"{'='*60}") # Create output directory name from keywords keywords_safe = [k.replace(" ", "_").replace("/", "_").replace("-", "_") for k in keywords] keyword_dir_name = "_AND_".join(keywords_safe) keyword_dir = self.output_dir / keyword_dir_name keyword_dir.mkdir(exist_ok=True) # Find records records = self.find_records_with_all_keywords(keywords) query_results = { "keywords": keywords, "query_type": "AND", "records_found": len(records), "samples": [], "audio_files_created": 0, } if not records: self.logger.warning(f"No records found with ALL keywords: {keywords}") return query_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_keywords = record.get("_matched_caption_keywords", {}) self.logger.info(f" [{idx}/{len(records)}] Processing: {record_id}") self.logger.info(f" Matched keywords: {matched_keywords}") self.logger.info(f" In stem: {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_keywords": matched_keywords, "audio_files": [], } # Save metadata metadata_file = record_dir / "metadata.json" with open(metadata_file, "w") as f: json.dump( { "id": record_id, "keywords_searched": keywords, "keywords_matched": matched_keywords, "matched_stem": matched_stem, "all_caption_keywords": record.get("_all_caption_keywords", []), "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") query_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) query_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") query_results["audio_files_created"] += 1 self.logger.info(f" ✓ Converted Vocals stem") query_results["samples"].append(sample_info) self.stats["keywords_processed"] = keywords self.stats["total_records_found"] = len(records) return query_results def generate_summary(self, results: 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_query": self.samples_per_query, }, "query": { "keywords": results["keywords"], "query_type": "AND", "query_string": " AND ".join(results["keywords"]), }, "statistics": { "keywords_searched": self.stats["keywords_processed"], "query_successful": self.stats["query_successful"], "total_records_found": self.stats["total_records_found"], "total_audio_converted": self.stats["total_audio_converted"], "errors": len(self.stats["errors"]), }, "detailed_results": results, } return summary def run(self, keywords: List[str]): """Run extraction for specified keywords with AND logic.""" self.logger.info("=" * 80) self.logger.info("AND KEYWORD EXTRACTION") self.logger.info("=" * 80) # Process keywords results = self.process_keywords(keywords) # Generate summary summary = self.generate_summary(results) # Save summary keywords_safe = [k.replace(" ", "_").replace("/", "_").replace("-", "_") for k in keywords] keyword_dir_name = "_AND_".join(keywords_safe) keyword_dir = self.output_dir / keyword_dir_name summary_file = keyword_dir / "extraction_summary.json" with open(summary_file, "w") as f: json.dump(summary, f, indent=2) # Create README self.create_readme(summary, keyword_dir) # Final report self.logger.info("\n" + "=" * 80) self.logger.info("EXTRACTION COMPLETE") self.logger.info("=" * 80) self.logger.info(f"Keywords searched: {' AND '.join(keywords)}") self.logger.info(f"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: {keyword_dir}") self.logger.info(f"Summary saved to: {summary_file}") if self.stats["total_records_found"] == 0: self.logger.warning(f"\nNo samples found with ALL keywords: {keywords}") def create_readme(self, summary: Dict, output_dir: Path): """Create README file.""" readme_file = output_dir / "README.md" with open(readme_file, "w") as f: f.write("# AND Keyword Extraction Results\n\n") f.write(f"Generated: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n\n") f.write("## Query\n\n") f.write(f"Keywords: **{' AND '.join(summary['query']['keywords'])}**\n") f.write("(Records must contain ALL keywords)\n\n") f.write("## Summary\n\n") stats = summary["statistics"] f.write(f"- Records found: {stats['total_records_found']}\n") f.write(f"- Audio files created: {stats['total_audio_converted']}\n\n") if summary["detailed_results"]["records_found"] > 0: f.write("## Samples\n\n") for idx, sample in enumerate(summary["detailed_results"]["samples"], 1): f.write(f"### Sample {idx}: {sample['id']}\n") f.write(f"- Matched stem: {sample['matched_stem']}\n") f.write(f"- Matched keywords: {sample['matched_keywords']}\n") f.write(f"- Audio files: {', '.join(sample['audio_files'])}\n\n") f.write("## Directory Structure\n\n") f.write("```\n") f.write("[keyword_AND_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(" ├── extraction_summary.json\n") f.write(" └── README.md\n") f.write("```\n") self.logger.info(f"README created: {readme_file}") def main(): """Main execution.""" parser = argparse.ArgumentParser( description="Extract audio samples that contain ALL specified keywords (AND logic)" ) 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 (ALL must be present)" ) parser.add_argument( "--samples", type=int, default=5, help="Number of samples to extract (default: 5)" ) args = parser.parse_args() # Create extractor extractor = AndKeywordExtractor( dataset_file=args.dataset_file, output_dir=args.output_dir, samples_per_query=args.samples ) # Run extraction extractor.run(args.keywords) if __name__ == "__main__": main()