weather-llm-initial / scripts /preprocess_wiki.py
NagacharanVemula
Initial code-first model repo publish.
78dea75
Raw
History Blame Contribute Delete
1.48 kB
#!/usr/bin/env python3
"""Stream English Wikipedia articles into chunked JSONL (Hugging Face `datasets`)."""
from __future__ import annotations
import argparse
import sys
from pathlib import Path
ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT / "src"))
from tqdm import tqdm
from weather_llm.config_loader import load_yaml, repo_root_from
from weather_llm.data.jsonl_utils import write_jsonl
from weather_llm.data.wiki import wikipedia_to_chunks
def main() -> int:
ap = argparse.ArgumentParser()
ap.add_argument("--config", type=Path, default=None)
args = ap.parse_args()
root = repo_root_from(ROOT)
cfg_path = args.config or (root / "configs/data_default.yaml")
cfg = load_yaml(cfg_path)
wiki = cfg["wikipedia"]
paths = cfg["paths"]
out_path = root / paths["processed"] / "wikipedia.jsonl"
stream = wikipedia_to_chunks(
dataset_name=str(wiki["dataset_name"]),
dataset_config=str(wiki["dataset_config"]),
max_articles=int(wiki["max_articles"]),
min_article_chars=int(wiki["min_article_chars"]),
chunk_target_chars=int(wiki["chunk_target_chars"]),
chunk_overlap_chars=int(wiki["chunk_overlap_chars"]),
)
rows = []
for row in tqdm(stream, desc="wiki chunks"):
rows.append(row)
write_jsonl(out_path, iter(rows))
print("Wrote", len(rows), "chunks to", out_path)
return 0
if __name__ == "__main__":
raise SystemExit(main())