| """ |
| Repair script 2: Re-shard oversized files (english 37x1.6GB, hindi-big 13x1.15GB) |
| Problem: 1GB shard_mb setting created viewer-breaking 1.6GB parquets (HTTP 500 + OOM). |
| Fix (Kaggle, streaming, delete-after-upload): |
| 1. Download one big shard -> read in batches -> re-write as ~150MB / 100k-row shards |
| under staging folder (english_fixed/, hindi_fixed/). |
| 2. Verify row counts, then delete originals and rename staging -> live in 2 commits. |
| 3. NEVER delete before verify ( HF has no trash for datasets). |
| |
| Run: |
| python data/scripts/repair_reshard_large.py --repo ViuAI/viu-mini-raw-pretrain --stream english --dry-run |
| python data/scripts/repair_reshard_large.py --repo ViuAI/viu-mini-raw-pretrain --stream english |
| python data/scripts/repair_reshard_large.py --repo ViuAI/viu-mini-raw-pretrain --stream hindi_big |
| |
| Note: 1.6GB download per file needs ~8GB disk+RAM; process ONE file at a time, |
| upload splits immediately, delete local. Total time ~2-4 hrs on Kaggle. |
| """ |
| import argparse |
| import gc |
| import tempfile |
| from pathlib import Path |
|
|
| TARGET_MB = 150 |
| TARGET_ROWS = 100_000 |
|
|
| def shard_plan(stream: str): |
| if stream == "english": |
| return [f"english/train-{i:05d}.parquet" for i in range(37)] |
| if stream == "hindi_big": |
| return [f"hindi/train-{i:05d}.parquet" for i in list(range(434, 447))] |
| raise ValueError("stream must be english or hindi_big") |
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--repo", default="ViuAI/viu-mini-raw-pretrain") |
| ap.add_argument("--stream", required=True, choices=["english", "hindi_big"]) |
| ap.add_argument("--dry-run", action="store_true") |
| ap.add_argument("--token", default=None) |
| ap.add_argument("--limit-files", type=int, default=None) |
| args = ap.parse_args() |
|
|
| from huggingface_hub import HfApi, hf_hub_download |
| from data.scripts.kaggle_collector import get_hf_token |
| import pandas as pd |
|
|
| token = args.token or get_hf_token() |
| api = HfApi(token=token) if not args.dry_run else None |
| files = shard_plan(args.stream) |
| if args.limit_files: |
| files = files[:args.limit_files] |
| stage = "english_fixed" if args.stream == "english" else "hindi_fixed" |
| print(f"[reshard] {args.stream}: {len(files)} big files -> staging {stage}/ (~{TARGET_MB}MB / {TARGET_ROWS} rows each)") |
|
|
| if args.dry_run: |
| for f in files: |
| print(f" would split {f}") |
| print("[dry-run] no download/upload. Run without --dry-run on Kaggle.") |
| return |
|
|
| tmp = Path(tempfile.mkdtemp(prefix="viu_reshard_")) |
| out_id = 0 |
| |
| existing = set(api.list_repo_files(repo_id=args.repo, repo_type="dataset")) |
| staged = sorted([f for f in existing if f.startswith(stage + "/")]) |
| if staged: |
| nums = [int(f.split("-")[-1].split(".")[0]) for f in staged] |
| out_id = max(nums) + 1 |
| print(f"[resume] staging has {len(staged)} shards, next out_id={out_id}") |
|
|
| for big in files: |
| print(f"\n[dl] {big} ...") |
| local = hf_hub_download(repo_id=args.repo, filename=big, repo_type="dataset") |
| import pyarrow.parquet as pq |
| pf = pq.ParquetFile(local) |
| print(f" rows={pf.metadata.num_rows:,}, row_groups={pf.metadata.num_row_groups}") |
| buf = [] |
| def flush(): |
| nonlocal out_id, buf |
| if not buf: |
| return |
| df = pd.DataFrame(buf, columns=["text", "lang", "source", "domain"]) |
| out_local = tmp / f"train-{out_id:05d}.parquet" |
| df.to_parquet(out_local, index=False, compression="snappy") |
| mb = out_local.stat().st_size / 1024 / 1024 |
| remote = f"{stage}/train-{out_id:05d}.parquet" |
| print(f" [up] {remote} {len(df):,} rows {mb:.1f}MB") |
| api.upload_file(path_or_fileobj=str(out_local), path_in_repo=remote, |
| repo_id=args.repo, repo_type="dataset", |
| commit_message=f"reshard {big} -> {remote} ({len(df):,} rows)") |
| out_local.unlink(missing_ok=True) |
| out_id += 1 |
| buf = [] |
| gc.collect() |
| for batch in pf.iter_batches(batch_size=20000): |
| for r in batch.to_pylist(): |
| |
| buf.append({k: r.get(k, "") for k in ("text", "lang", "source", "domain")}) |
| if len(buf) >= TARGET_ROWS: |
| flush() |
| flush() |
| print(f"[ok] split {big} done") |
| print(f"\n[done] staged {out_id} shards in {stage}/") |
| print("NEXT (manual, after verifying row counts in viewer):") |
| print(f" 1. Delete originals: {files[0]} ... ({len(files)} files)") |
| print(f" 2. Rename {stage}/* -> original folder via 2 commits (copy+delete).") |
| print(" See docs/HF_REPAIR_RUNBOOK.md for exact delete/rename commands.") |
|
|
| if __name__ == "__main__": |
| main() |
|
|
