"""Phase 1 dataset generation. Generates Bloom-level student questions and writes them to JSONL. """ import json import os import random import sys import time from dotenv import load_dotenv # Let this script run from project root. sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) import config # noqa: E402 from scripts.generators import build_source # noqa: E402 class BloomQuestionGenerator: """Generate labeled student questions by Bloom level.""" def __init__(self, generation_config, source): """Store generation config and provider source.""" self.cfg = generation_config self.source = source def _build_prompt(self, domain_label, topic, sublevel): """Build one prompt for a domain/topic/sublevel combo.""" readable_domain = domain_label.replace("_", " ") definition = config.SUBLEVEL_DEFINITIONS[sublevel] anchors = config.SUBLEVEL_FEWSHOT[sublevel] all_stems = config.SUBLEVEL_STEMS[sublevel] # Sample stems each time to vary phrasing. sampled_stems = random.sample(all_stems, min(6, len(all_stems))) template = random.choice(self.cfg.instruction_templates) instruction = template.format( n=self.cfg.batch_size, topic=topic, domain=readable_domain ) trap_count = max(1, round(self.cfg.batch_size * self.cfg.anti_shortcut_fraction)) anchor_block = "\n".join(f" - {a}" for a in anchors) stem_block = ", ".join(f'"{s}..."' for s in sampled_stems) prompt = ( f"{instruction}\n\n" f"Every question must sit at the cognitive level '{sublevel}', " f"meaning a question that {definition}.\n\n" f"Reference examples of '{sublevel}' questions (different subject, " f"shown only to convey the cognitive level, do not copy them):\n" f"{anchor_block}\n\n" f"Vary the phrasing widely. You may open questions in different ways, " f"for example: {stem_block}. Do not start most questions the same way.\n\n" f"Important diversity requirement: at least {trap_count} of the " f"{self.cfg.batch_size} questions must be 'counter-stereotypical', " f"meaning their opening word would NOT normally signal the " f"'{sublevel}' level, yet the question genuinely belongs to it. " f"For example a critical, limitation-probing question that begins " f"with 'What' or 'Why', or a factual recall question that begins " f"with 'Define' or 'List' rather than 'What'. The cognitive intent, " f"not the first word, determines the level.\n\n" f"Constraints:\n" f" - Each question must genuinely match the '{sublevel}' level.\n" f" - Vary length and structure across the questions.\n" f" - Keep each question to a single sentence.\n\n" f"Return a JSON object with a single key \"questions\" whose value is " f"an array of {self.cfg.batch_size} question strings. " f'Example: {{"questions": ["...", "..."]}}' ) return prompt def _parse_questions(self, raw_text): """Parse generated questions from JSON response text.""" text = raw_text.strip() if text.startswith("```"): text = text.strip("`") if text.startswith("json"): text = text[4:] parsed = json.loads(text) if isinstance(parsed, dict): questions = parsed.get("questions", []) elif isinstance(parsed, list): questions = parsed else: questions = [] return [q.strip() for q in questions if isinstance(q, str) and q.strip()] def _generate_batch(self, prompt): """Call provider with retry and backoff.""" for attempt in range(self.cfg.max_retries): try: raw = self.source.complete(prompt) return self._parse_questions(raw) except json.JSONDecodeError: wait = self.cfg.backoff_base_seconds * (2 ** attempt) print(f" JSON parse failed, retrying in {wait:.1f}s") time.sleep(wait) except Exception as error: wait = self.cfg.backoff_base_seconds * (2 ** attempt) print(f" API error ({type(error).__name__}), retrying in {wait:.1f}s") time.sleep(wait) print(" giving up on this batch after max retries") return [] @staticmethod def _completed_combos(output_path): """Return already finished (domain, sublevel) combos.""" done = set() if not os.path.exists(output_path): return done with open(output_path, "r", encoding="utf-8") as handle: for line in handle: try: record = json.loads(line) done.add((record["domain"], record["bloom_sublevel"])) except (json.JSONDecodeError, KeyError): continue return done def _write_records(self, records, output_path): """Append records to JSONL output.""" with open(output_path, "a", encoding="utf-8") as handle: for record in records: handle.write(json.dumps(record, ensure_ascii=False) + "\n") def _generate_for_combo(self, domain_label, subtopics, sublevel, n_batches, split): """Generate records for one (domain, sublevel).""" records = [] for batch_index in range(n_batches): topic = random.choice(subtopics) prompt = self._build_prompt(domain_label, topic, sublevel) questions = self._generate_batch(prompt) for question in questions: records.append({ "question": question, "bloom_sublevel": sublevel, "bloom_class": config.SUBLEVEL_TO_CLASS[sublevel], "domain": domain_label, "topic": topic, "split": split, "source": self.source.label, }) print(f" batch {batch_index + 1}/{n_batches} on '{topic}': " f"{len(questions)} questions") time.sleep(self.cfg.seconds_between_calls) return records def run(self): """Run full generation loop.""" output_path = self.cfg.output_path already_done = self._completed_combos(output_path) if already_done: print(f"Resuming: {len(already_done)} combos already complete.\n") plan = [ (config.TRAIN_DOMAINS, "train", self.cfg.train_batches_per_combo), (config.OOD_DOMAINS, "ood", self.cfg.ood_batches_per_combo), ] total_records = 0 for domain_dict, split, n_batches in plan: for domain_label, subtopics in domain_dict.items(): for sublevel in config.BLOOM_SUBLEVELS: if (domain_label, sublevel) in already_done: print(f"[skip] {domain_label} / {sublevel} (already done)") continue print(f"[{split}] {domain_label} / {sublevel}") records = self._generate_for_combo( domain_label, subtopics, sublevel, n_batches, split ) self._write_records(records, output_path) total_records += len(records) print(f" -> wrote {len(records)} records " f"(running total {total_records})\n") print(f"Done. Output at {output_path}") self._print_summary(output_path) @staticmethod def _print_summary(output_path): """Print simple distribution summary.""" class_counts, domain_counts, split_counts = {}, {}, {} with open(output_path, "r", encoding="utf-8") as handle: for line in handle: record = json.loads(line) class_counts[record["bloom_class"]] = class_counts.get(record["bloom_class"], 0) + 1 domain_counts[record["domain"]] = domain_counts.get(record["domain"], 0) + 1 split_counts[record["split"]] = split_counts.get(record["split"], 0) + 1 print("\nClass distribution:", class_counts) print("Domain distribution:", domain_counts) print("Split distribution:", split_counts) def main(): """Load env, build source, and run generation.""" load_dotenv() generation_config = config.GenerationConfig() source = build_source( generation_config.provider, generation_config.model_name, generation_config.temperature, ) print(f"Generating with source: {source.label}\n") generator = BloomQuestionGenerator(generation_config, source) generator.run() if __name__ == "__main__": main()