""" Between Events task generator for temporal reasoning dataset. Multi-hop task: Identifies or counts sounds occurring inside the temporal window between two anchor events in a sequential audio scene. """ import random from pathlib import Path from typing import Dict, List, Optional from utils import ( setup_logger, set_random_seed, concatenate_to_target_duration, generate_controlled_gap_durations, ) from tasks.multihop_base import MultihopBaseGenerator class BetweenEventsTaskGenerator(MultihopBaseGenerator): """Generates between_events task dataset samples.""" TASK_NAME = "between_events" def generate_sample( self, sample_id: int, target_duration_seconds: float = None, question_type: str = None, ) -> Optional[Dict]: """ Generate a single between_events sample. Pipeline: 1. Build sequential scene with 5-8 unique events 2. Pick two anchor events (sound1, sound2) with at least 1 event between 3. Identify or count events in the window 4. Generate question """ n_events = random.randint( self.task_config.get("min_events", 5), self.task_config.get("max_events", 8), ) # Need unique categories so "between" is unambiguous n_events = min(n_events, len(self.dataset.CATEGORIES)) categories = self.dataset.sample_categories(n_events) random.shuffle(categories) # Build audio final_audio, source_files, events_meta = self._build_scene( categories, target_duration_seconds ) output_path = self.audio_output / f"{sample_id}.wav" final_audio.export(str(output_path), format="wav") if question_type is None: question_type = random.choice(self.task_config["question_types"]) # Pick two anchors with at least one event between them if n_events < 3: self.logger.warning(f"Sample {sample_id}: Not enough events for between") return None idx1 = random.randint(0, n_events - 3) idx2 = random.randint(idx1 + 2, n_events - 1) sound1 = categories[idx1] sound2 = categories[idx2] between_cats = categories[idx1 + 1 : idx2] mcq_data, open_data, q_meta = self._generate_question( question_type, sound1, sound2, between_cats, categories ) if mcq_data is None: 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, "sound1": sound1, "sound2": sound2, "between_events": between_cats, "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 between_events sample {sample_id}: " f"{n_events} events, type={question_type}" ) return metadata def _build_scene(self, categories, target_duration_s): """Build sequential audio scene.""" from pydub import AudioSegment as PydubSegment n = len(categories) num_gaps = n - 1 per_event_s = max(self.source_clip_duration, (target_duration_s or 30) / n) 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) 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 = [] result = audio_segments[0] events_meta = [{"index": 0, "category": categories[0]}] for i in range(1, n): gap_ms = gap_durations[i - 1] result = result + PydubSegment.silent(duration=gap_ms) result = result + audio_segments[i] events_meta.append({"index": i, "category": categories[i]}) return result, source_files, events_meta def _generate_question( self, question_type, sound1, sound2, between_cats, all_cats ): """Generate between_events question.""" if question_type == "identify_between": if not between_cats: return None, None, {} correct = random.choice(between_cats) mcq_text = self.task_config["mcq_questions"]["identify_between"].format(sound1=sound1, sound2=sound2) open_text = self.task_config["open_text_questions"]["identify_between"].format(sound1=sound1, sound2=sound2) mcq_data = self.question_generator.generate_category_mcq( mcq_text, correct, between_cats, self.dataset.CATEGORIES ) open_data = self.question_generator.generate_category_open_text( open_text, correct ) q_meta = {"correct_category": correct} return mcq_data, open_data, q_meta elif question_type == "count_between": count = len(between_cats) mcq_text = self.task_config["mcq_questions"]["count_between"].format(sound1=sound1, sound2=sound2) open_text = self.task_config["open_text_questions"]["count_between"].format(sound1=sound1, sound2=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 = {"correct_count": count} return mcq_data, open_data, q_meta elif question_type == "yes_no_between": if not between_cats: return None, None, {} # 50% true, 50% false if random.random() < 0.5: target_sound = random.choice(between_cats) correct = True else: not_between = [c for c in all_cats if c not in between_cats and c != sound1 and c != sound2] if not not_between: target_sound = random.choice(between_cats) correct = True else: target_sound = random.choice(not_between) correct = False mcq_text = self.task_config["mcq_questions"]["yes_no_between"].format(target_sound=target_sound, sound1=sound1, sound2=sound2) open_text = self.task_config["open_text_questions"]["yes_no_between"].format(target_sound=target_sound, sound1=sound1, sound2=sound2) mcq_data = self.question_generator.generate_yes_no_mcq(mcq_text, correct) open_data = self.question_generator.generate_yes_no_open_text(open_text, correct) q_meta = {"target_sound": target_sound, "correct_answer": correct} return mcq_data, open_data, q_meta return None, None, {} def main(config_path: str = None): """Main entry point for between_events 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( "between_events_task", log_file=str(Path(config["output"]["base_path"]) / config["logging"]["log_file"]), level=config["logging"]["level"], console_output=config["logging"]["console_output"], ) generator = BetweenEventsTaskGenerator(config, logger) generator.generate_dataset() if __name__ == "__main__": main()