tinymind-native-colab-handoff/bundle/train/dataset_quality_governor.py

from __future__ import annotations

from collections import Counter, defaultdict
from dataclasses import dataclass, field
from datetime import datetime, timezone
import hashlib
import json
from pathlib import Path
import re
from typing import Any


ROOT = Path(__file__).resolve().parents[1]
SECRET_PATTERNS = [
    re.compile(r"\bhf_[A-Za-z0-9]{20,}\b"),
    re.compile(r"\bsk-or-v1-[A-Za-z0-9]{20,}\b"),
    re.compile(r"\b(?:api[_-]?key|access[_-]?token|secret)\s*[:=]\s*['\"]?[A-Za-z0-9_\-]{16,}", re.IGNORECASE),
]
BASE64ISH_RE = re.compile(r"^[A-Za-z0-9+/=\s]{512,}$")
SYMBOL_RE = re.compile(r"[^\w\s\u0E00-\u0E7F.,:;!?()\[\]{}<>/\\'\"`+=#@&%*\-]")
RAW_REASONING_RE = re.compile(
    r"<\s*(?:think|thinking|reasoning)\b[^>]*>.*?<\s*/\s*(?:think|thinking|reasoning)\s*>",
    re.IGNORECASE | re.DOTALL,
)


def _canonical_text(item: dict[str, Any]) -> str:
    if "messages" in item:
        return "\n".join(str(m.get("role", "")) + ":" + str(m.get("content", "")) for m in item.get("messages", []))
    return str(item.get("text", ""))


def _norm_for_hash(text: str) -> str:
    text = re.sub(r"\s+", " ", text).strip().lower()
    text = re.sub(r"https?://\S+", "<url>", text)
    return text


def _source(item: dict[str, Any]) -> str:
    if item.get("source"):
        return str(item["source"])
    meta = item.get("metadata") or {}
    if meta.get("record_kind"):
        return "sandbox_" + str(meta["record_kind"])
    return "unknown"


def _domain(item: dict[str, Any], source: str) -> str:
    low = source.lower()
    if "alignment_tool_sft" in low:
        meta_domain = str((item.get("metadata") or {}).get("domain", "")).lower()
        if meta_domain in {"alignment_constraint_following", "alignment_tool_calling"}:
            return meta_domain
        return "alignment_tool_sft"
    if "logic_agent_code" in low:
        meta_domain = str((item.get("metadata") or {}).get("domain", "")).lower()
        if meta_domain in {
            "instruction_following",
            "tool_grounding",
            "coding_python",
            "coding_cpp_rust",
            "reasoning_logic",
            "data_tooling",
        }:
            return "logic_" + meta_domain
        return "logic_agent_code"
    if "claude_reasoning_bucket" in low:
        category = str((item.get("metadata") or {}).get("category") or item.get("category") or "general").lower()
        return "claude_reasoning_" + re.sub(r"[^a-z0-9_]+", "_", category).strip("_")
    if "parsebench" in low:
        return "data_tooling"
    if "simpleqa" in low:
        return "general"
    if "multiloko" in low:
        return "thai_grounding"
    if "mgsm" in low:
        return "thai_grounding"
    if "livecodebench" in low:
        return "logic_coding_python"
    if "coverage_100k" in low or "coverage_100_axis" in low:
        return "coverage_100k"
    if "cve" in low:
        return "cve_intelligence"
    if "thai" in low or "mt_opus" in low:
        return "thai_grounding"
    if "sandbox" in low or "tool" in low:
        return "sandbox_tools"
    if "reverse" in low or "apk" in low or "ghidra" in low or "il2cpp" in low or "droid" in low:
        return "reverse_engineering"
    if "maxmind" in low:
        return "data_tooling"
    return "general"


def _semantic_hash(item: dict[str, Any], text: str, domain: str) -> str:
    metadata = item.get("metadata") or {}
    if domain == "coverage_100k" and metadata.get("fingerprint_sha256"):
        return "coverage_100k:" + str(metadata["fingerprint_sha256"])
    if domain.startswith("logic_") and metadata.get("fingerprint_sha256"):
        return domain + ":" + str(metadata["fingerprint_sha256"])
    if domain.startswith("alignment_") and metadata.get("fingerprint_sha256"):
        return domain + ":" + str(metadata["fingerprint_sha256"])
    return hashlib.sha256(_norm_for_hash(text).encode("utf-8", errors="ignore")).hexdigest()


def _source_shard(source: str, domain: str) -> str:
    if domain.startswith(("logic_", "alignment_", "claude_reasoning_")):
        return f"{source}:{domain}"
    return source


def _token_estimate(text: str) -> int:
    # Conservative multilingual approximation. Used only for filtering/reporting, not claims.
    return max(1, len(text) // 4)


def _quality_reject_reason(text: str, *, skip_repetition_scan: bool = False) -> str | None:
    compact = re.sub(r"\s+", " ", text).strip()
    if RAW_REASONING_RE.search(compact):
        return "raw_reasoning_trace"
    if any(pattern.search(compact) for pattern in SECRET_PATTERNS):
        return "secret_like_token"

    alnum = re.sub(r"[^A-Za-z0-9+/=]", "", text)
    unique_chars = len(set(alnum))
    if re.search(r"([A-Za-z0-9+/=])\1{127,}", text):
        return "encoded_blob"
    if len(alnum) >= 512 and unique_chars <= 8:
        return "encoded_blob"
    if BASE64ISH_RE.fullmatch(text.strip()) and len(alnum) >= 512 and len(alnum) / max(1, len(text)) > 0.92:
        return "encoded_blob"

    if not skip_repetition_scan:
        words = re.findall(r"[\w\u0E00-\u0E7F]{2,}", compact.lower())
        if len(words) >= 80:
            most_common_count = Counter(words).most_common(1)[0][1]
            if most_common_count / len(words) > 0.35:
                return "repetition_loop"
        for n in (3, 4, 5):
            if len(words) >= n * 16:
                grams = [" ".join(words[i : i + n]) for i in range(len(words) - n + 1)]
                if Counter(grams).most_common(1)[0][1] >= 12:
                    return "repetition_loop"

    visible = [ch for ch in text if not ch.isspace()]
    if len(visible) >= 80:
        symbol_count = sum(1 for ch in visible if SYMBOL_RE.match(ch))
        punctuation_count = sum(1 for ch in visible if not (ch.isalnum() or "\u0E00" <= ch <= "\u0E7F"))
        wordlike_count = sum(1 for ch in visible if ch.isalnum() or "\u0E00" <= ch <= "\u0E7F")
        if symbol_count / len(visible) > 0.45 and wordlike_count / len(visible) < 0.35:
            return "symbol_noise"
        if punctuation_count / len(visible) > 0.70 and wordlike_count / len(visible) < 0.20:
            return "symbol_noise"
    return None


def _read_jsonl(path: Path):
    decoder = json.JSONDecoder(strict=False)
    with path.open("r", encoding="utf-8", errors="replace") as f:
        for line_no, line in enumerate(f, start=1):
            if not line.strip():
                continue
            try:
                yield line_no, decoder.decode(line)
            except json.JSONDecodeError:
                yield line_no, None


@dataclass(frozen=True)
class DatasetQualityPolicy:
    max_records: int = 24_000
    max_estimated_tokens: int = 2048
    min_estimated_tokens: int = 8
    recipe_profile: str = "default"
    domain_caps: dict[str, int] = field(
        default_factory=lambda: {
            "cve_intelligence": 8_000,
            "thai_grounding": 5_000,
            "reverse_engineering": 7_000,
            "sandbox_tools": 1_000,
            "data_tooling": 1_000,
            "coverage_100k": 100_000,
            "logic_instruction_following": 12_000,
            "logic_tool_grounding": 12_000,
            "logic_coding_python": 12_000,
            "logic_coding_cpp_rust": 12_000,
            "logic_reasoning_logic": 12_000,
            "logic_data_tooling": 12_000,
            "alignment_constraint_following": 15_000,
            "alignment_tool_calling": 15_000,
            "claude_reasoning_coding": 2_000,
            "claude_reasoning_humanities": 1_000,
            "claude_reasoning_science": 1_000,
            "claude_reasoning_general": 800,
            "general": 2_000,
        }
    )

    @staticmethod
    def for_profile(profile: str, *, max_records: int, max_estimated_tokens: int) -> "DatasetQualityPolicy":
        profiles = {
            "default": {
                "cve_intelligence": 8_000,
                "thai_grounding": 5_000,
                "reverse_engineering": 7_000,
                "sandbox_tools": 1_000,
                "data_tooling": 1_000,
                "coverage_100k": 100_000,
                "logic_instruction_following": 12_000,
                "logic_tool_grounding": 12_000,
                "logic_coding_python": 12_000,
                "logic_coding_cpp_rust": 12_000,
                "logic_reasoning_logic": 12_000,
                "logic_data_tooling": 12_000,
                "alignment_constraint_following": 15_000,
                "alignment_tool_calling": 15_000,
                "claude_reasoning_coding": 2_000,
                "claude_reasoning_humanities": 1_000,
                "claude_reasoning_science": 1_000,
                "claude_reasoning_general": 800,
                "general": 2_000,
            },
            "balanced": {
                "coverage_100k": 42_000,
                "thai_grounding": 16_000,
                "cve_intelligence": 14_000,
                "reverse_engineering": 14_000,
                "sandbox_tools": 8_000,
                "data_tooling": 4_000,
                "general": 12_000,
                "logic_instruction_following": 0,
                "logic_tool_grounding": 0,
                "logic_coding_python": 0,
                "logic_coding_cpp_rust": 0,
                "logic_reasoning_logic": 0,
                "logic_data_tooling": 0,
                "alignment_constraint_following": 0,
                "alignment_tool_calling": 0,
                "claude_reasoning_coding": 0,
                "claude_reasoning_humanities": 0,
                "claude_reasoning_science": 0,
                "claude_reasoning_general": 0,
            },
            "frontier": {
                "coverage_100k": 36_000,
                "thai_grounding": 18_000,
                "cve_intelligence": 12_000,
                "reverse_engineering": 12_000,
                "sandbox_tools": 10_000,
                "data_tooling": 4_000,
                "general": 16_000,
                "logic_instruction_following": 0,
                "logic_tool_grounding": 0,
                "logic_coding_python": 0,
                "logic_coding_cpp_rust": 0,
                "logic_reasoning_logic": 0,
                "logic_data_tooling": 0,
                "alignment_constraint_following": 0,
                "alignment_tool_calling": 0,
                "claude_reasoning_coding": 0,
                "claude_reasoning_humanities": 0,
                "claude_reasoning_science": 0,
                "claude_reasoning_general": 0,
            },
            "surgery": {
                "coverage_100k": 5_000,
                "logic_instruction_following": 9_000,
                "logic_tool_grounding": 9_000,
                "logic_coding_python": 9_000,
                "logic_coding_cpp_rust": 9_000,
                "logic_reasoning_logic": 9_000,
                "logic_data_tooling": 9_000,
                "alignment_constraint_following": 14_000,
                "alignment_tool_calling": 14_000,
                "claude_reasoning_coding": 1_500,
                "claude_reasoning_humanities": 900,
                "claude_reasoning_science": 900,
                "claude_reasoning_math": 600,
                "claude_reasoning_physics": 600,
                "claude_reasoning_biology": 600,
                "claude_reasoning_chemistry": 600,
                "claude_reasoning_medicine": 600,
                "claude_reasoning_law": 600,
                "claude_reasoning_business": 500,
                "claude_reasoning_finance": 500,
                "claude_reasoning_linguistics": 500,
                "claude_reasoning_history": 500,
                "claude_reasoning_philosophy": 500,
                "claude_reasoning_psychology": 500,
                "claude_reasoning_economics": 500,
                "claude_reasoning_political_science": 500,
                "claude_reasoning_sociology": 500,
                "claude_reasoning_geography": 500,
                "claude_reasoning_literature": 500,
                "claude_reasoning_arts": 500,
                "claude_reasoning_earth_science": 500,
                "claude_reasoning_creative_writing": 300,
                "claude_reasoning_general": 800,
                "thai_grounding": 8_000,
                "cve_intelligence": 6_000,
                "reverse_engineering": 6_000,
                "sandbox_tools": 4_000,
                "data_tooling": 2_000,
                "general": 12_000,
            },
            "apex": {
                "coverage_100k": 3_000,
                "logic_instruction_following": 18_000,
                "logic_tool_grounding": 18_000,
                "logic_coding_python": 14_000,
                "logic_coding_cpp_rust": 14_000,
                "logic_reasoning_logic": 14_000,
                "logic_data_tooling": 10_000,
                "alignment_constraint_following": 20_000,
                "alignment_tool_calling": 20_000,
                "claude_reasoning_coding": 1_200,
                "claude_reasoning_humanities": 700,
                "claude_reasoning_science": 700,
                "claude_reasoning_math": 500,
                "claude_reasoning_physics": 500,
                "claude_reasoning_biology": 500,
                "claude_reasoning_chemistry": 500,
                "claude_reasoning_medicine": 500,
                "claude_reasoning_law": 500,
                "claude_reasoning_business": 350,
                "claude_reasoning_finance": 350,
                "claude_reasoning_linguistics": 350,
                "claude_reasoning_history": 350,
                "claude_reasoning_philosophy": 350,
                "claude_reasoning_psychology": 350,
                "claude_reasoning_economics": 350,
                "claude_reasoning_political_science": 350,
                "claude_reasoning_sociology": 350,
                "claude_reasoning_geography": 350,
                "claude_reasoning_literature": 350,
                "claude_reasoning_arts": 250,
                "claude_reasoning_earth_science": 350,
                "claude_reasoning_creative_writing": 0,
                "thai_grounding": 10_000,
                "cve_intelligence": 4_000,
                "reverse_engineering": 5_000,
                "sandbox_tools": 6_000,
                "data_tooling": 4_000,
                "general": 8_000,
            },
        }
        if profile not in profiles:
            raise ValueError(f"unknown recipe_profile {profile!r}; choose one of {sorted(profiles)}")
        return DatasetQualityPolicy(
            max_records=max_records,
            max_estimated_tokens=max_estimated_tokens,
            recipe_profile=profile,
            domain_caps=profiles[profile],
        )


class DatasetQualityGovernor:
    def __init__(self, policy: DatasetQualityPolicy | None = None):
        self.policy = policy or DatasetQualityPolicy()

    def build(self, input_jsonl: str | Path, out_dir: str | Path) -> dict[str, Any]:
        src = Path(input_jsonl)
        out = Path(out_dir)
        out.mkdir(parents=True, exist_ok=True)
        optimized = out / "tinymind_12b_quality_governed_mix.jsonl"

        seen: set[str] = set()
        kept_records = 0
        input_records_seen = 0
        domain_counts: Counter[str] = Counter()
        source_counts: Counter[str] = Counter()
        source_shard_counts: Counter[str] = Counter()
        reject_counts: Counter[str] = Counter()
        rejected_examples: dict[str, list[dict[str, Any]]] = defaultdict(list)

        with optimized.open("w", encoding="utf-8", newline="\n") as f:
            for line_no, item in _read_jsonl(src):
                input_records_seen += 1
                if item is None:
                    self._reject(reject_counts, rejected_examples, "invalid_json", line_no, None)
                    continue
                text = _canonical_text(item)
                tok = _token_estimate(text)
                source = _source(item)
                domain = _domain(item, source)
                h = _semantic_hash(item, text, domain)
                if h in seen:
                    self._reject(reject_counts, rejected_examples, "duplicate_semantic_hash", line_no, source)
                    continue
                if tok < self.policy.min_estimated_tokens:
                    self._reject(reject_counts, rejected_examples, "too_short", line_no, source)
                    continue
                if tok > self.policy.max_estimated_tokens:
                    self._reject(reject_counts, rejected_examples, "too_long", line_no, source)
                    continue
                quality_reason = _quality_reject_reason(
                    text,
                    skip_repetition_scan=domain == "coverage_100k" or domain.startswith("logic_") or domain.startswith("alignment_"),
                )
                if quality_reason:
                    self._reject(reject_counts, rejected_examples, quality_reason, line_no, source)
                    continue
                if domain_counts[domain] >= self.policy.domain_caps.get(domain, self.policy.domain_caps["general"]):
                    self._reject(reject_counts, rejected_examples, "domain_cap", line_no, source)
                    continue
                if kept_records >= self.policy.max_records:
                    self._reject(reject_counts, rejected_examples, "global_cap", line_no, source)
                    continue

                item.setdefault("quality_governor", {})
                item["quality_governor"].update(
                    {
                        "domain": domain,
                        "source": source,
                        "estimated_tokens": tok,
                        "semantic_sha256": h,
                        "loss_weight": self._loss_weight(domain, item),
                    }
                )
                f.write(json.dumps(item, ensure_ascii=False) + "\n")
                seen.add(h)
                kept_records += 1
                domain_counts[domain] += 1
                source_counts[source] += 1
                source_shard_counts[_source_shard(source, domain)] += 1

        purity_intensity_gate = self._purity_intensity_gate(
            domain_counts=domain_counts,
            source_counts=source_shard_counts,
            reject_counts=reject_counts,
            kept_records=kept_records,
        )
        balanced_ready = self._balanced_ready(domain_counts, kept_records)
        train_allowed = balanced_ready and (
            self.policy.recipe_profile != "apex" or purity_intensity_gate["training_intensity_ready"]
        )
        manifest = {
            "schema_version": "tinymind-dataset-quality-governor-v1",
            "created_at": datetime.now(timezone.utc).isoformat(),
            "input_jsonl": str(src),
            "optimized_jsonl": str(optimized),
            "input_records_seen": input_records_seen,
            "kept_records": kept_records,
            "rejected_records": sum(reject_counts.values()),
            "domain_counts": dict(sorted(domain_counts.items())),
            "source_counts_top": dict(source_counts.most_common(40)),
            "source_shard_counts_top": dict(source_shard_counts.most_common(40)),
            "reject_counts": dict(sorted(reject_counts.items())),
            "rejected_examples": rejected_examples,
            "policy": {
                "max_records": self.policy.max_records,
                "max_estimated_tokens": self.policy.max_estimated_tokens,
                "min_estimated_tokens": self.policy.min_estimated_tokens,
                "recipe_profile": self.policy.recipe_profile,
                "domain_caps": self.policy.domain_caps,
            },
            "purity_intensity_gate": purity_intensity_gate,
            "training_contract": self._training_contract(domain_counts, kept_records),
            "claim_gate": {
                "quality_governed_dataset_ready": kept_records > 0,
                "balanced_recipe_ready": balanced_ready,
                "train_allowed": train_allowed,
                "zero_waste_claim_allowed": False,
                "reason": "The governor reduces known waste patterns and records evidence, but cannot prove absolute zero waste.",
            },
        }
        manifest_path = out / "dataset_quality_governor_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

    def _reject(
        self,
        counts: Counter[str],
        examples: dict[str, list[dict[str, Any]]],
        reason: str,
        line_no: int,
        source: str | None,
    ) -> None:
        counts[reason] += 1
        if len(examples[reason]) < 8:
            examples[reason].append({"line": line_no, "source": source})

    def _balanced_ready(self, domain_counts: Counter[str], kept_records: int) -> bool:
        if self.policy.recipe_profile == "default" or kept_records <= 0:
            return False
        dominant_share = max(domain_counts.values(), default=0) / kept_records
        coverage_share = domain_counts.get("coverage_100k", 0) / kept_records
        logic_share = sum(count for domain, count in domain_counts.items() if domain.startswith("logic_")) / kept_records
        alignment_share = sum(count for domain, count in domain_counts.items() if domain.startswith("alignment_")) / kept_records
        if self.policy.recipe_profile == "surgery":
            return coverage_share <= 0.10 and (logic_share + alignment_share) >= 0.55 and alignment_share >= 0.20 and dominant_share <= 0.35
        if self.policy.recipe_profile == "apex":
            return coverage_share <= 0.05 and (logic_share + alignment_share) >= 0.60 and alignment_share >= 0.20 and dominant_share <= 0.30
        return dominant_share <= 0.50

    def _loss_weight(self, domain: str, item: dict[str, Any]) -> float:
        metadata = item.get("metadata") or {}
        if metadata.get("loss_weight") is not None:
            try:
                return float(metadata["loss_weight"])
            except (TypeError, ValueError):
                pass
        if domain == "coverage_100k":
            if self.policy.recipe_profile == "apex":
                return 0.08
            return 0.15
        if domain.startswith("logic_"):
            if self.policy.recipe_profile == "apex":
                return 1.35
            return 1.25
        if domain.startswith("alignment_"):
            if self.policy.recipe_profile == "apex":
                return 1.60
            return 1.45
        if domain.startswith("claude_reasoning_"):
            metadata = item.get("metadata") or {}
            if metadata.get("reasoning_blocks_stripped") is False:
                return 0.55
            return float(metadata.get("loss_weight", 0.9))
        if domain in {"sandbox_tools", "data_tooling"}:
            if self.policy.recipe_profile == "apex":
                return 1.35
            return 1.2
        return 1.0

    def _purity_intensity_gate(
        self,
        *,
        domain_counts: Counter[str],
        source_counts: Counter[str],
        reject_counts: Counter[str],
        kept_records: int,
    ) -> dict[str, Any]:
        coverage_share = domain_counts.get("coverage_100k", 0) / max(kept_records, 1)
        logic_share = sum(count for domain, count in domain_counts.items() if domain.startswith("logic_")) / max(kept_records, 1)
        alignment_share = sum(count for domain, count in domain_counts.items() if domain.startswith("alignment_")) / max(kept_records, 1)
        claude_share = sum(count for domain, count in domain_counts.items() if domain.startswith("claude_reasoning_")) / max(kept_records, 1)
        dominant_source_share = max(source_counts.values(), default=0) / max(kept_records, 1)
        reasoning_trace_free = reject_counts.get("raw_reasoning_trace", 0) >= 0
        source_dominance_passed = dominant_source_share <= (0.35 if self.policy.recipe_profile == "apex" else 0.50)
        training_intensity_ready = self._balanced_ready(domain_counts, kept_records) and source_dominance_passed and reasoning_trace_free
        return {
            "profile": self.policy.recipe_profile,
            "coverage_share": coverage_share,
            "logic_share": logic_share,
            "alignment_share": alignment_share,
            "claude_reasoning_share": claude_share,
            "dominant_source_share": dominant_source_share,
            "reasoning_trace_free": reasoning_trace_free,
            "source_dominance_passed": source_dominance_passed,
            "training_intensity_ready": training_intensity_ready,
            "raw_reasoning_rejected": reject_counts.get("raw_reasoning_trace", 0),
            "world_purest_dataset_claim_allowed": False,
        }

    def _training_contract(self, domain_counts: Counter[str], kept_records: int) -> dict[str, Any]:
        return {
            "loss_normalization": "per_sample_token_normalized",
            "assistant_targeting": "final_answer_only_text_field",
            "reasoning_trace_policy": "strip_or_reject_raw_traces_before_main_training",
            "domain_loss_weights": {
                "coverage_100k": self._loss_weight("coverage_100k", {}),
                "logic": self._loss_weight("logic_instruction_following", {}),
                "alignment": self._loss_weight("alignment_tool_calling", {}),
                "sandbox_tools": self._loss_weight("sandbox_tools", {}),
                "data_tooling": self._loss_weight("data_tooling", {}),
                "claude_reasoning_clean": self._loss_weight("claude_reasoning_coding", {"metadata": {"reasoning_blocks_stripped": True}}),
            },
            "curriculum": [
                {
                    "phase": "alignment_core",
                    "focus": ["alignment_constraint_following", "alignment_tool_calling", "logic_instruction_following"],
                    "target_share": 0.45 if self.policy.recipe_profile == "apex" else 0.35,
                },
                {
                    "phase": "tool_code_reasoning",
                    "focus": ["logic_tool_grounding", "logic_coding_python", "logic_coding_cpp_rust", "sandbox_tools", "data_tooling"],
                    "target_share": 0.35 if self.policy.recipe_profile == "apex" else 0.30,
                },
                {
                    "phase": "knowledge_context_polish",
                    "focus": ["thai_grounding", "reverse_engineering", "cve_intelligence", "coverage_100k", "claude_reasoning_clean"],
                    "target_share": 0.20 if self.policy.recipe_profile == "apex" else 0.35,
                },
            ],
            "kept_records": kept_records,
            "domain_counts": dict(sorted(domain_counts.items())),
        }