""" Task 1: Count - Generate counting questions This task joins multiple audio sources and asks questions about counting the number of unique sound sources in the audio. """ import csv import random from pathlib import Path from typing import Dict, List import sys sys.path.append(str(Path(__file__).parent.parent)) from utils import ( AudioProcessor, ESC50Dataset, QuestionGenerator, LLMQuestionGenerator, setup_logger, set_random_seed, generate_sample_durations_for_task, generate_single_clip_duration, build_count_task_audio, get_max_clip_num_to_be_joined ) class CountTaskGenerator: """Generator for counting task dataset.""" def __init__(self, config: Dict, logger): """ Initialize count task generator. Args: config: Configuration dictionary logger: Logger instance """ self.config = config self.logger = logger self.task_config = config['tasks']['count'] # Initialize components self.dataset = ESC50Dataset( config['esc50']['metadata_path'], config['esc50']['audio_path'], config # Pass config for class subset loading ) self.audio_processor = AudioProcessor( crossfade_duration=config['audio']['crossfade_duration'], silence_duration=config['audio']['silence_duration'], with_silence=config['audio']['with_silence'], normalize=config['audio']['normalize'], normalize_target_dBFS=config['audio']['normalize_target_dBFS'], synthetic_silence_path=config['synthetic_silence']['path'] ) self.question_generator = QuestionGenerator( num_options=config['mcq']['num_options'], option_labels=config['mcq']['option_labels'], distractor_strategy=config['mcq']['distractor_strategy'] ) # Initialize LLM question generator self.llm_enabled = config.get('llm', {}).get('enabled', False) self.llm_generator = LLMQuestionGenerator( enabled=self.llm_enabled, template_questions=self.task_config ) if self.llm_enabled: logger.info("LLM question generation enabled (local Llama 3.1 8B)") else: logger.info("Using template-based question generation") # Duration settings from config self.min_clip_duration = config['audio']['min_clip_duration'] self.max_clip_duration = config['audio']['max_clip_duration'] self.source_clip_duration = config['audio'].get('source_clip_duration', 5.0) self.min_silence_ms = config['audio'].get('min_silence_duration', 100) self.max_extra_silence_per_gap_ms = config['audio'].get('max_extra_silence_per_gap', 500) # Small crossfade within same-source repetitions (for consecutive mode) self.crossfade_within_source_ms = config['audio'].get('crossfade_within_source', 50) self.task_duration_hours = self.task_config['task_duration_size'] # Ordering mode: "random" or "consecutive" # random: Clips shuffled (A B A C B A C) - tests sound recognition # consecutive: Same-source grouped (AAA BBB CCC) - easier self.ordering_mode = self.task_config.get('ordering_mode', 'random') logger.info(f"Count task ordering mode: {self.ordering_mode}") # Set up output paths self.output_base = Path(config['output']['base_path']) / 'count' self.output_base.mkdir(parents=True, exist_ok=True) self.audio_output = self.output_base / 'audios' self.audio_output.mkdir(parents=True, exist_ok=True) def create_sampling_list(self, parent_list: List, n_sampling: int) -> List: """ Sample elements from parent list with replacement. Args: parent_list: List to sample from n_sampling: Number of samples Returns: List of sampled elements """ return [random.choice(parent_list) for _ in range(n_sampling)] def generate_sample(self, sample_id: int, target_unique_count: int = None, target_duration_seconds: float = None) -> Dict: """ Generate a single count task sample. Pipeline for COUNT task: 1. Use pre-generated target duration (or generate if not provided) 2. Calculate max clips that can fit 3. Pick N unique classes (N <= max_clips, since each source needs at least 1 clip) 4. For each class, sample one audio clip 5. Calculate repetitions to fill target duration 6. Based on ordering_mode: - "random": Shuffle clips (A B A C B A C) - tests recognition - "consecutive": Group same-class (AAA BBB CCC) - easier 7. Insert silences between clips 8. Distribute remainder as random extra silences Args: sample_id: Sample ID number target_unique_count: Target number of unique sounds (for balanced distribution) target_duration_seconds: Pre-generated target duration (from generate_sample_durations_for_task) Returns: Dictionary with sample metadata """ # Use pre-generated duration or generate one (backward compatibility) if target_duration_seconds is not None: clip_duration_seconds = target_duration_seconds else: clip_duration_seconds = generate_single_clip_duration( self.min_clip_duration, self.max_clip_duration ) # Calculate max clips that can fit in target duration max_clips, remainder_seconds = get_max_clip_num_to_be_joined( clip_duration_seconds, self.source_clip_duration, self.min_silence_ms ) # Ensure at least 1 clip max_clips = max(1, max_clips) max_clips_per_sample = self.task_config.get('max_clips_per_sample', 10) # Calculate valid range: n_unique_audios can be 1 to max_clips_per_sample # but cannot exceed what physically fits or available categories max_unique_for_sample = min(max_clips, max_clips_per_sample, len(self.dataset.CATEGORIES)) if max_unique_for_sample < 1: raise ValueError( f"Sample {sample_id}: Cannot generate sample - max_unique_for_sample={max_unique_for_sample}. " f"max_clips={max_clips}, max_clips_per_sample={max_clips_per_sample}, " f"available_categories={len(self.dataset.CATEGORIES)}, duration={clip_duration_seconds:.1f}s. " f"Increase min_clip_duration or reduce max_clips_per_sample." ) # Determine n_unique_audios - use target from balanced distribution or random if target_unique_count is not None: # Clamp target to what this specific sample duration can fit # Short samples can't fit all possible answers, so we clamp down n_unique_audios = min(target_unique_count, max_unique_for_sample) if n_unique_audios != target_unique_count: self.logger.debug( f"Sample {sample_id}: Clamped target from {target_unique_count} to {n_unique_audios} " f"(duration={clip_duration_seconds:.1f}s can only fit {max_clips} clips)" ) else: # No target specified - randomly select from valid range n_unique_audios = random.randint(1, max_unique_for_sample) self.logger.debug( f"Sample {sample_id}: target={clip_duration_seconds:.1f}s, max_clips={max_clips}, " f"n_unique_audios={n_unique_audios}" ) # Sample unique categories - use least-used categories for balanced distribution selected_categories = self.dataset.get_least_used_categories(n_unique_audios) # Track usage of all selected categories for cat in selected_categories: self.dataset.category_usage_counts[cat] += 1 # Sample one file from each unique category source_files = [] source_paths = [] source_categories = [] for category in selected_categories: filename, filepath = self.dataset.sample_file_from_category(category) source_files.append(filename) source_paths.append(filepath) source_categories.append(category) # Load unique source audios source_audios = [] for file_path in source_paths: audio = self.audio_processor.load_audio(file_path) source_audios.append(audio) # Build audio using configured ordering mode final_audio, clip_sequence, build_metadata = build_count_task_audio( source_audios, source_categories, clip_duration_seconds, ordering_mode=self.ordering_mode, source_clip_duration_seconds=self.source_clip_duration, min_silence_ms=self.min_silence_ms, max_extra_silence_per_gap_ms=self.max_extra_silence_per_gap_ms, crossfade_within_source_ms=self.crossfade_within_source_ms ) # Save the audio output_audio_path = self.audio_output / f"{sample_id}.wav" final_audio.export(str(output_audio_path), format="wav") # Generate questions (using LLM if enabled) if self.llm_enabled and self.llm_generator: llm_questions = self.llm_generator.generate_count_questions( correct_count=n_unique_audios, categories_present=list(set(clip_sequence)) ) mcq_question_text = llm_questions.get('mcq_question') open_text_question_text = llm_questions.get('open_text_question') else: mcq_question_text = random.choice(self.task_config['mcq_questions']) open_text_question_text = random.choice(self.task_config['open_text_questions']) # Generate MCQ with options mcq_data = self.question_generator.generate_count_mcq( mcq_question_text, n_unique_audios, self.dataset.CATEGORIES ) # Generate open-text answer open_text_data = self.question_generator.generate_count_open_text( open_text_question_text, n_unique_audios ) # Create metadata metadata = { 'id': sample_id, 'audio_path': str(output_audio_path.relative_to(self.output_base.parent)), 'n_unique_sounds': n_unique_audios, 'total_clips': build_metadata['total_clips'], 'repetitions_per_source': build_metadata['repetitions_per_source'], 'ordering_mode': self.ordering_mode, 'source_files': source_files, 'source_categories': source_categories, 'clip_sequence': clip_sequence, 'unique_categories': sorted(list(set(source_categories))), 'target_duration_seconds': clip_duration_seconds, 'actual_duration_seconds': len(final_audio) / 1000.0, 'mcq_question': mcq_data['question'], 'mcq_options': mcq_data['options'], 'mcq_correct_answer': mcq_data['correct_answer'], 'open_text_question': open_text_data['question'], 'open_text_answer': open_text_data['correct_answer'], 'llm_generated': self.llm_enabled } self.logger.info( f"Generated count sample {sample_id}: {n_unique_audios} unique sounds, " f"{build_metadata['total_clips']} clips, {len(final_audio)/1000:.1f}s" ) return metadata def generate_dataset(self) -> tuple: """ Generate the complete count task dataset. Returns: Tuple of (mcq_csv_path, open_text_csv_path) """ # Generate sample durations upfront to exactly fill target duration sample_durations = generate_sample_durations_for_task( self.task_duration_hours, self.min_clip_duration, self.max_clip_duration ) num_samples = len(sample_durations) self.logger.info(f"Generating {num_samples} count task samples (target: {self.task_duration_hours}h, actual: {sum(sample_durations)/3600:.2f}h)...") # Calculate max clips each sample can fit based on duration max_clips_per_sample = self.task_config.get('max_clips_per_sample', 10) sample_max_clips = [] for duration in sample_durations: max_clips, _ = get_max_clip_num_to_be_joined( duration, self.source_clip_duration, self.min_silence_ms ) # Limit to config max and available categories max_for_sample = min(max_clips, max_clips_per_sample, len(self.dataset.CATEGORIES)) sample_max_clips.append(max_for_sample) # Create balanced distribution by assigning targets based on sample capacity # Sort samples by capacity to assign higher targets to samples that can fit them possible_answers = list(range(1, max_clips_per_sample + 1)) samples_per_answer = num_samples // len(possible_answers) remainder = num_samples % len(possible_answers) # Create list of (sample_idx, duration, max_clips_capacity) sample_info = [(i, sample_durations[i], sample_max_clips[i]) for i in range(num_samples)] # Sort by capacity (descending) - assign high targets to high-capacity samples sample_info.sort(key=lambda x: x[2], reverse=True) # Assign targets: distribute each answer count across samples balanced_assignments = [None] * num_samples assignment_pool = [] for answer in possible_answers: count = samples_per_answer + (1 if remainder > 0 else 0) assignment_pool.extend([answer] * count) remainder = max(0, remainder - 1) # Reverse pool so we assign high targets first (to high-capacity samples) assignment_pool.sort(reverse=True) for idx, (sample_idx, duration, capacity) in enumerate(sample_info): # Assign target, clamped to sample's capacity target = min(assignment_pool[idx], capacity) balanced_assignments[sample_idx] = target # Log the actual distribution after capacity clamping from collections import Counter distribution = Counter(balanced_assignments) self.logger.info(f"Balanced answer distribution (after capacity-aware assignment): {dict(sorted(distribution.items()))}") all_metadata = [] for i in range(num_samples): metadata = self.generate_sample( i, target_unique_count=balanced_assignments[i], target_duration_seconds=sample_durations[i] ) all_metadata.append(metadata) # Save MCQ CSV mcq_csv_path = self.output_base / 'count_mcq.csv' self._save_mcq_csv(all_metadata, mcq_csv_path) # Save open-text CSV open_text_csv_path = self.output_base / 'count_open_text.csv' self._save_open_text_csv(all_metadata, open_text_csv_path) # Save metadata CSV metadata_csv_path = self.output_base / 'count_metadata.csv' self._save_metadata_csv(all_metadata, metadata_csv_path) self.logger.info(f"Count task dataset generation complete!") self.logger.info(f" - MCQ CSV: {mcq_csv_path}") self.logger.info(f" - Open-text CSV: {open_text_csv_path}") self.logger.info(f" - Metadata CSV: {metadata_csv_path}") self.logger.info(f" - Audio files: {self.audio_output}") return mcq_csv_path, open_text_csv_path def _save_mcq_csv(self, metadata_list: List[Dict], output_path: Path): """Save MCQ format CSV.""" with open(output_path, 'w', newline='') as f: writer = csv.writer(f) # Header writer.writerow([ 'question', 'id', 'audio_path', 'optionA', 'optionB', 'optionC', 'optionD', 'correct', 'source_wavs', 'source_categories' ]) # Data rows for meta in metadata_list: writer.writerow([ meta['mcq_question'], meta['id'], meta['audio_path'], meta['mcq_options']['A'], meta['mcq_options']['B'], meta['mcq_options']['C'], meta['mcq_options']['D'], meta['mcq_correct_answer'], str(meta['source_files']), str(meta['unique_categories']) ]) def _save_open_text_csv(self, metadata_list: List[Dict], output_path: Path): """Save open-text format CSV.""" with open(output_path, 'w', newline='') as f: writer = csv.writer(f) # Header writer.writerow([ 'question', 'id', 'audio_path', 'answer', 'source_wavs', 'source_categories' ]) # Data rows for meta in metadata_list: writer.writerow([ meta['open_text_question'], meta['id'], meta['audio_path'], meta['open_text_answer'], str(meta['source_files']), str(meta['unique_categories']) ]) def _save_metadata_csv(self, metadata_list: List[Dict], output_path: Path): """Save detailed metadata CSV.""" with open(output_path, 'w', newline='') as f: writer = csv.writer(f) # Header writer.writerow([ 'id', 'audio_path', 'total_clips', 'n_unique_sounds', 'source_files', 'source_categories', 'unique_categories', 'ordering_mode', 'target_duration_s', 'actual_duration_s', 'llm_generated' ]) # Data rows for meta in metadata_list: writer.writerow([ meta['id'], meta['audio_path'], meta['total_clips'], meta['n_unique_sounds'], str(meta['source_files']), str(meta['source_categories']), str(meta['unique_categories']), meta.get('ordering_mode', 'random'), meta.get('target_duration_seconds', 0), meta.get('actual_duration_seconds', 0), meta.get('llm_generated', False) ]) def main(config_path: str = None): """Main entry point for count task generation.""" import yaml # Load configuration if config_path is None: config_path = Path(__file__).parent.parent / 'config.yaml' with open(config_path, 'r') as f: config = yaml.safe_load(f) # Set random seed set_random_seed(config['random_seed']) # Setup logger logger = setup_logger( 'count_task', log_file=str(Path(config['output']['base_path']) / config['logging']['log_file']), level=config['logging']['level'], console_output=config['logging']['console_output'] ) # Generate dataset generator = CountTaskGenerator(config, logger) generator.generate_dataset() if __name__ == '__main__': main()