tinymind-native-8b-remote-handoff/bundle/data/kaggle_scicode_ingestor.py

"""Kaggle SciCode ingestion for TinyMind continued training.

SciCode is a benchmark, so this module is deliberately contamination-aware:
dev rows with released solutions can become supplemental scientific-code SFT,
while test rows are quarantined by default as evaluation evidence and are not
included in the main training JSONL.
"""

from __future__ import annotations

from dataclasses import dataclass
from datetime import datetime, timezone
import hashlib
import json
from pathlib import Path
import re
from typing import Any, Callable, Iterable


SCHEMA_VERSION = "tinymind-kaggle-scicode-ingest-v1"
DEFAULT_DATASET_SLUG = "open-benchmarks/scicode"
DEFAULT_LICENSE = "apache-2.0"
TRAIN_SOURCE = "logic_agent_code_scicode_kaggle_dev"
QUARANTINE_SOURCE = "scicode_kaggle_test_quarantine"


def _now() -> str:
    return datetime.now(timezone.utc).isoformat()


def _sha(payload: object) -> str:
    return hashlib.sha256(json.dumps(payload, ensure_ascii=False, sort_keys=True).encode("utf-8")).hexdigest()


def _norm(text: str) -> str:
    return re.sub(r"\s+", " ", text or "").strip()


def _as_text(value: object) -> str:
    if value is None:
        return ""
    if isinstance(value, str):
        return _norm(value)
    return _norm(json.dumps(value, ensure_ascii=False, sort_keys=True))


def _load_jsonl(path: Path) -> list[dict[str, Any]]:
    rows: list[dict[str, Any]] = []
    with path.open("r", encoding="utf-8", errors="replace") as handle:
        for line in handle:
            if not line.strip():
                continue
            rows.append(json.loads(line))
    return rows


def _record_id(prefix: str, row: dict[str, Any], suffix: str = "") -> str:
    digest = hashlib.sha1((_sha(row) + suffix).encode("utf-8")).hexdigest()[:16]
    return f"{prefix}-{digest}"


def _prompt_for_problem(row: dict[str, Any]) -> str:
    parts = [
        f"Problem: {_as_text(row.get('problem_name'))} ({_as_text(row.get('problem_id'))})",
        _as_text(row.get("problem_description_main")),
        _as_text(row.get("problem_background_main")),
        "I/O contract:\n" + _as_text(row.get("problem_io")),
        "Required dependencies:\n" + _as_text(row.get("required_dependencies")),
        "Write a correct Python scientific-computing solution. Keep the code executable and explain only essential numerical assumptions.",
    ]
    return "\n\n".join(part for part in parts if part.strip())


def _assistant_for_problem(row: dict[str, Any]) -> str:
    solution = _as_text(row.get("general_solution"))
    tests = _as_text(row.get("general_tests"))
    if tests:
        return f"{solution}\n\nValidation tests from source:\n{tests}"
    return solution


def _prompt_for_step(row: dict[str, Any], step: dict[str, Any]) -> str:
    parts = [
        f"Parent problem: {_as_text(row.get('problem_name'))} ({_as_text(row.get('problem_id'))})",
        _as_text(row.get("problem_description_main")),
        "Step: " + _as_text(step.get("step_number")),
        _as_text(step.get("step_description_prompt")),
        _as_text(step.get("step_background")),
        "Function header:\n" + _as_text(step.get("function_header")),
        "Required dependencies:\n" + _as_text(row.get("required_dependencies")),
        "Return only the implementation needed for this step unless tests require extra helpers.",
    ]
    return "\n\n".join(part for part in parts if part.strip())


def _assistant_for_step(step: dict[str, Any]) -> str:
    code = _as_text(step.get("ground_truth_code"))
    return_line = _as_text(step.get("return_line"))
    if return_line and return_line not in code:
        return f"{code}\n{return_line}"
    return code


def _chat_record(
    *,
    record_id: str,
    source: str,
    user: str,
    assistant: str,
    row: dict[str, Any],
    domain: str,
    split: str,
    license_value: str,
    kind: str,
    loss_weight: float,
) -> dict[str, Any] | None:
    user = _norm(user)
    assistant = _norm(assistant)
    if len(user) < 60 or len(assistant) < 40:
        return None
    fingerprint = _sha({"split": split, "kind": kind, "row": row, "user": user, "assistant": assistant})
    return {
        "id": record_id,
        "source": source,
        "license": license_value,
        "messages": [
            {
                "role": "system",
                "content": (
                    "You are TinyMind scientific-code tutor. Solve with grounded numerical reasoning, "
                    "clear Python, and source-verifiable constraints."
                ),
            },
            {"role": "user", "content": user},
            {"role": "assistant", "content": assistant},
        ],
        "metadata": {
            "domain": domain,
            "dataset": DEFAULT_DATASET_SLUG,
            "split": split,
            "record_kind": kind,
            "fingerprint_sha256": fingerprint,
            "source_problem_id": str(row.get("problem_id", "")),
            "source_problem_name": str(row.get("problem_name", "")),
            "loss_weight": loss_weight,
            "contamination_policy": "dev_solution_supplemental_only_test_split_quarantined",
        },
    }


@dataclass(frozen=True)
class SciCodeIngestPolicy:
    dataset_slug: str = DEFAULT_DATASET_SLUG
    file_path: str = ""
    max_records: int = 512
    include_problem_level: bool = True
    include_sub_steps: bool = True
    quarantine_test_split: bool = True
    license_value: str = DEFAULT_LICENSE
    loss_weight: float = 1.18


class SciCodeKaggleIngestor:
    def __init__(
        self,
        policy: SciCodeIngestPolicy | None = None,
        *,
        dataset_loader: Callable[[str, str], Any] | None = None,
        dataset_downloader: Callable[[str], str] | None = None,
    ):
        self.policy = policy or SciCodeIngestPolicy()
        self.dataset_loader = dataset_loader
        self.dataset_downloader = dataset_downloader

    def _download_dir(self) -> Path:
        if self.dataset_downloader:
            return Path(self.dataset_downloader(self.policy.dataset_slug))
        import kagglehub

        return Path(kagglehub.dataset_download(self.policy.dataset_slug))

    def _load_rows(self, file_path: str | None = None) -> tuple[list[dict[str, Any]], dict[str, Any]]:
        file_path = file_path if file_path is not None else self.policy.file_path
        if file_path:
            if self.dataset_loader:
                dataset = self.dataset_loader(self.policy.dataset_slug, file_path)
            else:
                import kagglehub
                from kagglehub import KaggleDatasetAdapter

                dataset = kagglehub.dataset_load(
                    KaggleDatasetAdapter.HUGGING_FACE,
                    self.policy.dataset_slug,
                    file_path,
                )
            rows = self._rows_from_dataset(dataset)
            return rows, {"loader": "kagglehub_hf_adapter", "file_path": file_path}

        root = self._download_dir()
        dev_path = root / "problems_dev.jsonl"
        if not dev_path.exists():
            candidates = sorted(root.glob("*.jsonl"))
            if not candidates:
                raise FileNotFoundError(f"No JSONL files found in {root}")
            dev_path = candidates[0]
        return _load_jsonl(dev_path), {"loader": "kagglehub_dataset_download", "file_path": str(dev_path)}

    def _load_quarantine_rows(self) -> tuple[list[dict[str, Any]], dict[str, Any]]:
        if not self.policy.quarantine_test_split:
            return [], {"loader": "disabled"}
        root = self._download_dir()
        test_path = root / "problems_test.jsonl"
        if not test_path.exists():
            return [], {"loader": "kagglehub_dataset_download", "file_path": str(test_path), "missing": True}
        return _load_jsonl(test_path), {"loader": "kagglehub_dataset_download", "file_path": str(test_path)}

    def _rows_from_dataset(self, dataset: Any) -> list[dict[str, Any]]:
        if hasattr(dataset, "to_list"):
            return list(dataset.to_list())
        if isinstance(dataset, dict):
            rows: list[dict[str, Any]] = []
            for split_rows in dataset.values():
                rows.extend(self._rows_from_dataset(split_rows))
            return rows
        return [dict(row) for row in dataset]

    def _normalize_train_rows(self, rows: Iterable[dict[str, Any]]) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
        records: list[dict[str, Any]] = []
        blocked: list[dict[str, Any]] = []
        for row in rows:
            if len(records) >= self.policy.max_records:
                break
            if self.policy.include_problem_level and row.get("general_solution"):
                record = _chat_record(
                    record_id=_record_id("scicode-problem", row, "problem"),
                    source=TRAIN_SOURCE,
                    user=_prompt_for_problem(row),
                    assistant=_assistant_for_problem(row),
                    row=row,
                    domain="coding_python",
                    split="dev",
                    license_value=self.policy.license_value,
                    kind="problem_solution",
                    loss_weight=self.policy.loss_weight,
                )
                if record:
                    records.append(record)
                else:
                    blocked.append({"problem_id": row.get("problem_id"), "reason": "problem_record_too_short"})

            if not self.policy.include_sub_steps:
                continue
            for step in row.get("sub_steps") or []:
                if len(records) >= self.policy.max_records:
                    break
                if not isinstance(step, dict) or not step.get("ground_truth_code"):
                    blocked.append({"problem_id": row.get("problem_id"), "reason": "missing_step_solution"})
                    continue
                record = _chat_record(
                    record_id=_record_id("scicode-step", {"row": row, "step": step}, str(step.get("step_number", ""))),
                    source=TRAIN_SOURCE,
                    user=_prompt_for_step(row, step),
                    assistant=_assistant_for_step(step),
                    row={**row, "sub_step": step.get("step_number")},
                    domain="coding_python",
                    split="dev",
                    license_value=self.policy.license_value,
                    kind="sub_step_solution",
                    loss_weight=self.policy.loss_weight,
                )
                if record:
                    records.append(record)
                else:
                    blocked.append({"problem_id": row.get("problem_id"), "step": step.get("step_number"), "reason": "step_record_too_short"})
        return records, blocked

    def _normalize_quarantine_rows(self, rows: Iterable[dict[str, Any]]) -> list[dict[str, Any]]:
        quarantine: list[dict[str, Any]] = []
        for row in rows:
            prompt = _prompt_for_problem(row)
            if len(prompt) < 60:
                continue
            quarantine.append(
                {
                    "id": _record_id("scicode-quarantine", row, "test"),
                    "source": QUARANTINE_SOURCE,
                    "license": self.policy.license_value,
                    "messages": [
                        {"role": "system", "content": "SciCode held-out evaluation prompt. Do not use for main training."},
                        {"role": "user", "content": prompt},
                    ],
                    "metadata": {
                        "dataset": DEFAULT_DATASET_SLUG,
                        "split": "test",
                        "record_kind": "heldout_prompt_only",
                        "fingerprint_sha256": _sha(row),
                        "source_problem_id": str(row.get("problem_id", "")),
                        "source_problem_name": str(row.get("problem_name", "")),
                        "main_training_allowed": False,
                        "reason": "SciCode test split is quarantined to preserve evaluation integrity.",
                    },
                }
            )
        return quarantine

    def write_jsonl(self, out_dir: str | Path) -> dict[str, Any]:
        out = Path(out_dir)
        out.mkdir(parents=True, exist_ok=True)

        rows, train_fetch = self._load_rows()
        records, blocked = self._normalize_train_rows(rows)
        quarantine_rows, quarantine_fetch = self._load_quarantine_rows()
        quarantine_records = self._normalize_quarantine_rows(quarantine_rows)

        train_path = out / "scicode_sft_train.jsonl"
        quarantine_path = out / "scicode_quarantine_eval_prompts.jsonl"
        train_path.write_text("\n".join(json.dumps(row, ensure_ascii=False, sort_keys=True) for row in records), encoding="utf-8")
        quarantine_path.write_text(
            "\n".join(json.dumps(row, ensure_ascii=False, sort_keys=True) for row in quarantine_records),
            encoding="utf-8",
        )

        manifest = {
            "schema_version": SCHEMA_VERSION,
            "created_at": _now(),
            "dataset": self.policy.dataset_slug,
            "license": self.policy.license_value,
            "source_url": "https://www.kaggle.com/datasets/open-benchmarks/scicode",
            "reference_url": "https://scicode-bench.github.io/",
            "train_jsonl": str(train_path),
            "quarantine_jsonl": str(quarantine_path),
            "train_records": len(records),
            "blocked_records": len(blocked),
            "quarantine_records": len(quarantine_records),
            "blocked": blocked[:100],
            "fetch": {"train": train_fetch, "quarantine": quarantine_fetch},
            "policy": {
                "max_records": self.policy.max_records,
                "include_problem_level": self.policy.include_problem_level,
                "include_sub_steps": self.policy.include_sub_steps,
                "quarantine_test_split": self.policy.quarantine_test_split,
                "loss_weight": self.policy.loss_weight,
            },
            "claim_gate": {
                "scicode_ingested": len(records) > 0,
                "main_training_allowed": len(records) > 0,
                "test_split_quarantined": self.policy.quarantine_test_split,
                "external_benchmark_claim_allowed": False,
                "world_best_claim_allowed": False,
                "reason": "SciCode dev solution rows are supplemental SFT only; held-out/test rows are quarantined and external claims require official evaluation.",
            },
            "dataset_sha256": hashlib.sha256(
                "\n".join(json.dumps(row, ensure_ascii=False, sort_keys=True) for row in records).encode("utf-8")
            ).hexdigest(),
        }
        manifest_path = out / "scicode_ingest_manifest.json"
        manifest["manifest_path"] = str(manifest_path)
        manifest_path.write_text(json.dumps(manifest, ensure_ascii=False, indent=2, sort_keys=True), encoding="utf-8")
        return manifest