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

"""Kaggle benchmark mix ingestion for TinyMind SFT supplements.

The sources handled here are valuable training/evaluation artifacts but also
benchmarks.  The ingestor therefore writes provenance, caps per source, and
keeps rank/official claims blocked.  Rows are suitable only as supplemental
continued-training data unless later excluded from the corresponding benchmark
evaluation.
"""

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, Iterable


SCHEMA_VERSION = "tinymind-kaggle-benchmark-mix-v1"


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: object) -> str:
    if text is None:
        return ""
    if not isinstance(text, str):
        text = json.dumps(text, ensure_ascii=False, sort_keys=True)
    return re.sub(r"\s+", " ", text).strip()


def _jsonl_rows(path: Path, limit: int | None = None) -> 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))
            if limit and len(rows) >= limit:
                break
    return rows


def _csv_rows(path: Path, limit: int | None = None) -> list[dict[str, Any]]:
    import pandas as pd

    df = pd.read_csv(path)
    if limit:
        df = df.head(limit)
    return [dict(row) for row in df.to_dict(orient="records")]


def _download(slug: str) -> Path:
    import kagglehub

    return Path(kagglehub.dataset_download(slug))


def _record(
    *,
    source: str,
    row: dict[str, Any],
    user: str,
    assistant: str,
    dataset: str,
    domain: str,
    license_value: str,
    loss_weight: float,
    record_kind: str,
    extra_meta: dict[str, Any] | None = None,
) -> dict[str, Any] | None:
    user = _norm(user)
    assistant = _norm(assistant)
    if len(user) < 30 or len(assistant) < 8:
        return None
    fingerprint = _sha({"dataset": dataset, "kind": record_kind, "row": row, "user": user, "assistant": assistant})
    metadata = {
        "dataset": dataset,
        "domain": domain,
        "record_kind": record_kind,
        "fingerprint_sha256": fingerprint,
        "loss_weight": loss_weight,
        "benchmark_contamination_policy": "supplemental_training_blocks_corresponding_official_claim",
    }
    if extra_meta:
        metadata.update(extra_meta)
    return {
        "id": f"kaggle-{hashlib.sha1(fingerprint.encode('utf-8')).hexdigest()[:16]}",
        "source": source,
        "license": license_value,
        "messages": [
            {
                "role": "system",
                "content": (
                    "You are TinyMind. Answer with grounded, concise, source-aware reasoning. "
                    "Keep benchmark-derived knowledge separate from official evaluation claims."
                ),
            },
            {"role": "user", "content": user},
            {"role": "assistant", "content": assistant},
        ],
        "metadata": metadata,
    }


@dataclass(frozen=True)
class KaggleBenchmarkMixPolicy:
    max_parsebench: int = 400
    max_simpleqa: int = 400
    max_multiloko: int = 800
    max_mgsm: int = 400
    max_livecodebench: int = 400
    multiloko_languages: tuple[str, ...] = (
        "english",
        "arabic",
        "bengali",
        "cantonese",
        "czech",
        "dutch",
        "farsi",
        "french",
        "german",
        "hindi",
        "indonesian",
        "italian",
        "japanese",
        "korean",
        "portuguese",
        "russian",
        "spanish",
        "thai",
    )


class KaggleBenchmarkMixIngestor:
    def __init__(
        self,
        policy: KaggleBenchmarkMixPolicy | None = None,
        *,
        roots: dict[str, str | Path] | None = None,
    ):
        self.policy = policy or KaggleBenchmarkMixPolicy()
        self.roots = {key: Path(value) for key, value in (roots or {}).items()}

    def _root(self, name: str, slug: str) -> Path:
        return self.roots.get(name) or _download(slug)

    def _parsebench(self) -> tuple[list[dict[str, Any]], dict[str, Any]]:
        root = self._root("parsebench", "llamaindex-org/parsebench")
        files = ["table.jsonl", "text_content.jsonl", "text_formatting.jsonl", "layout.jsonl", "chart.jsonl"]
        per_file = max(1, self.policy.max_parsebench // len(files))
        records: list[dict[str, Any]] = []
        for filename in files:
            path = root / filename
            if not path.exists():
                continue
            for row in _jsonl_rows(path, per_file):
                rule = _norm(row.get("rule"))
                expected = _norm(row.get("expected_markdown"))
                category = _norm(row.get("category"))
                row_type = _norm(row.get("type"))
                pdf = _norm(row.get("pdf"))
                if expected:
                    user = (
                        f"Convert the referenced document page into faithful markdown.\n"
                        f"PDF: {pdf}\nCategory: {category}\nRule type: {row_type}\n"
                        f"Preserve tables, headings, and content exactly when present."
                    )
                    assistant = expected
                else:
                    user = (
                        f"Explain the document parsing validation rule as an auditable checklist.\n"
                        f"PDF: {pdf}\nCategory: {category}\nRule type: {row_type}\nRule JSON: {rule}"
                    )
                    assistant = (
                        f"Validate category `{category}` with rule type `{row_type}`. "
                        f"Use this rule exactly: {rule}. Report mismatches with page-local evidence and avoid inventing missing document content."
                    )
                rec = _record(
                    source="parsebench_kaggle_document_tooling",
                    row=row,
                    user=user,
                    assistant=assistant,
                    dataset="llamaindex-org/parsebench",
                    domain="document_parsing",
                    license_value="apache-2.0",
                    loss_weight=1.12,
                    record_kind=f"parsebench_{category}_{row_type}",
                    extra_meta={"pdf": pdf, "category": category, "rule_type": row_type},
                )
                if rec:
                    records.append(rec)
                    if len(records) >= self.policy.max_parsebench:
                        return records, {"root": str(root), "records": len(records), "files": files}
        return records, {"root": str(root), "records": len(records), "files": files}

    def _simpleqa(self) -> tuple[list[dict[str, Any]], dict[str, Any]]:
        root = self._root("simpleqa", "deepmind/simpleqa-verified")
        path = root / "simpleqa_verified.csv"
        rows = _csv_rows(path, self.policy.max_simpleqa)
        records: list[dict[str, Any]] = []
        for row in rows:
            urls = _norm(row.get("urls"))
            user = (
                f"Answer the verified factual question exactly. If evidence is insufficient, say so.\n"
                f"Question: {_norm(row.get('problem'))}\n"
                f"Topic: {_norm(row.get('topic'))}\nAnswer type: {_norm(row.get('answer_type'))}\nSources: {urls}"
            )
            assistant = f"{_norm(row.get('answer'))}\n\nEvidence URLs: {urls}"
            rec = _record(
                source="simpleqa_verified_kaggle_factual_qa",
                row=row,
                user=user,
                assistant=assistant,
                dataset="deepmind/simpleqa-verified",
                domain="factual_qa",
                license_value="unknown-kaggle-dataset-card-required",
                loss_weight=0.85,
                record_kind="verified_factual_qa",
                extra_meta={
                    "topic": _norm(row.get("topic")),
                    "answer_type": _norm(row.get("answer_type")),
                    "multi_step": bool(row.get("multi_step")),
                    "requires_reasoning": bool(row.get("requires_reasoning")),
                    "external_urls": urls,
                },
            )
            if rec:
                records.append(rec)
        return records, {"root": str(root), "records": len(records), "file": str(path)}

    def _multiloko(self) -> tuple[list[dict[str, Any]], dict[str, Any]]:
        root = self._root("multiloko", "metaresearch/multiloko")
        base = root / "benchmark_data"
        records: list[dict[str, Any]] = []
        langs_seen: dict[str, int] = {}
        per_lang = max(1, self.policy.max_multiloko // max(1, len(self.policy.multiloko_languages)))
        for lang in self.policy.multiloko_languages:
            path = base / lang / "dev.jsonl"
            if not path.exists():
                continue
            count = 0
            for row in _jsonl_rows(path):
                user = (
                    f"Answer from the provided context in the same language as the question when natural.\n"
                    f"Language: {lang}\nContext:\n{_norm(row.get('text'))[:6000]}\n\nQuestion: {_norm(row.get('question'))}"
                )
                targets = row.get("targets") if isinstance(row.get("targets"), list) else []
                answer = _norm(row.get("target") or (targets[0] if targets else ""))
                rec = _record(
                    source="multiloko_kaggle_multilingual_grounding",
                    row=row,
                    user=user,
                    assistant=answer,
                    dataset="metaresearch/multiloko",
                    domain="multilingual_context_qa",
                    license_value="unknown-kaggle-dataset-card-required",
                    loss_weight=1.05,
                    record_kind="multilingual_context_qa",
                    extra_meta={"language": lang, "output_type": _norm(row.get("output_type")), "source_id": _norm(row.get("id"))},
                )
                if rec:
                    records.append(rec)
                    count += 1
                    if count >= per_lang or len(records) >= self.policy.max_multiloko:
                        break
            langs_seen[lang] = count
            if len(records) >= self.policy.max_multiloko:
                break
        return records, {"root": str(root), "records": len(records), "languages": langs_seen}

    def _mgsm(self) -> tuple[list[dict[str, Any]], dict[str, Any]]:
        root = self._root("mgsm", "open-benchmarks/mgsm-multilingual-grade-school-math-benchmark")
        records: list[dict[str, Any]] = []
        files = [("thai", root / "mgsm_th.tsv"), ("english", root / "mgsm_en.tsv")]
        per_file = max(1, self.policy.max_mgsm // len(files))
        for lang, path in files:
            if not path.exists():
                continue
            import pandas as pd

            df = pd.read_csv(path, sep="\t", header=None)
            for _, row in df.head(per_file).iterrows():
                question = _norm(row.iloc[0])
                answer = _norm(row.iloc[1])
                user = (
                    f"Solve this grade-school math problem. Show compact arithmetic reasoning, then give the final numeric answer.\n"
                    f"Language: {lang}\nProblem: {question}"
                )
                assistant = f"Reasoning: compute the quantities step by step from the problem statement.\nFinal answer: {answer}"
                rec = _record(
                    source="mgsm_kaggle_multilingual_math",
                    row={"language": lang, "question": question, "answer": answer},
                    user=user,
                    assistant=assistant,
                    dataset="open-benchmarks/mgsm-multilingual-grade-school-math-benchmark",
                    domain="thai_math" if lang == "thai" else "math",
                    license_value="unknown-kaggle-dataset-card-required",
                    loss_weight=1.22 if lang == "thai" else 1.08,
                    record_kind="multilingual_grade_school_math",
                    extra_meta={"language": lang},
                )
                if rec:
                    records.append(rec)
                if len(records) >= self.policy.max_mgsm:
                    return records, {"root": str(root), "records": len(records), "files": [str(p) for _, p in files]}
        return records, {"root": str(root), "records": len(records), "files": [str(p) for _, p in files]}

    def _livecodebench(self) -> tuple[list[dict[str, Any]], dict[str, Any]]:
        root = self._root("livecodebench", "open-benchmarks/livecodebench")
        path = root / "test6.jsonl"
        if not path.exists():
            candidates = sorted(root.glob("test*.jsonl"), key=lambda p: p.stat().st_size)
            if not candidates:
                return [], {"root": str(root), "records": 0, "missing": True}
            path = candidates[0]
        records: list[dict[str, Any]] = []
        for row in _jsonl_rows(path):
            title = _norm(row.get("question_title"))
            content = _norm(row.get("question_content"))
            public_tests = _norm(row.get("public_test_cases"))
            difficulty = _norm(row.get("difficulty"))
            user = (
                f"Write a Python solution for this programming problem. Use stdin/stdout and satisfy the public tests.\n"
                f"Title: {title}\nDifficulty: {difficulty}\nProblem:\n{content[:6000]}\nPublic tests: {public_tests[:1600]}"
            )
            assistant = (
                "Plan: parse stdin, derive the invariant from the statement, implement an efficient Python solution, "
                "and verify against the public tests. Do not assume hidden private-test data."
            )
            rec = _record(
                source="livecodebench_kaggle_code_prompting",
                row=row,
                user=user,
                assistant=assistant,
                dataset="open-benchmarks/livecodebench",
                domain="coding_python",
                license_value="unknown-kaggle-dataset-card-required",
                loss_weight=0.72,
                record_kind="livecodebench_prompt_strategy",
                extra_meta={
                    "question_id": _norm(row.get("question_id")),
                    "contest_id": _norm(row.get("contest_id")),
                    "difficulty": difficulty,
                    "main_training_allowed": True,
                    "solution_label_available": False,
                },
            )
            if rec:
                records.append(rec)
            if len(records) >= self.policy.max_livecodebench:
                break
        return records, {"root": str(root), "records": len(records), "file": str(path)}

    def write_jsonl(self, out_dir: str | Path) -> dict[str, Any]:
        out = Path(out_dir)
        out.mkdir(parents=True, exist_ok=True)
        groups = {
            "parsebench": self._parsebench(),
            "simpleqa": self._simpleqa(),
            "multiloko": self._multiloko(),
            "mgsm": self._mgsm(),
            "livecodebench": self._livecodebench(),
        }
        records: list[dict[str, Any]] = []
        fetch: dict[str, Any] = {}
        seen: set[str] = set()
        for name, (rows, report) in groups.items():
            fetch[name] = report
            for row in rows:
                key = row["metadata"]["fingerprint_sha256"]
                if key in seen:
                    continue
                seen.add(key)
                records.append(row)

        train_path = out / "kaggle_benchmark_mix_sft.jsonl"
        train_path.write_text("\n".join(json.dumps(row, ensure_ascii=False, sort_keys=True) for row in records), encoding="utf-8")
        manifest = {
            "schema_version": SCHEMA_VERSION,
            "created_at": _now(),
            "train_jsonl": str(train_path),
            "records_written": len(records),
            "source_counts": {name: report["records"] for name, (_, report) in groups.items()},
            "fetch": fetch,
            "policy": {
                "max_parsebench": self.policy.max_parsebench,
                "max_simpleqa": self.policy.max_simpleqa,
                "max_multiloko": self.policy.max_multiloko,
                "max_mgsm": self.policy.max_mgsm,
                "max_livecodebench": self.policy.max_livecodebench,
                "multiloko_languages": list(self.policy.multiloko_languages),
            },
            "claim_gate": {
                "main_training_allowed": len(records) > 0,
                "official_benchmark_claim_allowed": False,
                "world_best_claim_allowed": False,
                "reason": "These Kaggle datasets are benchmark-derived supplements; official claims require separate held-out submissions.",
            },
            "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 / "kaggle_benchmark_mix_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