scripts/make_dataset.py

"""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()