""" Conditional Count task generator for temporal reasoning dataset. Multi-hop task: First apply a temporal condition (before, after, between), then count events satisfying that condition. """ import random from collections import Counter from pathlib import Path from typing import Dict, List, Optional from tasks.multihop_base import MultihopBaseGenerator class ConditionalCountTaskGenerator(MultihopBaseGenerator): """Generates conditional_count task dataset samples.""" TASK_NAME = "conditional_count" def generate_sample( self, sample_id: int, target_duration_seconds: float = None, question_type: str = None, ) -> Optional[Dict]: """ Generate a single conditional_count sample. Pipeline: 1. Build a sequential scene with 4-8 events (some categories repeated) 2. Pick a random question type (count_after, count_before, count_between) 3. Select anchor sound(s) and target sound from the sequence 4. Compute the correct count 5. Generate MCQ and open-text questions """ n_events = random.randint( self.task_config.get("min_events", 4), self.task_config.get("max_events", 8), ) # Build scene — we allow repeated categories for counting n_unique = random.randint(max(2, n_events // 2), min(n_events, len(self.dataset.CATEGORIES))) categories_pool = self.dataset.sample_categories(n_unique) # Fill sequence to n_events by repeating some categories = list(categories_pool) while len(categories) < n_events: categories.append(random.choice(categories_pool)) random.shuffle(categories) # Build audio scene final_audio, source_files, events_meta = self._build_scene_from_categories( categories, target_duration_seconds ) # Save audio output_path = self.audio_output / f"{sample_id}.wav" final_audio.export(str(output_path), format="wav") # Select question type if question_type is None: question_type = random.choice(self.task_config["question_types"]) # Generate question mcq_data, open_data, answer, q_meta = self._generate_question( question_type, categories, events_meta ) if mcq_data is None: self.logger.warning(f"Sample {sample_id}: Failed to generate question") return None metadata = { "id": sample_id, "audio_path": str(output_path.relative_to(self.output_base.parent)), "n_events": n_events, "categories": categories, "source_files": source_files, "question_type": question_type, "target_duration_s": target_duration_seconds, "actual_duration_s": 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_data["question"], "open_text_answer": open_data["correct_answer"], **q_meta, } self.logger.info( f"Generated conditional_count sample {sample_id}: " f"{n_events} events, type={question_type}, answer={answer}" ) return metadata def _build_scene_from_categories( self, categories: List[str], target_duration_s: float ): """Build audio scene from pre-specified category sequence.""" from pydub import AudioSegment as PydubSegment from utils import concatenate_to_target_duration, generate_controlled_gap_durations n_events = len(categories) num_gaps = n_events - 1 # Estimate per-event duration gap_total_est = num_gaps * 500 # ~500ms avg gap avail = max(n_events * self.source_clip_duration, target_duration_s - gap_total_est / 1000) per_event_s = max(self.source_clip_duration, avail / n_events) source_files = [] audio_segments = [] for cat in categories: fname, fpath = self.dataset.sample_file_from_category(cat) audio = self.audio_processor.load_audio(fpath) audio_segments.append(audio) source_files.append(fname) # Gaps if num_gaps > 0: gap_durations = generate_controlled_gap_durations( num_gaps, min_gap_ms=300, max_gap_ms=1500, gap_multiplier=2.0 ) else: gap_durations = [] # Assemble result = audio_segments[0] current_ms = 0 events_meta = [ { "index": 0, "category": categories[0], "start_ms": 0, "end_ms": len(audio_segments[0]), "duration_ms": len(audio_segments[0]), } ] current_ms = len(audio_segments[0]) for i in range(1, n_events): gap_ms = gap_durations[i - 1] result = result + PydubSegment.silent(duration=gap_ms) current_ms += gap_ms event_start = current_ms result = result + audio_segments[i] current_ms += len(audio_segments[i]) events_meta.append( { "index": i, "category": categories[i], "start_ms": event_start, "end_ms": current_ms, "duration_ms": len(audio_segments[i]), "gap_before_ms": gap_ms, } ) return result, source_files, events_meta def _generate_question( self, question_type: str, categories: List[str], events_meta: List[Dict], ): """Generate question and compute answer for conditional count.""" unique_cats = list(set(categories)) n = len(categories) if question_type in ("count_after", "count_before"): # Pick anchor and target anchor_idx = random.randint(1, n - 2) if n > 2 else 0 anchor_sound = categories[anchor_idx] # Pick target sound (one that appears in the sequence) target_sound = random.choice(unique_cats) if question_type == "count_after": count = sum(1 for i in range(anchor_idx + 1, n) if categories[i] == target_sound) else: count = sum(1 for i in range(0, anchor_idx) if categories[i] == target_sound) # Format question mcq_templates = self.task_config["mcq_questions"][question_type] open_templates = self.task_config["open_text_questions"][question_type] mcq_text = mcq_templates.format( target_sound=target_sound, anchor_sound=anchor_sound ) open_text = open_templates.format( target_sound=target_sound, anchor_sound=anchor_sound ) mcq_data = self.question_generator.generate_count_mcq( mcq_text, count, self.dataset.CATEGORIES ) open_data = self.question_generator.generate_count_open_text(open_text, count) q_meta = { "anchor_sound": anchor_sound, "anchor_index": anchor_idx, "target_sound": target_sound, "correct_count": count, } return mcq_data, open_data, count, q_meta elif question_type == "count_between": if n < 3: return None, None, None, {} # Pick two anchor indices idx1, idx2 = sorted(random.sample(range(n), 2)) sound1 = categories[idx1] sound2 = categories[idx2] # Count events strictly between idx1 and idx2 between_cats = categories[idx1 + 1 : idx2] # Pick a target — or count all if random.random() < 0.5 and between_cats: target_sound = random.choice(list(set(between_cats))) count = between_cats.count(target_sound) mcq_text = self.task_config["mcq_questions"]["count_between"].format(target_sound=target_sound, sound1=sound1, sound2=sound2) open_text = self.task_config["open_text_questions"]["count_between"].format(target_sound=target_sound, sound1=sound1, sound2=sound2) else: target_sound = "all" count = len(between_cats) mcq_text = f"How many sounds occur between {sound1} and {sound2}?" open_text = f"How many sounds occur between {sound1} and {sound2}?" mcq_data = self.question_generator.generate_count_mcq( mcq_text, count, self.dataset.CATEGORIES ) open_data = self.question_generator.generate_count_open_text(open_text, count) q_meta = { "sound1": sound1, "sound2": sound2, "target_sound": target_sound, "correct_count": count, } return mcq_data, open_data, count, q_meta elif question_type == "count_during": # Simplified: pick anchor, count other events (all overlap conceptually # since events are sequential — we interpret "during" as overlapping time # window centered on the anchor) anchor_idx = random.randint(0, n - 1) anchor_sound = categories[anchor_idx] # For sequential audio, "during" is approximated as the anchor's neighbors # We count how many of target_sound appear adjacent to anchor target_sound = random.choice(unique_cats) count = categories.count(target_sound) - (1 if target_sound == anchor_sound else 0) count = max(0, count) mcq_text = self.task_config["mcq_questions"]["count_during"].format(target_sound=target_sound, anchor_sound=anchor_sound) open_text = self.task_config["open_text_questions"]["count_during"].format(target_sound=target_sound, anchor_sound=anchor_sound) mcq_data = self.question_generator.generate_count_mcq( mcq_text, count, self.dataset.CATEGORIES ) open_data = self.question_generator.generate_count_open_text(open_text, count) q_meta = { "anchor_sound": anchor_sound, "target_sound": target_sound, "correct_count": count, } return mcq_data, open_data, count, q_meta elif question_type == "count_overlap": # In sequential audio, overlap count is 0 by construction unless # we build overlap scenes. Simplify: count events near anchor. anchor_idx = random.randint(0, n - 1) anchor_sound = categories[anchor_idx] # Count adjacent events as "overlapping" conceptually neighbors = set() if anchor_idx > 0: neighbors.add(anchor_idx - 1) if anchor_idx < n - 1: neighbors.add(anchor_idx + 1) count = len(neighbors) mcq_text = self.task_config["mcq_questions"]["count_overlap"].format(anchor_sound=anchor_sound, target_sound=anchor_sound) open_text = self.task_config["open_text_questions"]["count_overlap"].format(anchor_sound=anchor_sound, target_sound=anchor_sound) mcq_data = self.question_generator.generate_count_mcq( mcq_text, count, self.dataset.CATEGORIES ) open_data = self.question_generator.generate_count_open_text(open_text, count) q_meta = {"anchor_sound": anchor_sound, "correct_count": count} return mcq_data, open_data, count, q_meta return None, None, None, {} def main(config_path: str = None): """Main entry point for conditional_count task generation.""" import yaml 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(config["random_seed"]) logger = setup_logger( "conditional_count_task", log_file=str(Path(config["output"]["base_path"]) / config["logging"]["log_file"]), level=config["logging"]["level"], console_output=config["logging"]["console_output"], ) generator = ConditionalCountTaskGenerator(config, logger) generator.generate_dataset() if __name__ == "__main__": main()