Spaces:
Sleeping
Sleeping
| """Materialize every corpus in corpora.yaml into data/corpora/<name>.txt. | |
| python src/prepare_data.py # download what's missing | |
| python src/prepare_data.py --force # re-download everything | |
| Sources: `url` (single text file), `hf` (HuggingFace datasets), `path` (already | |
| local — just validated). | |
| """ | |
| from __future__ import annotations | |
| import argparse | |
| import os | |
| import sys | |
| import urllib.request | |
| import config | |
| # The machine-wide HuggingFace cache (~/.cache/huggingface) is root-owned and | |
| # not writable for us, so keep datasets' cache inside the project. | |
| os.environ.setdefault("HF_HOME", os.path.join(config.ROOT, ".hf_cache")) | |
| def _from_url(url: str) -> str: | |
| with urllib.request.urlopen(url, timeout=60) as resp: | |
| return resp.read().decode("utf-8", errors="replace") | |
| def _from_hf(src: dict) -> str: | |
| from datasets import load_dataset | |
| ds = load_dataset(src["dataset"], src.get("config"), split=src.get("split", "train")) | |
| col = src.get("text_column", "text") | |
| # WikiText ships lots of blank lines; keep non-empty ones. | |
| parts = [t for t in ds[col] if t and t.strip()] | |
| return "".join(parts) | |
| def prepare_one(corpus: dict, force: bool = False) -> None: | |
| name = corpus["name"] | |
| src = corpus["source"] | |
| out_path = config.corpus_text_path(name) | |
| if src["type"] == "path": | |
| local = src["path"] if os.path.isabs(src["path"]) else os.path.join(config.ROOT, src["path"]) | |
| if not os.path.exists(local): | |
| print(f" [{name}] MISSING local file: {local}", file=sys.stderr) | |
| return | |
| if os.path.abspath(local) != os.path.abspath(out_path): | |
| text = open(local, encoding="utf-8", errors="replace").read() | |
| _write(out_path, text) | |
| print(f" [{name}] using local file ({os.path.getsize(out_path)/1e6:.2f} MB)") | |
| return | |
| if os.path.exists(out_path) and not force: | |
| print(f" [{name}] already present ({os.path.getsize(out_path)/1e6:.2f} MB) — skip") | |
| return | |
| print(f" [{name}] fetching via {src['type']} ...") | |
| text = _from_url(src["url"]) if src["type"] == "url" else _from_hf(src) | |
| _write(out_path, text) | |
| print(f" [{name}] wrote {os.path.getsize(out_path)/1e6:.2f} MB to {out_path}") | |
| def _write(path: str, text: str) -> None: | |
| os.makedirs(os.path.dirname(path), exist_ok=True) | |
| with open(path, "w", encoding="utf-8") as f: | |
| f.write(text) | |
| def main() -> None: | |
| ap = argparse.ArgumentParser() | |
| ap.add_argument("--force", action="store_true", help="re-download even if present") | |
| args = ap.parse_args() | |
| os.makedirs(config.DATA_DIR, exist_ok=True) | |
| print(f"Preparing corpora -> {config.DATA_DIR}") | |
| for corpus in config.corpora(): | |
| try: | |
| prepare_one(corpus, force=args.force) | |
| except Exception as e: # noqa: BLE001 — report and continue with others | |
| print(f" [{corpus['name']}] FAILED: {e}", file=sys.stderr) | |
| print("Done.") | |
| if __name__ == "__main__": | |
| main() | |