File size: 4,912 Bytes
51b7902 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 | """
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
# resume staging counter
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():
# keep only unified cols (drop stray cols from legacy files)
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()
|