14.2 Data: Sourcing, Filtering, Dedup, Tokenize & Pack ~20B Tokens¶

"""
capstone/stacklm/data/synthetic.py

Source registry + offline synthetic corpus (Ch. 14.2). Each entry carries the
FULL coordinates needed to load it: repo id, config name, optional data_dir, and
the column the text actually lives in -- these differ per dataset, and getting
them wrong is the most common way a data pipeline silently produces zero
documents.
"""
import hashlib
import random
from dataclasses import dataclass
from typing import Iterator, Optional


@dataclass(frozen=True)
class DataMixEntry:
    name: str                      # short id, e.g. "fineweb_edu"
    hf_path: str                   # HuggingFace repo id
    weight: float                  # fraction of the 20B-token budget
    domain: str                    # "web"|"synthetic"|"code"|"math" -- routes the filter
    hf_config: Optional[str] = None    # config/subset name (`name=` in load_dataset)
    hf_data_dir: Optional[str] = None  # directory-sharded repos (starcoderdata)
    text_column: str = "text"          # NOT always "text" (starcoderdata: "content")
    revision: str = "main"             # pin to a commit sha for a reproducible corpus
    gated: bool = False                # requires huggingface_hub.login()


# The exact mix fixed in capstone/PLAN.md sec. 2. Weights sum to 1.0.
STACK100M_MIX = [
    DataMixEntry("fineweb_edu", "HuggingFaceFW/fineweb-edu", 0.70, "web",
                 hf_config="sample-100BT"),
    DataMixEntry("cosmopedia_v2", "HuggingFaceTB/smollm-corpus", 0.15, "synthetic",
                 hf_config="cosmopedia-v2"),
    DataMixEntry("starcoder", "bigcode/starcoderdata", 0.10, "code",
                 hf_data_dir="python", text_column="content", gated=True),
    DataMixEntry("finemath", "HuggingFaceTB/finemath", 0.05, "math",
                 hf_config="finemath-4plus"),
]

TOTAL_TOKEN_BUDGET = 20_000_000_000  # ~20B tokens, ~200 tok/param (PLAN.md sec. 2)


def load_hf_stream(entry: DataMixEntry):
    """Open one source as a streaming `datasets.IterableDataset`.

    Raises rather than returning an empty stream: a mistyped config or column is
    a bug, and a pipeline that silently yields nothing is the worst possible
    failure mode -- you find out 20 GPU-hours later.
    """
    from datasets import load_dataset  # heavy optional dependency; not in CI

    ds = load_dataset(
        entry.hf_path,
        name=entry.hf_config,
        data_dir=entry.hf_data_dir,
        split="train",
        streaming=True,
        revision=entry.revision,
    )
    cols = getattr(ds, "column_names", None)
    if cols is not None and entry.text_column not in cols:
        raise KeyError(
            f"{entry.name}: text column {entry.text_column!r} not in {cols}. "
            "Most HF text corpora use 'text', but bigcode/starcoderdata "
            "uses 'content'."
        )
    return ds


def stream_hf(entry: DataMixEntry, probe: int = 8) -> Iterator[dict]:
    """Yield normalized {"text","source","domain"} docs from the real dataset.
    Asserts that the first `probe` rows are not all empty, so a misconfigured
    source fails loudly instead of contributing zero tokens."""
    ds = load_hf_stream(entry)
    n_seen = n_nonempty = 0
    for row in ds:
        text = row.get(entry.text_column) or ""
        if n_seen < probe:
            n_seen += 1
            n_nonempty += bool(text)
            if n_seen == probe and n_nonempty == 0:
                raise ValueError(
                    f"{entry.name}: first {probe} rows had an empty "
                    f"{entry.text_column!r} field -- wrong column or wrong config?"
                )
        if text:
            yield {"text": text, "source": entry.name, "domain": entry.domain}

def stream_source(entry: DataMixEntry, offline: bool = False,
                  n_docs: int = 2000) -> Iterator[dict]:
    """Real HF stream when `offline=False`, synthetic fallback otherwise.

    Only the *opening* of the stream is guarded: if `datasets` is missing or the
    Hub is unreachable we fall back to the synthetic corpus, but once the stream
    is open, errors propagate. Swallowing mid-stream exceptions would silently
    truncate the corpus.
    """
    if not offline:
        try:
            gen = stream_hf(entry)
            first = next(gen)
        except (ImportError, OSError, ConnectionError):
            pass                                   # no `datasets` / no network
        else:
            yield first
            yield from gen
            return
    yield from synthetic_corpus(entry, n_docs=n_docs)

_dup_cache: dict = {}

_VOCAB = {
    "web": ["photosynthesis", "converts", "sunlight", "into", "chemical", "energy",
            "plants", "use", "carbon", "dioxide", "and", "water", "to", "produce",
            "glucose", "the", "process", "occurs", "inside", "chloroplasts"],
    "synthetic": ["chapter", "one", "introduces", "the", "concept", "of", "gravity",
                  "as", "a", "force", "that", "attracts", "objects", "with", "mass",
                  "toward", "each", "other", "consider", "an", "example"],
    "code": ["def", "compute", "(", "x", ")", ":", "return", "x", "*", "x", "+", "1",
             "for", "i", "in", "range", "(", "10", ")", ":", "print", "(", "i", ")"],
    "math": ["let", "f", "(", "x", ")", "=", "x^2", "then", "the", "derivative",
             "is", "2x", "solve", "for", "x", "when", "3x", "+", "5", "=", "20"],
}


def synthetic_corpus(entry: DataMixEntry, n_docs: int = 2000) -> Iterator[dict]:
    """Deterministic in-process corpus with injected exact (every 97th) and near
    (every 53rd) duplicates, so the dedup stages have something real to catch.
    It teaches the model nothing -- it exists so every downstream stage is
    exercised by CI as a hermetic stand-in for the real streams."""
    seed = int(hashlib.blake2b(entry.name.encode(), digest_size=4).hexdigest(), 16)
    rng = random.Random(seed)
    vocab = _VOCAB[entry.domain]
    for i in range(n_docs):
        if entry.domain in _dup_cache and i % 97 == 0:
            text = _dup_cache[entry.domain]                        # exact duplicate
        elif entry.domain in _dup_cache and i % 53 == 0:
            base = _dup_cache[entry.domain].split()                # near-duplicate
            for _ in range(max(1, len(base) // 20)):
                base[rng.randrange(len(base))] = rng.choice(vocab)
            text = " ".join(base)
        else:
            n_words = rng.randint(60, 400)
            text = " ".join(rng.choice(vocab) for _ in range(n_words)) + "."
            _dup_cache[entry.domain] = text
        doc_id = hashlib.sha1(f"{entry.name}-{i}".encode()).hexdigest()[:12]
        yield {"text": text, "source": entry.name, "domain": entry.domain, "doc_id": doc_id}

"""
capstone/stacklm/data/filters.py

Fast, dependency-free quality gates, domain-routed. A lean subset of Ch. 3.2,
tuned so filtering compute doesn't compete with training compute at 20B tokens.
"""
import hashlib
import json
import re

_WORD_RE = re.compile(r"\S+")

FILTER_CONFIG = {
    "web": {"min_words": 50, "max_words": 100_000, "min_mean_word_len": 3.0,
            "max_mean_word_len": 10.0, "min_alpha_frac": 0.60, "max_digit_frac": 0.20,
            "max_repeat_line_frac": 0.30},
    "code": {"min_chars": 20, "max_chars": 200_000, "max_char_frac": 0.30},
    "math": {"min_words": 20, "min_digit_frac": 0.03, "min_markers": 3},
}


def filter_config_hash() -> str:
    """Stable 12-hex-char digest of FILTER_CONFIG, recorded in the manifest."""
    blob = json.dumps(FILTER_CONFIG, sort_keys=True).encode("utf-8")
    return hashlib.blake2b(blob, digest_size=6).hexdigest()


def basic_stats(text: str) -> dict:
    words = _WORD_RE.findall(text)
    n_words = len(words) or 1
    n_chars = len(text) or 1
    alpha = sum(c.isalpha() for c in text)
    digit = sum(c.isdigit() for c in text)
    lines = text.splitlines() or [text]
    uniq = len(set(lines))
    return dict(
        n_words=n_words,
        alpha_frac=alpha / n_chars,
        digit_frac=digit / n_chars,
        mean_word_len=sum(len(w) for w in words) / n_words,
        dup_line_frac=1.0 - uniq / len(lines),   # boilerplate / nav-bar detector
    )


def passes_web_filter(text: str) -> bool:
    """Generic prose gate for FineWeb-Edu / Cosmopedia documents."""
    c = FILTER_CONFIG["web"]
    s = basic_stats(text)
    return (
        c["min_words"] <= s["n_words"] <= c["max_words"]
        and c["min_mean_word_len"] <= s["mean_word_len"] <= c["max_mean_word_len"]
        and s["alpha_frac"] >= c["min_alpha_frac"]
        and s["digit_frac"] <= c["max_digit_frac"]
        and s["dup_line_frac"] <= c["max_repeat_line_frac"]
    )


def passes_code_filter(text: str) -> bool:
    """Loose gate for StarCoder: reject empty/binary/minified-looking files
    (dominated by one repeated character); keep everything else, since code has
    a very different character distribution than prose and would be wrongly
    rejected by `passes_web_filter`."""
    c = FILTER_CONFIG["code"]
    if not (c["min_chars"] <= len(text) <= c["max_chars"]):
        return False
    head = text[:2000]
    most_common_frac = max(head.count(ch) for ch in set(head)) / max(len(head), 1)
    return most_common_frac <= c["max_char_frac"]


def passes_math_filter(text: str) -> bool:
    """FineMath gate: require mathematical density (digits, operators, LaTeX-ish
    markers), not generic prose fluency."""
    c = FILTER_CONFIG["math"]
    s = basic_stats(text)
    markers = sum(text.count(m) for m in ("=", "\\frac", "$", "^", "\\sum"))
    return s["n_words"] >= c["min_words"] and (
        s["digit_frac"] >= c["min_digit_frac"] or markers >= c["min_markers"]
    )


_FILTERS = {
    "web": passes_web_filter,
    "synthetic": passes_web_filter,
    "code": passes_code_filter,
    "math": passes_math_filter,
}


def quality_filter(doc: dict) -> bool:
    """Route a document to the filter for its source domain."""
    fn = _FILTERS.get(doc.get("domain", "web"), passes_web_filter)
    return fn(doc["text"])

"""
capstone/stacklm/data/dedup.py

Two-stage deduplication:
  1. Exact dedup  -- blake2b over normalized text, streaming.
  2. Near dedup   -- MinHash (Broder, 1997) + LSH banding over character
                     5-shingles, streaming with a bounded index.

Implemented from scratch (stdlib + numpy) so the mechanism is visible. For a
real 20B-token corpus use `datatrove`'s Minhash* pipeline -- see "Production
path" below for the throughput and memory arithmetic that forces it.
"""
import hashlib
import logging
import random
import re
from typing import Iterable, Iterator, List

import numpy as np

log = logging.getLogger(__name__)

_WS_RE = re.compile(r"\s+")
_MERSENNE_31 = (1 << 31) - 1  # keeps a*h+b inside int64 for vectorized minhashing


def normalize(text: str) -> str:
    """Lowercase + collapse whitespace, for a stable exact-dup hash key."""
    return _WS_RE.sub(" ", text.lower()).strip()


def exact_dedup(docs: Iterable[dict]) -> Iterator[dict]:
    """Drop documents whose normalized-text digest has been seen before.

    The payload is 16 bytes per *unique* document, but a CPython `set` of
    `bytes` costs ~80-90 B/element once the hash-table slot and object header
    are counted: budget ~1.7 GB at 20M unique documents, not the 320 MB the
    payload alone suggests.
    """
    seen: set = set()
    for doc in docs:
        h = hashlib.blake2b(normalize(doc["text"]).encode("utf-8"), digest_size=16).digest()
        if h in seen:
            continue
        seen.add(h)
        yield doc


def shingles(text: str, k: int = 5) -> List[str]:
    """Character k-shingles: robust to the small word-level edits that
    near-duplicates are made of, unlike whole-word shingles."""
    t = normalize(text)
    if len(t) < k:
        return [t]
    return list({t[i:i + k] for i in range(len(t) - k + 1)})


class MinHasher:
    """`num_perm` hash functions h_i(x) = (a_i*x + b_i) mod p give a signature of
    `num_perm` minima over a shingle set. Broder (1997):
    P(min_i(A) == min_i(B)) = Jaccard(A, B), so the fraction of matching
    signature positions is an unbiased estimator of the true Jaccard similarity
    -- without ever materializing the shingle sets to compare them.

    The permutation sweep is vectorized: the (num_perm, n_shingles) matrix is
    built once and reduced with `.min(axis=1)`, ~12x faster than the Python loop
    (measured ~8 ms vs ~104 ms on a 5.5 KB document).
    """

    def __init__(self, num_perm: int = 128, seed: int = 1234):
        self.num_perm = num_perm
        self._p = _MERSENNE_31
        rng = random.Random(seed)
        self.a = np.array([rng.randrange(1, self._p) for _ in range(num_perm)], dtype=np.int64)
        self.b = np.array([rng.randrange(0, self._p) for _ in range(num_perm)], dtype=np.int64)

    def _shingle_hashes(self, shingle_list) -> np.ndarray:
        raw = b"".join(
            hashlib.blake2b(s.encode("utf-8"), digest_size=4).digest() for s in shingle_list
        )
        h = np.frombuffer(raw, dtype=">u4").astype(np.int64)
        return h % self._p

    def signature(self, shingle_list) -> np.ndarray:
        """(num_perm,) uint32 signature. Deliberately NOT a tuple of Python
        ints: a 128-int tuple measures ~4.9 KB, this row is 128*4 = 512 B."""
        if not shingle_list:
            return np.zeros(self.num_perm, dtype=np.uint32)
        h = self._shingle_hashes(shingle_list)                       # (n_shingles,)
        mixed = (self.a[:, None] * h[None, :] + self.b[:, None]) % self._p
        return mixed.min(axis=1).astype(np.uint32)                   # (num_perm,)


class SignatureStore:
    """Contiguous (n, num_perm) uint32 store, grown by doubling up to `capacity`.

    Row cost is exactly num_perm*4 bytes -- 512 B at num_perm=128, versus ~4.9 KB
    for the tuple-of-ints it replaces.
    """

    def __init__(self, num_perm: int, capacity: int, initial: int = 4096):
        self.num_perm, self.capacity = num_perm, capacity
        self._buf = np.empty((max(1, min(initial, capacity)), num_perm), dtype=np.uint32)
        self.n = 0

    def __len__(self) -> int:
        return self.n

    def __getitem__(self, i: int) -> np.ndarray:
        return self._buf[i]

    def append(self, sig: np.ndarray) -> None:
        assert self.n < self.capacity, "SignatureStore is full"
        if self.n == self._buf.shape[0]:
            grown = np.empty((min(2 * self.n, self.capacity), self.num_perm),
                             dtype=np.uint32)
            grown[: self.n] = self._buf
            self._buf = grown
        self._buf[self.n] = sig
        self.n += 1


class LSHIndex:
    """Banding: split the signature into `bands` bands of `rows` rows. Two docs
    are *candidates* if any band matches exactly, turning an O(n^2) all-pairs
    scan into near-linear bucket lookups at the cost of a probabilistic
    threshold governed by the (1/bands)^(1/rows) S-curve."""

    def __init__(self, num_perm: int = 128, bands: int = 16):
        assert num_perm % bands == 0
        self.bands, self.rows = bands, num_perm // bands
        self.buckets: list = [dict() for _ in range(bands)]

    def _band_keys(self, sig) -> list:
        """32-byte `bytes` keys, not tuples of Python ints: ~65 B vs ~340 B each,
        and byte-identical across processes and runs."""
        sig = np.asarray(sig, dtype=np.uint32)
        return [sig[i * self.rows:(i + 1) * self.rows].tobytes() for i in range(self.bands)]

    def query_candidates(self, sig) -> set:
        cands: set = set()
        for b, key in enumerate(self._band_keys(sig)):
            cands.update(self.buckets[b].get(key, ()))
        return cands

    def insert(self, doc_idx: int, sig) -> None:
        for b, key in enumerate(self._band_keys(sig)):
            self.buckets[b].setdefault(key, []).append(doc_idx)


def estimate_jaccard(sig_a, sig_b) -> float:
    """Fraction of agreeing signature positions -- an unbiased estimator of the
    true Jaccard similarity (Broder, 1997)."""
    return float(np.mean(np.asarray(sig_a) == np.asarray(sig_b)))


def lsh_candidate_prob(jaccard: float, bands: int, rows: int) -> float:
    """P(at least one band collides) at true similarity J -- the LSH S-curve."""
    return 1.0 - (1.0 - jaccard ** rows) ** bands


def near_dedup_stream(docs: Iterable[dict], num_perm: int = 128, bands: int = 16,
                      threshold: float = 0.8, index_capacity: int = 500_000
                      ) -> Iterator[dict]:
    """Streaming near-dedup: yield documents that are not near-duplicates of an
    already-kept document. The corpus is never resident; memory is bounded by
    the index, measured at ~512 B/signature plus ~3.3 KB/document of LSH buckets
    (16 bands), i.e. ~1.9 GB at the 500k default.

    HARD CEILING: past `index_capacity` the index stops growing and near-dup
    recall for every later document silently drops to zero. We log loudly once.
    At 20B tokens (~20M documents) this ceiling is what forces the `datatrove`
    path -- do not just raise the number.
    """
    hasher = MinHasher(num_perm=num_perm)
    index = LSHIndex(num_perm=num_perm, bands=bands)
    store = SignatureStore(num_perm, index_capacity)
    warned = False
    for doc in docs:
        sig = hasher.signature(shingles(doc["text"]))
        if any(estimate_jaccard(sig, store[c]) >= threshold
               for c in index.query_candidates(sig)):
            continue
        if len(store) < index_capacity:
            index.insert(len(store), sig)
            store.append(sig)
        elif not warned:
            warned = True
            log.warning("near_dedup_stream: index_capacity=%d reached; near-dup "
                        "recall is now ZERO for the rest of this stream. Shard the "
                        "input or switch to datatrove's MinhashDedup* pipeline.",
                        index_capacity)
        yield doc


def near_dedup(docs, num_perm: int = 128, bands: int = 16,
               threshold: float = 0.8) -> list:
    """List-returning convenience wrapper around `near_dedup_stream`."""
    return list(near_dedup_stream(docs, num_perm=num_perm, bands=bands,
                                  threshold=threshold))

"""
capstone/scripts/dedup_datatrove.py  (production path -- needs `pip install datatrove[all]`)

MinHash deduplication of the Stack-100M corpus, using the same 4-stage pipeline
HuggingFace used for FineWeb. Not run in CI (heavy, multi-process, on-disk).
"""
from datatrove.executor.local import LocalPipelineExecutor
from datatrove.pipeline.dedup import (
    MinhashDedupSignature, MinhashDedupBuckets, MinhashDedupCluster, MinhashDedupFilter,
)
from datatrove.pipeline.dedup.minhash import MinhashConfig
from datatrove.pipeline.readers import JsonlReader
from datatrove.pipeline.writers.jsonl import JsonlWriter
from datatrove.utils.hashing import HashConfig

# FineWeb's settings: 5-grams, 112 permutations as 14 buckets x 8 hashes.
# The band/row S-curve of (1/14)^(1/8) ~ 0.72 targets documents ~75%+ similar --
# the same knob analyzed in the pitfall box above, just at production defaults.
cfg = MinhashConfig(
    hash_config=HashConfig(precision=64),
    num_buckets=14,
    hashes_per_bucket=8,
    n_grams=5,
)

IN, WORK, OUT = "s3://.../filtered", "/scratch/minhash", "/scratch/deduped"
TASKS = 64  # one task per CPU core; SlurmPipelineExecutor scales this to a cluster

stage1 = LocalPipelineExecutor(
    pipeline=[JsonlReader(IN),
              MinhashDedupSignature(output_folder=f"{WORK}/signatures", config=cfg)],
    tasks=TASKS, logging_dir=f"{WORK}/logs/sig")

stage2 = LocalPipelineExecutor(          # one task per bucket: buckets are independent
    pipeline=[MinhashDedupBuckets(input_folder=f"{WORK}/signatures",
                                  output_folder=f"{WORK}/buckets", config=cfg)],
    tasks=cfg.num_buckets, logging_dir=f"{WORK}/logs/buckets", depends=stage1)

stage3 = LocalPipelineExecutor(          # union-find over all candidate pairs: single task
    pipeline=[MinhashDedupCluster(input_folder=f"{WORK}/buckets",
                                  output_folder=f"{WORK}/remove_ids", config=cfg)],
    tasks=1, logging_dir=f"{WORK}/logs/cluster", depends=stage2)

stage4 = LocalPipelineExecutor(          # re-read the corpus, drop the flagged ids
    pipeline=[JsonlReader(IN),
              MinhashDedupFilter(input_folder=f"{WORK}/remove_ids"),
              JsonlWriter(OUT)],
    tasks=TASKS, logging_dir=f"{WORK}/logs/filter", depends=stage3)

if __name__ == "__main__":
    stage4.run()   # `depends` chains the whole DAG; running the last stage runs all four

"""
capstone/stacklm/data/pack.py

Pack tokenized documents into fixed-length SEQ_LEN=2048 windows with
document-aware segmentation. Every packed chunk begins with <bos>, so the token
array alone carries the document boundaries. Position ids are DERIVED, never
stored -- see `segments_from_bos`. The same "pack without cross-contamination"
recipe reappears for instruction data in Chat Templates, Data Formatting &
Sequence Packing (Ch. 5.2).
"""
from typing import Iterable, Iterator, Protocol
import numpy as np

SEQ_LEN = 2048  # Stack-100M pretraining max_seq_len (capstone/PLAN.md sec. 1)


class Tokenizer(Protocol):
    bos_id: int
    eos_id: int
    pad_id: int
    def encode(self, text: str) -> list: ...


def pack_documents(docs: Iterable[dict], tokenizer, seq_len: int = SEQ_LEN) -> Iterator[tuple]:
    """Greedily concatenate `<bos> body <eos>` across documents into fixed-length
    windows. Yields (input_ids, position_ids) per window; the final partial window
    is padded so no tokens are silently dropped. A document dict may carry
    pre-computed `ids` (the corpus builder tokenizes once, for budgeting).

    Documents longer than the window are chunked into seq_len-sized pieces, each
    treated as its own "document" for position-reset purposes -- keeping every
    position id inside [0, seq_len), which is both the range RoPE's theta=10000
    base (capstone/PLAN.md sec. 1) is tuned for and the valid index range of the
    model's precomputed rope_cos/rope_sin tables.
    """
    buf_ids: list = []
    buf_pos: list = []
    max_body = seq_len - 2  # room for <bos> and <eos> in every chunk

    for doc in docs:
        raw = doc["ids"] if "ids" in doc else tokenizer.encode(doc["text"])
        chunks = [raw[i:i + max_body] for i in range(0, len(raw), max_body)] or [[]]
        for chunk in chunks:
            toks = [tokenizer.bos_id, *chunk, tokenizer.eos_id]
            pos = list(range(len(toks)))  # this chunk's own position clock, from 0
            buf_ids.extend(toks)
            buf_pos.extend(pos)
            while len(buf_ids) >= seq_len:
                yield buf_ids[:seq_len], buf_pos[:seq_len]
                buf_ids, buf_pos = buf_ids[seq_len:], buf_pos[seq_len:]

    if buf_ids:  # flush a final, padded window
        pad_n = seq_len - len(buf_ids)
        buf_ids.extend([tokenizer.pad_id] * pad_n)
        buf_pos.extend([0] * pad_n)
        yield buf_ids, buf_pos


def segments_from_bos(input_ids: np.ndarray, bos_id: int) -> tuple:
    """Derive (seq_ids, position_ids) from a packed window's tokens alone.

    `seq_ids[i]` is the index of the document token i belongs to (-1 for a
    leading fragment continued from the previous window); `position_ids[i]` is
    the offset of token i inside its document. Storing these on disk would
    double the corpus for information the token array already contains.
    """
    starts = input_ids == bos_id
    seq_ids = np.cumsum(starts) - 1                       # -1 for a leading tail
    idx = np.arange(input_ids.shape[0])
    seg_start = np.maximum.accumulate(np.where(starts, idx, -1))
    return seq_ids, idx - np.maximum(seg_start, 0)


def build_intra_doc_causal_mask(position_ids: np.ndarray) -> np.ndarray:
    """Reconstruct the boolean (seq_len, seq_len) attention mask -- True means
    "token i may attend to token j". Two conditions: (a) causal, j <= i;
    (b) same document, i.e. i and j fall in the same position-reset segment.

    In production this dense mask is replaced by FlashAttention's `cu_seqlens`
    varlen API for O(seq_len) memory instead of O(seq_len^2) -- see
    FlashAttention I: IO-Awareness & The Online Softmax (Ch. 4.2) and Exercise 2.
    This dense version is for teaching, unit tests, and small-scale CPU eval.
    """
    seq_len = position_ids.shape[0]
    doc_id = np.cumsum(position_ids == 0)  # monotonically increasing per document
    causal = np.tril(np.ones((seq_len, seq_len), dtype=bool))
    same_doc = doc_id[:, None] == doc_id[None, :]
    return causal & same_doc


def segment_ids_from_positions(position_ids: np.ndarray) -> np.ndarray:
    """The per-token segment id (0,1,2,...) that `Stack100M.forward(seq_ids=...)`
    consumes to build its (B, 1, T, T) block-diagonal mask."""
    return np.cumsum(position_ids == 0) - 1

"""
capstone/stacklm/data/shard.py

One shard = `shard_XXXXX.tokens.bin` (flat uint16, n_seq x seq_len) plus a tiny
`.meta.bin` holding [n_seq, seq_len]. Position ids are NOT stored: they are
derived on read from `input_ids == bos_id`, halving the corpus on disk and in
page cache.
"""
import json
import numpy as np
from pathlib import Path

DTYPE = np.uint16  # vocab_size=32768 fits; see worked example


class ShardWriter:
    """Buffers packed windows and flushes a new shard every `seqs_per_shard`
    sequences (default target: ~100M tokens/shard, see the worked example)."""

    def __init__(self, out_dir: str, seq_len: int = 2048,
                 tokens_per_shard: int = 100_000_000,
                 store_positions: bool = False):
        self.out_dir = Path(out_dir)
        self.out_dir.mkdir(parents=True, exist_ok=True)
        self.seq_len = seq_len
        self.seqs_per_shard = max(1, tokens_per_shard // seq_len)
        self.store_positions = store_positions
        self._buf_ids: list = []
        self._buf_pos: list = []
        self._shard_idx = 0
        self.n_sequences = 0

    def add(self, input_ids, position_ids=None) -> None:
        self._buf_ids.append(np.asarray(input_ids, dtype=DTYPE))
        if self.store_positions:
            self._buf_pos.append(np.asarray(position_ids, dtype=DTYPE))
        self.n_sequences += 1
        if len(self._buf_ids) >= self.seqs_per_shard:
            self._flush()

    def _flush(self) -> None:
        if not self._buf_ids:
            return
        ids = np.stack(self._buf_ids)  # (n_seq, seq_len)
        stem = str(self.out_dir / f"shard_{self._shard_idx:05d}")
        ids.tofile(stem + ".tokens.bin")
        if self.store_positions:
            np.stack(self._buf_pos).tofile(stem + ".pos.bin")
        np.array([ids.shape[0], ids.shape[1]], dtype=np.int64).tofile(stem + ".meta.bin")
        self._shard_idx += 1
        self._buf_ids.clear()
        self._buf_pos.clear()

    def close(self) -> None:
        self._flush()  # flush the trailing partial shard

    def write_manifest(self, tokenizer=None, extra: dict = None) -> dict:
        """Write manifest.json beside the shards. `bos_id` is the load-bearing
        field: the dataset needs it to recover document boundaries on read."""
        man = {
            "seq_len": self.seq_len,
            "n_shards": self._shard_idx,
            "n_sequences": self.n_sequences,
            "n_tokens": self.n_sequences * self.seq_len,
            "store_positions": self.store_positions,
        }
        if tokenizer is not None:
            man.update(bos_id=int(tokenizer.bos_id), eos_id=int(tokenizer.eos_id),
                       pad_id=int(tokenizer.pad_id))
        if extra:
            man.update(extra)
        (self.out_dir / "manifest.json").write_text(json.dumps(man, indent=2))
        return man


def build_shards(docs, tokenizer, out_dir: str, seq_len: int = 2048,
                 tokens_per_shard: int = 100_000_000,
                 store_positions: bool = False, manifest_extra: dict = None) -> int:
    """End-to-end: pack a document stream, write shards + manifest. Returns the
    number of shards written."""
    from .pack import pack_documents
    writer = ShardWriter(out_dir, seq_len=seq_len, tokens_per_shard=tokens_per_shard,
                         store_positions=store_positions)
    for input_ids, position_ids in pack_documents(docs, tokenizer, seq_len=seq_len):
        writer.add(input_ids, position_ids)
    writer.close()
    writer.write_manifest(tokenizer=tokenizer, extra=manifest_extra)
    return writer._shard_idx

"""
capstone/stacklm/data/dataset.py

torch Dataset over the sharded uint16 .bin files. Only tokens are stored;
`seq_ids` (the segment id `Stack100M.forward` consumes for document-aware
masking) and `position_ids` (which index its RoPE tables) are derived per item
from `input_ids == bos_id`.
"""
import bisect
import json
import numpy as np
import torch
from pathlib import Path
from torch.utils.data import Dataset

from .pack import segments_from_bos


class PackedMemmapDataset(Dataset):
    def __init__(self, shard_dir: str, bos_id: int = None):
        self.shard_dir = Path(shard_dir)
        self._shards = []       # list of (tokens_memmap, pos_memmap_or_None)
        self._cum_seqs = [0]    # prefix sums of sequence counts, for indexing
        self.seq_len = None     # None, not a loop variable: an empty dir must not crash

        man_path = self.shard_dir / "manifest.json"
        self.manifest = json.loads(man_path.read_text()) if man_path.exists() else {}
        self.bos_id = bos_id if bos_id is not None else self.manifest.get("bos_id")

        for meta_path in sorted(self.shard_dir.glob("shard_*.meta.bin")):
            n_seq, seq_len = (int(x) for x in np.fromfile(meta_path, dtype=np.int64))
            stem = str(meta_path)[: -len(".meta.bin")]
            tok_mm = np.memmap(stem + ".tokens.bin", dtype=np.uint16, mode="r",
                               shape=(n_seq, seq_len))
            pos_path = Path(stem + ".pos.bin")
            pos_mm = (np.memmap(pos_path, dtype=np.uint16, mode="r", shape=(n_seq, seq_len))
                      if pos_path.exists() else None)
            self._shards.append((tok_mm, pos_mm))
            self._cum_seqs.append(self._cum_seqs[-1] + n_seq)
            self.seq_len = seq_len

        if self._shards and self._shards[0][1] is None and self.bos_id is None:
            raise ValueError(
                f"{shard_dir}: no .pos.bin and no bos_id (manifest.json missing?). "
                "Pass PackedMemmapDataset(dir, bos_id=tok.bos_id)."
            )

    def __len__(self) -> int:
        return self._cum_seqs[-1]

    def _locate(self, idx: int) -> tuple:
        """(shard, row) for a global sequence index. Binary search, not a linear
        scan: with 200 shards and ~10M samples/epoch a scan would cost billions
        of pointless Python comparisons per epoch."""
        if not 0 <= idx < self._cum_seqs[-1]:
            raise IndexError(idx)
        s = bisect.bisect_right(self._cum_seqs, idx) - 1
        return s, idx - self._cum_seqs[s]

    def __getitem__(self, idx: int) -> dict:
        s, row = self._locate(idx)
        tok_mm, pos_mm = self._shards[s]
        ids_np = tok_mm[row].astype(np.int64)
        if pos_mm is not None:                        # legacy shards with .pos.bin
            pos_np = pos_mm[row].astype(np.int64)
            seq_np = np.cumsum(pos_np == 0) - 1
        else:                                         # derive from the tokens alone
            seq_np, pos_np = segments_from_bos(ids_np, self.bos_id)
        ids = torch.from_numpy(ids_np)
        # Standard next-token target: shift by one. The training loop (Ch. 14.7)
        # masks the loss on pad targets, so the trailing padded window costs
        # nothing beyond the (small) wasted compute of its forward pass.
        return {
            "input_ids": ids[:-1],
            "position_ids": torch.from_numpy(np.ascontiguousarray(pos_np))[:-1],
            "seq_ids": torch.from_numpy(np.ascontiguousarray(seq_np))[:-1],
            "targets": ids[1:],
        }

"""
capstone/stacklm/data/build_corpus.py

The corpus driver: turn the four sources into sharded, packed `train/` + `val/`
corpora honouring the 70/15/10/5 mix and a total token budget.

Stages, per source: stream -> quality filter -> exact dedup -> near dedup ->
tokenize (once, in batches) -> stop at that source's token budget. The four
streams are then interleaved by weight, pushed through a reservoir shuffle
buffer, split into a deterministic document-level held-out set, packed, and
written as uint16 shards with a `manifest.json` recording the realized mix.
"""
import hashlib
import json
import random
from pathlib import Path

from .dedup import exact_dedup, near_dedup_stream
from .filters import filter_config_hash, quality_filter
from .pack import pack_documents
from .shard import ShardWriter
from .synthetic import STACK100M_MIX, TOTAL_TOKEN_BUDGET, stream_source


def encode_batched(docs, tokenizer, batch: int = 1024):
    """Attach `ids` to each document, encoding in batches.

    Every fast encoder in the ecosystem is a BATCH API -- Ch. 14.3's
    `encode_corpus` (multiprocessing.Pool over documents), HF `tokenizers`'
    `encode_batch`, `tiktoken`'s `encode_ordinary_batch`. Calling `encode` one
    document at a time leaves the process-parallel speedup on the table, and at
    84 GB of text that is ~4 core-hours instead of ~35 minutes.
    """
    encode_batch = getattr(tokenizer, "encode_batch", None)
    buf = []

    def flush():
        texts = [d["text"] for d in buf]
        id_lists = (encode_batch(texts) if encode_batch is not None
                    else [tokenizer.encode(t) for t in texts])
        for d, ids in zip(buf, id_lists):
            yield {**d, "ids": list(ids)}
        buf.clear()

    for doc in docs:
        buf.append(doc)
        if len(buf) >= batch:
            yield from flush()
    if buf:
        yield from flush()


def _source_pipeline(entry, tokenizer, budget_tokens, offline, dedup_kwargs, stats,
                     encode_batch_size: int = 1024):
    """Filtered, deduplicated, tokenized documents from one source, capped at
    `budget_tokens`. Tokenizing here (once) is what makes the budget exact."""
    docs = stream_source(entry, offline=offline)
    docs = (d for d in docs if quality_filter(d))
    docs = exact_dedup(docs)
    docs = near_dedup_stream(docs, **dedup_kwargs)
    used = 0
    for doc in encode_batched(docs, tokenizer, batch=encode_batch_size):
        if not doc["ids"]:
            continue
        used += len(doc["ids"]) + 2              # +2 for <bos>/<eos> added at pack time
        stats[entry.name] = used
        yield doc
        if used >= budget_tokens:
            return


def interleave_budgeted(entries, tokenizer, total_tokens, offline=True,
                        seed=1337, dedup_kwargs=None, stats=None):
    """Weighted round-robin over the per-source pipelines. Each source gets
    `weight * total_tokens`; a source that runs dry is dropped and the remaining
    weights renormalize -- deliberately UNLIKE
    `datasets.interleave_datasets(stopping_strategy=...)`, whose
    `"first_exhausted"` truncates the whole mix and whose `"all_exhausted"`
    oversamples the drained source. See the comparison at the end of this
    section."""
    stats = {} if stats is None else stats
    dedup_kwargs = dedup_kwargs or {}
    rng = random.Random(seed)
    gens, weights = {}, {}
    for e in entries:
        budget = int(round(e.weight * total_tokens))
        gens[e.name] = _source_pipeline(e, tokenizer, budget, offline, dedup_kwargs, stats)
        weights[e.name] = e.weight
        stats.setdefault(e.name, 0)
    alive = list(gens)
    while alive:
        name = rng.choices(alive, weights=[weights[n] for n in alive], k=1)[0]
        try:
            yield next(gens[name])
        except StopIteration:
            alive.remove(name)


def shuffle_buffer(docs, size=100_000, seed=1337):
    """Reservoir shuffle: without it, shard 0 would be ordered exactly as the
    sources were interleaved, and any resume-from-shard-k restart would see a
    biased slice of the mix."""
    rng = random.Random(seed)
    buf = []
    for doc in docs:
        if len(buf) < size:
            buf.append(doc)
            continue
        j = rng.randrange(size)
        yield buf[j]
        buf[j] = doc
    rng.shuffle(buf)
    yield from buf


def is_holdout(doc, per_mille: int = 1) -> bool:
    """Deterministic document-level held-out assignment. Hashing the *text*
    (not a counter) means the same document lands in the same split on every
    rebuild, and a document is never split across train and val -- the failure
    mode that silently contaminates a held-out perplexity."""
    h = int(hashlib.blake2b(doc["text"].encode("utf-8"), digest_size=8).hexdigest(), 16)
    return (h % 1000) < per_mille


def build_corpus(out_dir, tokenizer, total_tokens=TOTAL_TOKEN_BUDGET, entries=None,
                 seq_len=2048, tokens_per_shard=100_000_000, offline=True,
                 holdout_per_mille=1, holdout_tokens=10_000_000,
                 shuffle_size=100_000, seed=1337, dedup_kwargs=None) -> dict:
    """Build `out_dir/train` and `out_dir/val` shards. Returns the manifest."""
    entries = entries if entries is not None else STACK100M_MIX
    assert abs(sum(e.weight for e in entries) - 1.0) < 1e-9, "mix weights must sum to 1"
    out = Path(out_dir)
    stats: dict = {}
    val_docs: list = []
    held = {"tokens": 0}

    docs = interleave_budgeted(entries, tokenizer, total_tokens, offline=offline,
                               seed=seed, dedup_kwargs=dedup_kwargs, stats=stats)
    docs = shuffle_buffer(docs, size=shuffle_size, seed=seed)

    def train_stream():
        for doc in docs:
            if is_holdout(doc, holdout_per_mille):
                if held["tokens"] < holdout_tokens:
                    val_docs.append(doc)
                    held["tokens"] += len(doc["ids"]) + 2
                continue                     # held-out docs NEVER enter training
            yield doc

    train_writer = ShardWriter(out / "train", seq_len=seq_len,
                               tokens_per_shard=tokens_per_shard)
    for ids, pos in pack_documents(train_stream(), tokenizer, seq_len=seq_len):
        train_writer.add(ids, pos)
    train_writer.close()

    val_writer = ShardWriter(out / "val", seq_len=seq_len,
                             tokens_per_shard=tokens_per_shard)
    for ids, pos in pack_documents(iter(val_docs), tokenizer, seq_len=seq_len):
        val_writer.add(ids, pos)
    val_writer.close()

    realized = sum(stats.values()) or 1
    provenance = {
        "seed": seed,
        "token_budget": total_tokens,
        "filter_config_hash": filter_config_hash(),
        "dedup": {"num_perm": 128, "bands": 16, "threshold": 0.8, **(dedup_kwargs or {})},
        "sources": [
            {"name": e.name, "hf_path": e.hf_path, "hf_config": e.hf_config,
             "hf_data_dir": e.hf_data_dir, "revision": e.revision,
             "target_weight": e.weight,
             "realized_tokens": stats.get(e.name, 0),
             "realized_weight": round(stats.get(e.name, 0) / realized, 4)}
            for e in entries
        ],
        "holdout": {"per_mille": holdout_per_mille, "tokens": held["tokens"]},
        "offline_synthetic": offline,
    }
    man = train_writer.write_manifest(tokenizer=tokenizer, extra=provenance)
    val_writer.write_manifest(tokenizer=tokenizer, extra=provenance)
    (out / "manifest.json").write_text(json.dumps(man, indent=2))
    return man

from datasets import interleave_datasets, load_dataset

streams = [load_dataset(e.hf_path, name=e.hf_config, data_dir=e.hf_data_dir,
                        split="train", streaming=True, revision=e.revision)
           for e in STACK100M_MIX]

mixed = interleave_datasets(
    streams,
    probabilities=[e.weight for e in STACK100M_MIX],   # 0.70 / 0.15 / 0.10 / 0.05
    seed=1337,
    stopping_strategy="all_exhausted",  # keep sampling until every source is drained
).shuffle(seed=1337, buffer_size=100_000)              # streaming reservoir shuffle

"""
Toy end-to-end run of the Ch. 14.2 pipeline. No network, no GPU, ~10 seconds.
Swap ByteTokenizer for `stacklm.tokenizer.load_tokenizer()` for a real corpus.
"""
import tempfile
from stacklm.data import build_corpus, PackedMemmapDataset, STACK100M_MIX


class ByteTokenizer:
    """Dependency-free stand-in for Ch. 14.3's BPE tokenizer (vocab_size=32768).
    NOT what Stack-100M trains with -- it exists so this chapter's code runs
    before that artifact exists. It has no `encode_batch`, so `encode_batched`
    transparently falls back to the per-document path."""
    bos_id, eos_id, pad_id = 256, 257, 258
    vocab_size = 259

    def encode(self, text: str) -> list:
        return list(text.encode("utf-8"))


tok = ByteTokenizer()
out = tempfile.mkdtemp(prefix="stack100m_toy_")
man = build_corpus(out, tok,
                   total_tokens=200_000,      # 20B in the real run
                   seq_len=128,               # 2048 in the real run
                   tokens_per_shard=128 * 64,
                   offline=True,              # synthetic sources, hermetic
                   holdout_per_mille=20, holdout_tokens=20_000,
                   shuffle_size=256)

train = PackedMemmapDataset(f"{out}/train")   # bos_id comes from manifest.json
val = PackedMemmapDataset(f"{out}/val")
batch = train[0]

print(f"{man['n_shards']} shard(s), {len(train)} train seqs, {len(val)} val seqs")
print("realized mix:", {s["name"]: s["realized_weight"] for s in man["sources"]})
print("batch keys:", sorted(batch.keys()))    # input_ids, position_ids, seq_ids, targets

assert abs(sum(e.weight for e in STACK100M_MIX) - 1.0) < 1e-9
assert batch["input_ids"].shape == batch["targets"].shape == (man["seq_len"] - 1,)
assert int(batch["position_ids"].max()) < man["seq_len"]   # positions stay in range
assert int(batch["seq_ids"].max()) >= 0                    # at least one document