scripts/build_dapt_dataset.py

#!/usr/bin/env python3
"""Build a deterministic, provenance-preserving ControlAI DAPT dataset.

The default profile targets roughly 20 million high-quality tokens. Lower-value
supplemental books and research papers are capped so they cannot swamp canonical
books, courses, and mathematical foundations. Proprietary MathWorks text is not
included; it remains a local RAG/tool-reference layer.
"""

from __future__ import annotations

import argparse
import hashlib
import json
import re
from collections import Counter, defaultdict
from dataclasses import dataclass
from pathlib import Path
from typing import Iterator

from transformers import AutoTokenizer


PROJECT_ROOT = Path(__file__).resolve().parents[1]
DEFAULT_OUTPUT = PROJECT_ROOT / "data" / "training" / "dapt_v1"
DEFAULT_TOKENIZER = "Qwen/Qwen3-4B-Instruct-2507"


@dataclass(frozen=True)
class Pool:
    name: str
    path: Path
    token_cap: int | None
    priority: int


DEFAULT_POOLS = (
    Pool(
        "canonical_books",
        PROJECT_ROOT / "data/processed/core_books_chunks/knowledge_chunks.jsonl",
        None,
        0,
    ),
    Pool(
        "control_courses_and_code",
        PROJECT_ROOT / "data/processed/chunks/knowledge_chunks.jsonl",
        None,
        1,
    ),
    Pool(
        "advanced_control_courses",
        PROJECT_ROOT / "data/processed/advanced_control_chunks/knowledge_chunks.jsonl",
        None,
        2,
    ),
    Pool(
        "university_web_collections",
        PROJECT_ROOT / "data/processed/web_collections_chunks/knowledge_chunks.jsonl",
        None,
        3,
    ),
    Pool(
        "math_signal_numerical_foundations",
        PROJECT_ROOT / "data/processed/general_foundations_chunks/knowledge_chunks.jsonl",
        6_000_000,
        4,
    ),
    Pool(
        "control_research",
        PROJECT_ROOT / "data/processed/arxiv_chunks/knowledge_chunks.jsonl",
        4_000_000,
        5,
    ),
    Pool(
        "supplemental_open_books",
        PROJECT_ROOT / "data/processed/open_books_chunks/knowledge_chunks.jsonl",
        5_000_000,
        6,
    ),
)


def canonical_text(text: str) -> str:
    text = text.replace("\u00ad", "")
    text = re.sub(r"(?<=\w)-\s*\n\s*(?=\w)", "", text)
    text = re.sub(r"[ \t]+", " ", text)
    text = re.sub(r"\n{3,}", "\n\n", text)
    return text.strip()


def text_hash(text: str) -> str:
    normalized = re.sub(r"\s+", " ", text).casefold().strip()
    return hashlib.sha256(normalized.encode("utf-8")).hexdigest()


def stable_fraction(value: str) -> float:
    raw = hashlib.sha256(value.encode("utf-8")).digest()[:8]
    return int.from_bytes(raw, "big") / 2**64


def simhash64(text: str) -> int:
    """Conservative near-duplicate fingerprint over sampled word 4-grams."""
    words = re.findall(r"[a-z0-9]+", text.casefold())
    features = [" ".join(words[index : index + 4]) for index in range(0, len(words) - 3, 2)]
    if not features:
        features = words
    accumulator = [0] * 64
    for feature in features:
        value = int.from_bytes(
            hashlib.blake2b(feature.encode("utf-8"), digest_size=8).digest(), "big"
        )
        for bit in range(64):
            accumulator[bit] += 1 if value & (1 << bit) else -1
    fingerprint = 0
    for bit, score in enumerate(accumulator):
        if score >= 0:
            fingerprint |= 1 << bit
    return fingerprint


def remove_near_duplicates(
    rows: list[dict], max_hamming: int
) -> tuple[list[dict], Counter[str]]:
    """Keep priority-first rows; use four 16-bit LSH bands for candidates."""
    kept: list[dict] = []
    fingerprints: list[int] = []
    bands: dict[tuple[int, int], list[int]] = defaultdict(list)
    removed: Counter[str] = Counter()
    mask = (1 << 16) - 1
    for row in rows:
        fingerprint = simhash64(row["text"])
        candidate_indices: set[int] = set()
        for band in range(4):
            value = (fingerprint >> (band * 16)) & mask
            candidate_indices.update(bands[(band, value)])
        duplicate = any(
            (fingerprint ^ fingerprints[index]).bit_count() <= max_hamming
            for index in candidate_indices
        )
        if duplicate:
            removed[row["source_pool"]] += 1
            continue
        row["simhash64"] = f"{fingerprint:016x}"
        index = len(kept)
        kept.append(row)
        fingerprints.append(fingerprint)
        for band in range(4):
            value = (fingerprint >> (band * 16)) & mask
            bands[(band, value)].append(index)
    return kept, removed


def quality_reason(text: str, token_count: int) -> str | None:
    if token_count < 80:
        return "too_short"
    if not text:
        return "empty"
    printable = sum(character.isprintable() or character in "\n\t" for character in text)
    if printable / len(text) < 0.98:
        return "nonprintable_noise"
    alphabetic = sum(character.isalpha() for character in text)
    if alphabetic / len(text) < 0.35:
        return "low_alphabetic_ratio"
    if re.search(r"(.)\1{12,}", text):
        return "repeated_character_noise"
    words = re.findall(r"[A-Za-z]{2,}", text)
    if len(words) < 25:
        return "too_few_words"
    return None


def rows(path: Path) -> Iterator[dict]:
    with path.open(encoding="utf-8") as stream:
        for line_number, line in enumerate(stream, start=1):
            if not line.strip():
                continue
            row = json.loads(line)
            if not isinstance(row, dict) or not isinstance(row.get("text"), str):
                raise ValueError(f"{path}:{line_number}: invalid chunk record")
            yield row


def select_pool(pool: Pool, tokenizer) -> tuple[list[dict], Counter[str]]:
    candidates = []
    rejected: Counter[str] = Counter()
    for row in rows(pool.path):
        text = canonical_text(row["text"])
        token_count = len(tokenizer.encode(text, add_special_tokens=False))
        reason = quality_reason(text, token_count)
        if reason:
            rejected[reason] += 1
            continue
        candidates.append(
            {
                "text": text,
                "token_count": token_count,
                "chunk_id": row.get("chunk_id"),
                "document_id": row.get("document_id"),
                "source_id": row.get("source_id"),
                "source_title": row.get("source_title"),
                "source_pool": pool.name,
                "source_priority": pool.priority,
                "text_sha256": text_hash(text),
            }
        )

    # Stable hash sampling avoids a leading-document bias when a pool is capped.
    candidates.sort(key=lambda row: stable_fraction(row["text_sha256"]))
    if pool.token_cap is None:
        return candidates, rejected
    selected = []
    used = 0
    for row in candidates:
        if used + row["token_count"] > pool.token_cap and selected:
            continue
        selected.append(row)
        used += row["token_count"]
        if used >= pool.token_cap:
            break
    return selected, rejected


def write_jsonl(path: Path, dataset: list[dict]) -> None:
    with path.open("w", encoding="utf-8") as stream:
        for row in dataset:
            stream.write(json.dumps(row, ensure_ascii=False) + "\n")


def document_key(row: dict) -> str:
    """Namespace document ids because independent source pools reuse filenames."""
    identity = row.get("document_id") or row["text_sha256"]
    return f"{row['source_pool']}::{identity}"


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--output-dir", type=Path, default=DEFAULT_OUTPUT)
    parser.add_argument("--tokenizer", default=DEFAULT_TOKENIZER)
    parser.add_argument("--allow-download", action="store_true")
    parser.add_argument(
        "--validation-percent",
        type=float,
        default=1.0,
        help="Document-level validation percentage",
    )
    parser.add_argument(
        "--near-duplicate-hamming",
        type=int,
        default=3,
        help="Drop SimHash candidates at or below this Hamming distance",
    )
    args = parser.parse_args()
    if not 0 < args.validation_percent < 20:
        parser.error("--validation-percent must be between 0 and 20")
    if not 0 <= args.near_duplicate_hamming <= 8:
        parser.error("--near-duplicate-hamming must be between 0 and 8")

    missing = [str(pool.path) for pool in DEFAULT_POOLS if not pool.path.exists()]
    if missing:
        raise FileNotFoundError("Missing chunk pools:\n" + "\n".join(missing))

    tokenizer = AutoTokenizer.from_pretrained(
        args.tokenizer, local_files_only=not args.allow_download
    )
    selected = []
    pool_stats = {}
    for pool in DEFAULT_POOLS:
        pool_rows, rejected = select_pool(pool, tokenizer)
        selected.extend(pool_rows)
        pool_stats[pool.name] = {
            "rows_before_cross_pool_dedup": len(pool_rows),
            "tokens_before_cross_pool_dedup": sum(
                row["token_count"] for row in pool_rows
            ),
            "rejected": dict(sorted(rejected.items())),
            "token_cap": pool.token_cap,
        }
        print(
            f"{pool.name}: {len(pool_rows):,} rows, "
            f"{pool_stats[pool.name]['tokens_before_cross_pool_dedup']:,} tokens"
        )

    # Higher-priority canonical sources win cross-pool exact duplicates.
    selected.sort(key=lambda row: (row["source_priority"], row["text_sha256"]))
    unique = []
    seen_hashes: set[str] = set()
    duplicates_by_pool: Counter[str] = Counter()
    for row in selected:
        if row["text_sha256"] in seen_hashes:
            duplicates_by_pool[row["source_pool"]] += 1
            continue
        seen_hashes.add(row["text_sha256"])
        unique.append(row)

    unique, near_duplicates_by_pool = remove_near_duplicates(
        unique, args.near_duplicate_hamming
    )

    # Stratify by source pool and keep entire documents together. A pure hash
    # threshold can accidentally leave a small canonical pool unrepresented.
    documents_by_pool: dict[str, set[str]] = defaultdict(set)
    for row in unique:
        documents_by_pool[row["source_pool"]].add(document_key(row))
    valid_documents: set[str] = set()
    for pool_name, document_ids in documents_by_pool.items():
        ordered = sorted(
            document_ids,
            key=lambda value: stable_fraction(f"validation:{pool_name}:{value}"),
        )
        count = max(1, round(len(ordered) * args.validation_percent / 100.0))
        valid_documents.update(ordered[:count])

    train = []
    valid = []
    for row in unique:
        identity = document_key(row)
        destination = valid if identity in valid_documents else train
        destination.append(row)

    # Stable shuffle prevents source blocks while preserving reproducibility.
    train.sort(key=lambda row: stable_fraction(f"train:{row['text_sha256']}"))
    valid.sort(key=lambda row: stable_fraction(f"valid:{row['text_sha256']}"))

    args.output_dir.mkdir(parents=True, exist_ok=True)
    write_jsonl(args.output_dir / "train.jsonl", train)
    write_jsonl(args.output_dir / "valid.jsonl", valid)
    summary = {
        "schema_version": 1,
        "tokenizer": args.tokenizer,
        "profile": "high_quality_approximately_20m",
        "proprietary_mathworks_included": False,
        "validation_split": "source_pool_stratified_document_hash",
        "validation_percent": args.validation_percent,
        "train_rows": len(train),
        "train_tokens": sum(row["token_count"] for row in train),
        "valid_rows": len(valid),
        "valid_tokens": sum(row["token_count"] for row in valid),
        "unique_documents": len({document_key(row) for row in unique}),
        "cross_pool_exact_duplicates_removed": sum(duplicates_by_pool.values()),
        "duplicates_removed_by_pool": dict(sorted(duplicates_by_pool.items())),
        "near_duplicate_hamming_threshold": args.near_duplicate_hamming,
        "near_duplicates_removed": sum(near_duplicates_by_pool.values()),
        "near_duplicates_removed_by_pool": dict(
            sorted(near_duplicates_by_pool.items())
        ),
        "pools": pool_stats,
    }
    (args.output_dir / "summary.json").write_text(
        json.dumps(summary, indent=2) + "\n", encoding="utf-8"
    )
    print(json.dumps(summary, indent=2))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())