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

"""Strict-provenance Claude Opus reasoning data registry and normalizer.

The main training output is clean SFT only: final answers in OpenAI chat JSONL.
Raw thinking/reasoning traces are written to a separate quarantine artifact and
are never included in the default Sandbox Model Bridge mix.
"""

from __future__ import annotations

from collections import Counter
from dataclasses import asdict, dataclass
from datetime import datetime, timezone
import hashlib
import json
import os
from pathlib import Path
import re
from typing import Any, Iterable, Iterator


SCHEMA_VERSION = "tinymind-claude-reasoning-strict-v2"
BUCKET_URI = "hf://buckets/bbkdevops/claude-opus-4.6-4.7-reasoning-8.7k-bucket"
ALLOWED_MODELS = {"claude-opus-4-6", "claude-opus-4-7"}

INSTRUCTIONAL_CATEGORIES = (
    "coding",
    "math",
    "physics",
    "biology",
    "chemistry",
    "earth_science",
    "science",
    "history",
    "philosophy",
    "psychology",
    "political_science",
    "sociology",
    "economics",
    "geography",
    "literature",
    "humanities",
    "arts",
    "finance",
    "medicine",
    "law",
    "business",
    "linguistics",
    "creative_writing",
    "general",
)

CREATIVE_ROLEPLAY_CATEGORIES = (
    "roleplay_hero",
    "roleplay_villain",
    "roleplay_crossover",
    "narrative_prose",
)

CATEGORY_COUNTS = {
    "coding": {"examples": 1628, "tokens": 2545221, "multi_turn_pct": 30.4},
    "humanities": {"examples": 862, "tokens": 1849708, "multi_turn_pct": 32.5},
    "science": {"examples": 737, "tokens": 1681346, "multi_turn_pct": 37.4},
    "roleplay_hero": {"examples": 419, "tokens": 640084, "multi_turn_pct": 63.5},
    "roleplay_villain": {"examples": 378, "tokens": 635984, "multi_turn_pct": 60.8},
    "narrative_prose": {"examples": 377, "tokens": 710807, "multi_turn_pct": 43.0},
    "roleplay_crossover": {"examples": 315, "tokens": 581188, "multi_turn_pct": 56.8},
    "creative_writing": {"examples": 281, "tokens": 532504, "multi_turn_pct": 30.6},
    "medicine": {"examples": 280, "tokens": 519662, "multi_turn_pct": 22.1},
    "biology": {"examples": 277, "tokens": 541013, "multi_turn_pct": 21.3},
    "general": {"examples": 276, "tokens": 284696, "multi_turn_pct": 37.0},
    "arts": {"examples": 245, "tokens": 576170, "multi_turn_pct": 41.2},
    "chemistry": {"examples": 221, "tokens": 508546, "multi_turn_pct": 52.9},
    "physics": {"examples": 220, "tokens": 512196, "multi_turn_pct": 56.8},
    "math": {"examples": 212, "tokens": 394907, "multi_turn_pct": 54.2},
    "geography": {"examples": 155, "tokens": 358321, "multi_turn_pct": 42.6},
    "history": {"examples": 155, "tokens": 348822, "multi_turn_pct": 41.3},
    "economics": {"examples": 155, "tokens": 380372, "multi_turn_pct": 42.6},
    "political_science": {"examples": 154, "tokens": 374901, "multi_turn_pct": 38.3},
    "sociology": {"examples": 154, "tokens": 378261, "multi_turn_pct": 42.2},
    "business": {"examples": 152, "tokens": 315065, "multi_turn_pct": 38.2},
    "earth_science": {"examples": 152, "tokens": 358209, "multi_turn_pct": 41.4},
    "finance": {"examples": 151, "tokens": 328607, "multi_turn_pct": 38.4},
    "philosophy": {"examples": 150, "tokens": 335514, "multi_turn_pct": 41.3},
    "linguistics": {"examples": 150, "tokens": 306889, "multi_turn_pct": 39.3},
    "literature": {"examples": 150, "tokens": 299606, "multi_turn_pct": 38.7},
    "psychology": {"examples": 150, "tokens": 339565, "multi_turn_pct": 39.3},
    "law": {"examples": 150, "tokens": 375360, "multi_turn_pct": 41.3},
}

MODEL_COUNTS = {
    "claude-opus-4-6": {"examples": 4675, "share": 53.7, "tokens": 6304169},
    "claude-opus-4-7": {"examples": 4031, "share": 46.3, "tokens": 10709363},
}

OVERALL_STATS = {
    "examples": 8706,
    "estimated_tokens": 17013533,
    "avg_tokens_per_example": 1954,
    "with_reasoning": 8706,
    "with_reasoning_pct": 100.0,
    "multi_turn": 3454,
    "single_turn": 5252,
}

REASONING_RE = re.compile(r"<(?:reasoning|think)>\s*.*?\s*</(?:reasoning|think)>\s*", re.IGNORECASE | re.DOTALL)
FINAL_RE = re.compile(r"\bFinal answer\s*:\s*", re.IGNORECASE)
SECRET_RE = re.compile(r"\b(?:hf_[A-Za-z0-9]{20,}|sk-[A-Za-z0-9_\-]{20,}|sk-or-v1-[A-Za-z0-9]{20,})\b")


@dataclass(frozen=True)
class ClaudeReasoningSource:
    repo_id: str
    split: str = "train"
    license: str = "apache-2.0"
    model_provenance: str = "claude-opus-4-7"
    schema_type: str = "auto"
    row_estimate: int = 0
    allowed_tier: str = "strict_main"
    notes: str = ""


STRICT_SOURCE_REGISTRY: tuple[ClaudeReasoningSource, ...] = (
    ClaudeReasoningSource(
        repo_id="angrygiraffe/claude-opus-4.6-4.7-reasoning-8.7k",
        model_provenance="claude-opus-4-6|claude-opus-4-7",
        schema_type="category_model_messages",
        row_estimate=8706,
        notes="Prefer *_no_reasoning JSONL files when using hf download.",
    ),
    ClaudeReasoningSource(
        repo_id="lordx64/reasoning-distill-claude-opus-4-7-max",
        model_provenance="claude-opus-4-7",
        schema_type="thinking_response_messages",
        row_estimate=8124,
    ),
    ClaudeReasoningSource(
        repo_id="lordx64/reasoning-distill-opus-4-7-max-sft",
        model_provenance="claude-opus-4-7",
        schema_type="sft_text_with_think",
        row_estimate=7823,
    ),
    ClaudeReasoningSource(
        repo_id="TeichAI/Claude-Opus-4.6-Reasoning-887x",
        model_provenance="claude-opus-4-6",
        schema_type="messages_thinking_content",
        row_estimate=887,
    ),
    ClaudeReasoningSource(
        repo_id="LEGENDQ/Claude-Opus-4.6-Reasoning-Dataset",
        model_provenance="claude-opus-4-6",
        schema_type="auto",
        row_estimate=2326,
        notes="Strict candidate, but schema must pass local inspection before training.",
    ),
    ClaudeReasoningSource(
        repo_id="Verdugie/opus-4.6-training-catalog",
        model_provenance="mixed-claude-opus-4-6",
        schema_type="registry_reference",
        row_estimate=168301,
        allowed_tier="registry_reference",
        notes="Catalog has mixed upstream licenses; never included in strict main output by default.",
    ),
)


@dataclass(frozen=True)
class ClaudeReasoningPolicy:
    include_creative_roleplay: bool = False
    keep_reasoning_blocks: bool = False
    max_records: int = 8706
    max_records_per_source: int = 6000
    source_registry: tuple[ClaudeReasoningSource, ...] = STRICT_SOURCE_REGISTRY


def _now() -> str:
    return datetime.now(timezone.utc).replace(microsecond=0).isoformat()


def _sha(value: str) -> str:
    return hashlib.sha256(value.encode("utf-8", errors="ignore")).hexdigest()


def _read_jsonl(path: Path) -> Iterable[dict[str, Any]]:
    decoder = json.JSONDecoder(strict=False)
    with path.open("r", encoding="utf-8", errors="replace") as f:
        for line in f:
            if line.strip():
                try:
                    yield decoder.decode(line)
                except json.JSONDecodeError:
                    continue


def _category_files(source_dir: Path) -> list[Path]:
    categories = source_dir / "categories"
    root = categories if categories.exists() else source_dir
    return sorted(root.rglob("*.jsonl"))


def _clean_text(text: str) -> tuple[str, bool]:
    cleaned = REASONING_RE.sub("", text).strip()
    stripped = cleaned != text.strip()
    if FINAL_RE.search(cleaned):
        cleaned = FINAL_RE.split(cleaned, maxsplit=1)[-1].strip()
        stripped = True
    return cleaned, stripped


def _reasoning_blocks(text: str) -> list[str]:
    return [match.group(0) for match in REASONING_RE.finditer(text)]


def _strip_reasoning(messages: list[dict[str, Any]]) -> tuple[list[dict[str, str]], list[dict[str, Any]], bool]:
    stripped = False
    quarantine: list[dict[str, Any]] = []
    out: list[dict[str, str]] = []
    for idx, message in enumerate(messages):
        role = str(message.get("role", "user"))
        content = str(message.get("content", ""))
        if message.get("thinking"):
            quarantine.append({"message_index": idx, "role": role, "thinking": str(message["thinking"])})
            stripped = True
        if role == "assistant":
            for block in _reasoning_blocks(content):
                quarantine.append({"message_index": idx, "role": role, "reasoning_block": block})
            content, removed = _clean_text(content)
            stripped = stripped or removed
        out.append({"role": role, "content": content})
    return out, quarantine, stripped


def _source_for_row(row: dict[str, Any], source_hint: str | None) -> ClaudeReasoningSource | None:
    if source_hint:
        for source in STRICT_SOURCE_REGISTRY:
            if source.repo_id == source_hint or source.repo_id.replace("/", "_").lower() in source_hint.lower():
                return source
    repo = str(row.get("source_repo") or row.get("repo_id") or row.get("dataset") or row.get("source_dataset") or "")
    for source in STRICT_SOURCE_REGISTRY:
        if repo == source.repo_id or source.repo_id in repo:
            return source
    return None


def _model_from_row(row: dict[str, Any], source: ClaudeReasoningSource | None) -> str:
    model = str(row.get("model") or row.get("source_model") or "").strip()
    if model in ALLOWED_MODELS:
        return model
    if source and source.model_provenance in ALLOWED_MODELS:
        return source.model_provenance
    return model or "unknown"


def _category(row: dict[str, Any], source: ClaudeReasoningSource | None) -> str:
    value = str(row.get("category") or row.get("domain") or "general").lower()
    value = re.sub(r"[^a-z0-9_]+", "_", value).strip("_")
    if value in INSTRUCTIONAL_CATEGORIES or value in CREATIVE_ROLEPLAY_CATEGORIES:
        return value
    if source and "coding" in source.repo_id.lower():
        return "coding"
    return "general"


def _license_ok(source: ClaudeReasoningSource | None) -> bool:
    return bool(source and source.license.lower() in {"apache-2.0", "mit"})


def _provenance_ok(model: str, source: ClaudeReasoningSource | None) -> bool:
    if model in ALLOWED_MODELS:
        return True
    return bool(source and all(part in ALLOWED_MODELS for part in source.model_provenance.split("|")))


def _loss_weight(category: str) -> float:
    if category in {"coding", "math", "physics", "chemistry", "medicine", "law"}:
        return 1.12
    if category in CREATIVE_ROLEPLAY_CATEGORIES:
        return 0.35
    if category == "creative_writing":
        return 0.75
    return 1.0


def _messages_from_raw(row: dict[str, Any]) -> tuple[list[dict[str, Any]] | None, list[dict[str, Any]]]:
    if row.get("thinking") is not None and row.get("response") is not None:
        messages = list(row.get("messages") or [])
        if not messages and row.get("question"):
            messages = [{"role": "user", "content": str(row["question"])}]
        if not messages and row.get("prompt"):
            messages = [{"role": "user", "content": str(row["prompt"])}]
        messages.append({"role": "assistant", "content": str(row.get("response", ""))})
        return messages, [{"role": "assistant", "thinking": str(row.get("thinking", ""))}]

    if isinstance(row.get("messages"), list):
        return row["messages"], []

    if row.get("text") is not None:
        clean, removed = _clean_text(str(row["text"]))
        quarantine = [{"role": "assistant", "thinking": str(row["text"])}] if removed else []
        return [{"role": "user", "content": "Provide the final answer for this SFT item."}, {"role": "assistant", "content": clean}], quarantine

    if row.get("question") is not None and row.get("answer") is not None:
        quarantine = [{"role": "assistant", "thinking": str(row.get("thought", ""))}] if row.get("thought") else []
        return [{"role": "user", "content": str(row["question"])}, {"role": "assistant", "content": str(row["answer"])}], quarantine

    return None, []


def normalize_row(
    row: dict[str, Any],
    policy: ClaudeReasoningPolicy,
    *,
    source_hint: str | None = None,
) -> tuple[dict[str, Any] | None, dict[str, Any] | None, str | None]:
    source = _source_for_row(row, source_hint)
    if source and source.allowed_tier != "strict_main":
        return None, None, "source_not_main_tier"

    category = _category(row, source)
    if category in CREATIVE_ROLEPLAY_CATEGORIES and not policy.include_creative_roleplay:
        return None, None, "creative_roleplay_excluded"

    model = _model_from_row(row, source)
    if not _provenance_ok(model, source):
        return None, None, "missing_strict_model_provenance"
    if not _license_ok(source):
        return None, None, "license_review_failed"

    messages, raw_quarantine = _messages_from_raw(row)
    if not messages:
        return None, None, "unsupported_schema"

    if policy.keep_reasoning_blocks:
        normalized_messages = [
            {"role": str(message.get("role", "user")), "content": str(message.get("content", ""))}
            for message in messages
        ]
        thinking_quarantine = raw_quarantine
        reasoning_stripped = False
    else:
        normalized_messages, thinking_quarantine, reasoning_stripped = _strip_reasoning(messages)
        thinking_quarantine.extend(raw_quarantine)

    text = "\n".join(message["content"] for message in normalized_messages)
    if SECRET_RE.search(text):
        return None, None, "secret_like_token"
    if not text.strip():
        return None, None, "empty_after_strip"

    metadata = {
        "domain": f"claude_reasoning_{category}",
        "category": category,
        "source_model": model,
        "source_repo": source.repo_id if source else source_hint,
        "source_dataset": row.get("source_dataset"),
        "source_idx": row.get("source_idx"),
        "license": source.license if source else "unknown",
        "bucket_or_repo": source.repo_id if source else source_hint,
        "schema_type": source.schema_type if source else "auto",
        "reasoning_blocks_stripped": reasoning_stripped or bool(thinking_quarantine),
        "strict_provenance_passed": True,
        "license_review_passed": True,
        "loss_weight": _loss_weight(category),
        "fingerprint_sha256": _sha(json.dumps(normalized_messages, ensure_ascii=False, sort_keys=True)),
    }
    normalized = {
        "messages": normalized_messages,
        "source": "claude_reasoning_bucket",
        "category": category,
        "model": model,
        "metadata": metadata,
    }
    quarantine = None
    if thinking_quarantine:
        quarantine = {
            "source": "claude_reasoning_bucket_quarantine",
            "metadata": metadata | {"raw_trace_quarantine_only": True},
            "raw_trace": thinking_quarantine,
        }
    return normalized, quarantine, None


def iter_hf_rows(source: ClaudeReasoningSource) -> Iterator[dict[str, Any]]:
    """Stream rows from Hugging Face with env-only auth.

    Public datasets may not technically require a token, but TinyMind requires
    HF_TOKEN for network ingestion so scripts never pass credentials on the CLI.
    """

    if not os.environ.get("HF_TOKEN"):
        raise RuntimeError("HF_TOKEN env var is required for HF network ingestion; do not pass tokens on the command line.")
    from datasets import load_dataset

    ds = load_dataset(source.repo_id, split=source.split, token=os.environ["HF_TOKEN"], streaming=True)
    for row in ds:
        row = dict(row)
        row.setdefault("source_repo", source.repo_id)
        yield row


def build_static_profile(out_dir: str | Path) -> dict[str, Any]:
    out = Path(out_dir)
    out.mkdir(parents=True, exist_ok=True)
    manifest = _base_manifest()
    manifest["outputs"] = {}
    manifest["claim_gate"].update(
        {
            "local_rows_normalized": False,
            "training_ready": False,
            "main_training_allowed": False,
            "reason": "Static profile only. Run with a downloaded local source directory or network ingestion to normalize rows.",
        }
    )
    path = out / "claude_reasoning_bucket_manifest.json"
    manifest["manifest_path"] = str(path)
    path.write_text(json.dumps(manifest, ensure_ascii=False, indent=2, sort_keys=True), encoding="utf-8")
    return manifest


def build_claude_reasoning_dataset(
    out_dir: str | Path,
    source_dir: str | Path | None = None,
    policy: ClaudeReasoningPolicy | None = None,
    *,
    ingest_hf: bool = False,
    source_repo: list[str] | None = None,
) -> dict[str, Any]:
    policy = policy or ClaudeReasoningPolicy()
    out = Path(out_dir)
    out.mkdir(parents=True, exist_ok=True)

    if source_dir is None and not ingest_hf:
        return build_static_profile(out)

    train_path = out / "claude_reasoning_train.jsonl"
    quarantine_path = out / "claude_reasoning_quarantine.jsonl"
    category_counts: Counter[str] = Counter()
    model_counts: Counter[str] = Counter()
    source_counts: Counter[str] = Counter()
    rejected: Counter[str] = Counter()
    kept = 0
    quarantined = 0

    with train_path.open("w", encoding="utf-8", newline="\n") as train_f, quarantine_path.open("w", encoding="utf-8", newline="\n") as quarantine_f:
        for row, hint in _iter_candidate_rows(source_dir, policy, ingest_hf=ingest_hf, source_repo=source_repo):
            if kept >= policy.max_records:
                rejected["global_cap"] += 1
                continue
            normalized, quarantine, reason = normalize_row(row, policy, source_hint=hint)
            if reason:
                rejected[reason] += 1
                continue
            if normalized is None:
                rejected["unknown_reject"] += 1
                continue
            source_key = str(normalized["metadata"].get("source_repo"))
            if source_counts[source_key] >= policy.max_records_per_source:
                rejected["source_cap"] += 1
                continue
            train_f.write(json.dumps(normalized, ensure_ascii=False) + "\n")
            kept += 1
            category_counts[normalized["category"]] += 1
            model_counts[normalized["model"]] += 1
            source_counts[source_key] += 1
            if quarantine is not None:
                quarantine_f.write(json.dumps(quarantine, ensure_ascii=False) + "\n")
                quarantined += 1

    manifest = _base_manifest()
    manifest.update(
        {
            "source_dir": str(source_dir) if source_dir else None,
            "ingest_hf": ingest_hf,
            "outputs": {"train_jsonl": str(train_path), "quarantine_jsonl": str(quarantine_path)},
            "summary": {
                "records_written": kept,
                "quarantine_records": quarantined,
                "category_counts": dict(sorted(category_counts.items())),
                "model_counts": dict(sorted(model_counts.items())),
                "source_counts": dict(sorted(source_counts.items())),
                "rejected": dict(sorted(rejected.items())),
            },
            "policy": {
                "include_creative_roleplay": policy.include_creative_roleplay,
                "keep_reasoning_blocks": policy.keep_reasoning_blocks,
                "max_records": policy.max_records,
                "max_records_per_source": policy.max_records_per_source,
            },
        }
    )
    manifest["claim_gate"].update(
        {
            "local_rows_normalized": kept > 0,
            "strict_provenance_passed": kept > 0 and not any(k in rejected for k in ("missing_strict_model_provenance", "license_review_failed")),
            "reasoning_trace_stripped": kept > 0 and not policy.keep_reasoning_blocks,
            "license_review_passed": kept > 0 and rejected.get("license_review_failed", 0) == 0,
            "main_training_allowed": kept > 0 and not policy.keep_reasoning_blocks,
            "training_ready": kept > 0 and not policy.keep_reasoning_blocks,
            "raw_trace_quarantine_only": quarantined >= 0,
            "reasoning_trace_training_allowed": policy.keep_reasoning_blocks,
            "creative_roleplay_included": policy.include_creative_roleplay,
        }
    )
    path = out / "claude_reasoning_bucket_manifest.json"
    manifest["manifest_path"] = str(path)
    path.write_text(json.dumps(manifest, ensure_ascii=False, indent=2, sort_keys=True), encoding="utf-8")
    return manifest


def _iter_candidate_rows(
    source_dir: str | Path | None,
    policy: ClaudeReasoningPolicy,
    *,
    ingest_hf: bool,
    source_repo: list[str] | None,
) -> Iterator[tuple[dict[str, Any], str | None]]:
    if source_dir is not None:
        source = Path(source_dir)
        for path in _category_files(source):
            hint = _hint_from_path(path)
            for row in _read_jsonl(path):
                row.setdefault("source_repo", hint)
                yield row, hint
    if ingest_hf:
        selected = [
            source
            for source in policy.source_registry
            if source.allowed_tier == "strict_main" and (not source_repo or source.repo_id in source_repo)
        ]
        for source in selected:
            for row in iter_hf_rows(source):
                yield row, source.repo_id


def _hint_from_path(path: Path) -> str | None:
    parts = [part.lower() for part in path.parts]
    joined = "/".join(parts)
    for source in STRICT_SOURCE_REGISTRY:
        repo_key = source.repo_id.replace("/", "_").lower()
        if repo_key in joined or source.repo_id.lower() in joined:
            return source.repo_id
    if "categories" in parts or path.name in {"full_train.jsonl", "instruct_train.jsonl", "code_train.jsonl"}:
        return "angrygiraffe/claude-opus-4.6-4.7-reasoning-8.7k"
    return None


def _base_manifest() -> dict[str, Any]:
    return {
        "schema_version": SCHEMA_VERSION,
        "created_at": _now(),
        "bucket_uri": BUCKET_URI,
        "source_registry": [asdict(source) for source in STRICT_SOURCE_REGISTRY],
        "overall": OVERALL_STATS,
        "category_sets": {
            "instructional": list(INSTRUCTIONAL_CATEGORIES),
            "creative_roleplay": list(CREATIVE_ROLEPLAY_CATEGORIES),
        },
        "category_counts_from_card": CATEGORY_COUNTS,
        "model_counts_from_card": MODEL_COUNTS,
        "format": {
            "chat_format": "OpenAI JSONL messages",
            "metadata_fields": ["category", "model", "metadata.source_repo", "metadata.license"],
            "fine_tuning_reads_messages_only": True,
        },
        "safety_policy": {
            "default_includes_creative_roleplay": False,
            "default_keeps_reasoning_blocks": False,
            "reason": "Avoid overfitting to hidden reasoning traces and named-roleplay style imitation unless explicitly enabled.",
        },
        "claim_gate": {
            "external_source_profile_ready": True,
            "world_best_claim_allowed": False,
        },
    }