JetLaggedByData/scifi-forge

1"""2data/prepare_dataset.py3Step 2: Data preparation pipeline for QLoRA fine-tuning (V2).4 5Reads:  data/raw/internet_archive_scifi_v3.txt6Writes: data/chunks/scifi_train.jsonl7        data/chunks/scifi_val.jsonl8        data/chunks/dataset_stats.json9 10Output format (one JSON object per line):11  {12    "instruction": "Continue this science fiction story:",13    "input":  "<512-token context passage>",14    "output": "<128-token completion passage>"15  }16 17Run:18  python data/prepare_dataset.py19  python data/prepare_dataset.py --dry-run   # process 1000 samples only20"""21 22import re23import json24import random25import argparse26from pathlib import Path27from typing import Iterator28 29from dataset_config import (30    RAW_CORPUS, CHUNKS_DIR, TRAIN_JSONL, VAL_JSONL, STATS_JSON,31    CORPUS_START, CORPUS_END,32    CONTEXT_CHARS, COMPLETION_CHARS, WINDOW_STEP,33    TARGET_SAMPLES, VAL_FRACTION,34    INSTRUCTION, MIN_ALPHA_RATIO, MIN_WORDS, MAX_REPEAT_CHAR,35)36 37 38# ── Text cleaning ─────────────────────────────────────────────────────────39 40def clean_text(text: str) -> str:41    """42    Normalise corpus text:43    - Collapse runs of whitespace / tabs to single space44    - Normalise smart quotes and dashes to ASCII equivalents45    - Strip leading/trailing whitespace46    """47    text = text.replace("\t", " ")48    text = re.sub(r" {2,}", " ", text)49    text = text.replace("\u2018", "'").replace("\u2019", "'")50    text = text.replace("\u201c", '"').replace("\u201d", '"')51    text = text.replace("\u2013", "-").replace("\u2014", "-")52    text = re.sub(r"\n{3,}", "\n\n", text)53    return text.strip()54 55 56# ── Quality filters ───────────────────────────────────────────────────────57 58def is_quality_passage(text: str) -> bool:59    """60    Return True if passage meets minimum quality thresholds.61    Filters out header/footer junk, OCR artefacts, and repetitive noise.62    """63    if len(text.split()) < MIN_WORDS:64        return False65 66    total = len(text)67    if total == 0:68        return False69 70    alpha_ratio = sum(c.isalpha() for c in text) / total71    if alpha_ratio < MIN_ALPHA_RATIO:72        return False73 74    for char in set(text):75        if text.count(char) / total > MAX_REPEAT_CHAR:76            return False77 78    return True79 80 81# ── Sliding window chunker ────────────────────────────────────────────────82 83def sliding_window_samples(84    text: str,85    context_chars: int = CONTEXT_CHARS,86    completion_chars: int = COMPLETION_CHARS,87    step: int = WINDOW_STEP,88) -> Iterator[dict]:89    """90    Yield instruction-format dicts using a sliding window over `text`.91    Each sample: CONTEXT_CHARS of input + COMPLETION_CHARS of output.92    Window advances by `step` characters.93    """94    window = context_chars + completion_chars95    pos = 096    while pos + window <= len(text):97        chunk = text[pos: pos + window]98 99        # Split at nearest sentence boundary within ±50 chars of context end100        split_point = context_chars101        search_start = max(0, context_chars - 50)102        search_end = min(window, context_chars + 50)103        sentence_end = chunk.rfind(". ", search_start, search_end)104        if sentence_end != -1:105            split_point = sentence_end + 2       # include the space after period106 107        input_text  = chunk[:split_point].strip()108        output_text = chunk[split_point:split_point + completion_chars].strip()109 110        if input_text and output_text and is_quality_passage(input_text):111            yield {112                "instruction": INSTRUCTION,113                "input":  input_text,114                "output": output_text,115            }116 117        pos += step118 119 120# ── Train / val split ─────────────────────────────────────────────────────121 122def split_samples(123    samples: list[dict],124    val_fraction: float = VAL_FRACTION,125) -> tuple[list[dict], list[dict]]:126    """Shuffle then split into train / val."""127    random.shuffle(samples)128    val_n = max(1, int(len(samples) * val_fraction))129    return samples[val_n:], samples[:val_n]130 131 132# ── Writers ───────────────────────────────────────────────────────────────133 134def write_jsonl(samples: list[dict], path: Path) -> None:135    """Write list of dicts to a .jsonl file, one JSON object per line."""136    path.parent.mkdir(parents=True, exist_ok=True)137    with path.open("w", encoding="utf-8") as f:138        for sample in samples:139            f.write(json.dumps(sample, ensure_ascii=False) + "\n")140    print(f"  Wrote {len(samples):,} samples → {path}")141 142 143def write_stats(train: list, val: list, path: Path) -> None:144    """Save dataset statistics for later reference in benchmarks."""145    all_samples = train + val146    avg_input_len  = sum(len(s["input"])  for s in all_samples) / len(all_samples)147    avg_output_len = sum(len(s["output"]) for s in all_samples) / len(all_samples)148    stats = {149        "total_samples":      len(all_samples),150        "train_samples":      len(train),151        "val_samples":        len(val),152        "avg_input_chars":    round(avg_input_len,  1),153        "avg_output_chars":   round(avg_output_len, 1),154        "context_chars":      CONTEXT_CHARS,155        "completion_chars":   COMPLETION_CHARS,156        "window_step":        WINDOW_STEP,157        "corpus_chars_used":  CORPUS_END - CORPUS_START,158        "instruction":        INSTRUCTION,159    }160    path.write_text(json.dumps(stats, indent=2))161    print(f"\n── Dataset Stats ────────────────────────")162    for k, v in stats.items():163        print(f"  {k}: {v}")164    print(f"\n  Saved stats → {path}")165 166 167# ── Main pipeline ─────────────────────────────────────────────────────────168 169def run(dry_run: bool = False) -> None:170    """171    Full pipeline:172      1. Load and clean corpus slice (first 10M chars)173      2. Generate sliding-window samples174      3. Filter to TARGET_SAMPLES with quality checks175      4. Split train/val and write JSONL files176    """177    print(f"Loading corpus: {RAW_CORPUS}")178    if not RAW_CORPUS.exists():179        raise FileNotFoundError(180            f"Corpus not found at {RAW_CORPUS}\n"181            "Place internet_archive_scifi_v3.txt in data/raw/"182        )183 184    raw = RAW_CORPUS.read_text(encoding="utf-8")185    corpus = clean_text(raw[CORPUS_START:CORPUS_END])186    print(f"Corpus slice: {len(corpus):,} characters after cleaning")187 188    limit = 1_000 if dry_run else TARGET_SAMPLES189    print(f"Generating samples (target: {limit:,}, dry_run={dry_run})...")190 191    samples: list[dict] = []192    for sample in sliding_window_samples(corpus):193        samples.append(sample)194        if len(samples) % 5_000 == 0:195            print(f"  {len(samples):,} samples collected...")196        if len(samples) >= limit:197            break198 199    print(f"Total quality samples collected: {len(samples):,}")200 201    random.seed(42)202    train_samples, val_samples = split_samples(samples)203 204    print("\nWriting JSONL files...")205    write_jsonl(train_samples, TRAIN_JSONL)206    write_jsonl(val_samples,   VAL_JSONL)207    write_stats(train_samples, val_samples, STATS_JSON)208 209    print("\n✅ Data prep complete.")210    print(f"   Train: {TRAIN_JSONL}")211    print(f"   Val:   {VAL_JSONL}")212    print(f"   Stats: {STATS_JSON}")213 214 215if __name__ == "__main__":216    parser = argparse.ArgumentParser(description="SciFi Forge — data prep pipeline")217    parser.add_argument(218        "--dry-run", action="store_true",219        help="Generate only 1,000 samples for fast testing"220    )221    args = parser.parse_args()222    run(dry_run=args.dry_run)223