File size: 2,728 Bytes
ef18673 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 | """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()
|