| """PUBLIC INTERFACE of Data_Module. Other modules may import ONLY this file. |
| |
| run_clean(cfg) -> dict |
| 1. surgically cleans chat.jsonl -> cleaned_data/sft_clean.jsonl (+ post-write verification) |
| 2. streams every SOURCE_REPOS entry (dataset first, model fallback), blacklist -> garbage filter |
| -> context-aware dedup (seeded with chat.jsonl for cross-deduplication) |
| 3. mixes PRETRAIN_CONV_RATIO of chat samples into cleaned_data/pretrain_clean.jsonl |
| """ |
| from __future__ import annotations |
|
|
| import gc |
| import json |
| from pathlib import Path |
| from typing import Dict, List, Optional |
|
|
| from .internal.blacklist_gatekeeper import BlacklistGatekeeper |
| from .internal.context_deduplicator import ContextDeduplicator |
| from .internal.hf_universal_streamer import HFUniversalStreamer, RepoTelemetry |
| from .internal.pretrain_garbage_filter import PretrainGarbageFilter |
| from .internal.sft_surgical_cleaner import ( |
| build_pretrain_chat_samples, clean_chat_file, verify_pretrain_jsonl, verify_sft_jsonl, write_starter_chat, |
| ) |
|
|
|
|
| def cleaned_paths(cfg: Dict) -> Dict[str, Path]: |
| out = Path(cfg["root"]) / "cleaned_data" |
| return {"dir": out, "sft": out / "sft_clean.jsonl", "pretrain": out / "pretrain_clean.jsonl", |
| "rejected": out / "sft_rejected.jsonl"} |
|
|
|
|
| def ensure_chat_file(cfg: Dict) -> Path: |
| """Create the 40-line starter chat.jsonl (37 clean + 3 intentionally broken lines) if it does not exist.""" |
| chat_file = Path(cfg["chat_file"]) |
| if write_starter_chat(chat_file): |
| print(f"[Data] {chat_file.name} not found -> generated the 40-line starter dataset at {chat_file}") |
| print(" (3 lines are broken/duplicate on purpose: the cleaner below will excise them and re-verify)") |
| return chat_file |
|
|
|
|
| def cleaned_data_ready(cfg: Dict) -> bool: |
| p = cleaned_paths(cfg) |
| return p["sft"].exists() and p["pretrain"].exists() and p["sft"].stat().st_size > 0 |
|
|
|
|
| def _parse_limit(value) -> Optional[int]: |
| if value is None: |
| return None |
| if str(value).strip().lower() in ("full", "all", "none", ""): |
| return None |
| try: |
| return max(1, int(value)) |
| except ValueError: |
| return None |
|
|
|
|
| def _print_table(rows: List[RepoTelemetry]) -> None: |
| header = ["Repo", "Type", "[Total Discovered]", "[Skipped by Blacklist]", |
| "[Destroyed Garbage/Dup]", "[Kept Clean Samples]"] |
| body = [[r.repo, r.repo_type, str(r.total_discovered), str(r.skipped_blacklist), |
| str(r.destroyed), str(r.kept)] for r in rows] |
| widths = [max(len(header[i]), *(len(row[i]) for row in body)) if body else len(header[i]) |
| for i in range(len(header))] |
| line = "+" + "+".join("-" * (w + 2) for w in widths) + "+" |
| print("\n TELEMETRY [Total Discovered] -> [Skipped by Blacklist] -> [Destroyed Garbage/Dup] -> [Kept]") |
| print(" " + line) |
| print(" | " + " | ".join(h.ljust(w) for h, w in zip(header, widths)) + " |") |
| print(" " + line) |
| for row in body: |
| print(" | " + " | ".join(c.ljust(w) for c, w in zip(row, widths)) + " |") |
| print(" " + line) |
| for r in rows: |
| print(f" {r.repo}: garbage-gate={r.garbage_gate} too-large={r.too_large} read-err={r.read_errors} " |
| f"filter-destroyed={r.filter_destroyed} dup-files={r.duplicate_files} " |
| f"partial-dedup(sub-sections excised)={r.partial_dedup}" |
| + (f" ERROR: {r.error}" if r.error else "")) |
|
|
|
|
| def run_clean(cfg: Dict) -> Dict: |
| paths = cleaned_paths(cfg) |
| paths["dir"].mkdir(parents=True, exist_ok=True) |
| token = cfg.get("hf_token") or None |
| chat_file = Path(cfg["chat_file"]) |
|
|
| print("\n[Data] === STEP 1/3: SFT surgical cleaning of chat.jsonl ===") |
| records, sft_report = clean_chat_file(chat_file, paths["sft"], paths["rejected"], verbose=True) |
| print(f" chat.jsonl: {sft_report['raw_lines']} lines -> {sft_report['kept']} clean " |
| f"({sft_report['removed']} removed, {sft_report['repaired_modes']} mode repairs)") |
| if not records: |
| raise RuntimeError("No usable SFT records after cleaning chat.jsonl") |
|
|
| print("\n[Data] === STEP 2/3: streaming + filtering source repos ===") |
| gate = BlacklistGatekeeper(cfg.get("extra_blocked_extensions", [])) |
| garbage_filter = PretrainGarbageFilter(json_quota=cfg.get("json_quota", 0.15)) |
| dedup = ContextDeduplicator() |
| for rec in records: |
| dedup.seed_text("\n".join([rec["user"], rec["thinking"], rec["assistant"]])) |
|
|
| telemetry: List[RepoTelemetry] = [] |
| repo_tmp = paths["dir"] / "_repo_text.tmp.jsonl" |
| repo_docs = 0 |
| try: |
| streamer = HFUniversalStreamer(token, gate) |
| except Exception as exc: |
| streamer = None |
| print(f" !! could not start HF streamer: {exc}") |
|
|
| with open(repo_tmp, "w", encoding="utf-8") as sink: |
| for entry in cfg.get("source_repos", []): |
| repo = entry.get("repo") if isinstance(entry, dict) else str(entry) |
| override = cfg.get("limit_override") |
| limit = _parse_limit(override if override not in (None, "") else |
| (entry.get("limit", 150) if isinstance(entry, dict) else 150)) |
| tel = RepoTelemetry(repo=repo) |
| telemetry.append(tel) |
| if streamer is None: |
| tel.error = "streamer unavailable" |
| continue |
| try: |
| for item in streamer.stream(repo, limit, tel): |
| cleaned = garbage_filter.process(item["path"], item["ext"], item["text"]) |
| if cleaned is None: |
| tel.filter_destroyed += 1 |
| continue |
| doc, info = dedup.process(item["path"], cleaned) |
| if doc is None: |
| tel.duplicate_files += 1 |
| continue |
| if info["status"] == "partial": |
| tel.partial_dedup += 1 |
| sink.write(json.dumps({"text": doc, "source": repo, "path": item["path"]}, |
| ensure_ascii=False) + "\n") |
| tel.kept += 1 |
| repo_docs += 1 |
| except Exception as exc: |
| tel.error = f"{type(exc).__name__}: {exc}" |
| print(f" !! stream error on {repo}: {tel.error}") |
| _print_table(telemetry) |
| if garbage_filter.stats: |
| print(" filter stats: " + ", ".join(f"{k}={v}" for k, v in sorted(garbage_filter.stats.items()))) |
| if repo_docs == 0: |
| print(" !! WARNING: no repo documents survived - pre-training will use chat.jsonl data only.") |
|
|
| print("\n[Data] === STEP 3/3: building pretrain_clean.jsonl (repo text + chat mix) ===") |
| ratio = float(cfg.get("conv_ratio", 0.25)) |
| chat_samples = build_pretrain_chat_samples(records, repo_docs, ratio) |
| total = repo_docs + len(chat_samples) |
| gap = max(1, repo_docs // len(chat_samples)) if chat_samples else 0 |
| written_chat = 0 |
| with open(paths["pretrain"], "w", encoding="utf-8") as out, open(repo_tmp, "r", encoding="utf-8") as src: |
| for i, line in enumerate(src, 1): |
| out.write(line) |
| if chat_samples and written_chat < len(chat_samples) and i % gap == 0: |
| out.write(json.dumps({"text": chat_samples[written_chat], "source": "chat.jsonl"}, |
| ensure_ascii=False) + "\n") |
| written_chat += 1 |
| while written_chat < len(chat_samples): |
| out.write(json.dumps({"text": chat_samples[written_chat], "source": "chat.jsonl"}, |
| ensure_ascii=False) + "\n") |
| written_chat += 1 |
| try: |
| repo_tmp.unlink() |
| except OSError: |
| pass |
| check = verify_pretrain_jsonl(paths["pretrain"]) |
| print(f" [PRETRAIN-VERIFY] {'PASSED' if check['ok'] else 'FAILED'}: {check['lines']} lines " |
| f"({repo_docs} repo docs + {written_chat} chat samples = " |
| f"{(100.0 * written_chat / max(total, 1)):.0f}% conversation share)") |
| if not check["ok"]: |
| raise RuntimeError("pretrain_clean.jsonl failed verification: " + "; ".join(check["errors"][:5])) |
| sft_check = verify_sft_jsonl(paths["sft"]) |
| if not sft_check["ok"]: |
| raise RuntimeError("sft_clean.jsonl failed final verification") |
|
|
| summary = {"sft_records": len(records), "repo_docs": repo_docs, "chat_pretrain_samples": written_chat, |
| "sft_path": str(paths["sft"]), "pretrain_path": str(paths["pretrain"])} |
| del records, chat_samples, dedup, garbage_filter, telemetry |
| gc.collect() |
| print(f"[Data] done -> {paths['sft'].name}, {paths['pretrain'].name} (RAM released)") |
| return summary |
