data/ingest.py

"""Raw corpus ingestion utilities."""

from __future__ import annotations

import hashlib
import json
from dataclasses import dataclass
from pathlib import Path
from typing import Iterable, Iterator


@dataclass(frozen=True)
class SourceSpec:
    """Describes one raw corpus source."""

    name: str
    domain_tag: str
    quality_tier: str
    license_category: str
    estimated_tokens: int
    path: str
    text_key: str = "text"


SOURCE_REGISTRY: tuple[SourceSpec, ...] = (
    SourceSpec("general_web", "general", "medium", "permissive", 20_000_000_000, "data/raw/general_web.jsonl"),
    SourceSpec("code", "code", "high", "permissive", 8_000_000_000, "data/raw/code.jsonl"),
    SourceSpec("math_science", "math", "high", "permissive", 4_000_000_000, "data/raw/math_science.jsonl"),
    SourceSpec("books_longform", "general", "high", "restricted", 5_000_000_000, "data/raw/books.jsonl"),
    SourceSpec("multilingual", "multilingual", "medium", "permissive", 3_000_000_000, "data/raw/multilingual.jsonl"),
    SourceSpec("synthetic", "reasoning", "high", "permissive", 1_000_000_000, "data/raw/synthetic.jsonl"),
)


def iter_jsonl(path: Path, text_key: str = "text") -> Iterator[dict[str, object]]:
    """Yield JSONL records from disk."""
    with path.open("r", encoding="utf-8") as handle:
        for line in handle:
            line = line.strip()
            if not line:
                continue
            payload = json.loads(line)
            text = payload.get(text_key)
            if not isinstance(text, str) or not text.strip():
                continue
            yield payload


def stream_source(spec: SourceSpec) -> Iterator[dict[str, object]]:
    """Yield normalized records for one configured source."""
    path = Path(spec.path)
    if not path.exists():
        return iter(())
    return (
        {
            "id": stable_record_id(spec.name, record[spec.text_key]),
            "text": record[spec.text_key],
            "domain_tag": spec.domain_tag,
            "quality_tier": spec.quality_tier,
            "license_category": spec.license_category,
            "source_name": spec.name,
        }
        for record in iter_jsonl(path, spec.text_key)
    )


def stream_all_sources(sources: Iterable[SourceSpec] = SOURCE_REGISTRY) -> Iterator[dict[str, object]]:
    """Yield records from every source in the registry."""
    for spec in sources:
        yield from stream_source(spec)


def stable_record_id(source_name: str, text: str) -> str:
    """Hash a source+text pair into a stable content id."""
    digest = hashlib.sha256()
    digest.update(source_name.encode("utf-8"))
    digest.update(b"\0")
    digest.update(text.encode("utf-8"))
    return digest.hexdigest()