ViuMini-Dense-360M / data /scripts /kaggle_collector.py

feat: 100% overnight-resilient batch repair with 999 retries and persistent rate-limit handling

338f372 verified about 1 month ago

21.2 kB

"""
Kaggle Cloud Data Collector for ViuMini-MoE-242M.
Streams massive open datasets, cleans with stable blake2b hash dedup,
chunks into ~150MB Parquet shards (configurable via --shard_mb, default 150),
uploads to Hugging Face dataset repo `ViuAI/viu-mini-raw-pretrain`, and immediately deletes
local temporary files to keep Kaggle disk usage < 2GB.
Features:
  - Smart Auto-Resume: Automatically queries HF dataset repo, finds existing shards,
    starts numbering at (max_shard_id + 1), and fast-skips already processed rows.
  - ~150MB Shards: keeps HF viewer + DataLoader stable (old 1GB shards caused 1.6GB
    english files + viewer HTTP 500). Row cap 100k also enforced.
  - Zero Sub-Sampling: Ingests full corpora without artificial truncation.
Usage on Kaggle:
  python data/scripts/kaggle_collector.py --stream all --repo ViuAI/viu-mini-raw-pretrain --shard_mb 150
  python data/scripts/kaggle_collector.py --stream hindi --repo ViuAI/viu-mini-raw-pretrain --shard_mb 150
"""
import argparse
import gc
import gzip
import hashlib
import itertools
import json
import os
from pathlib import Path
import re
import sys
import tempfile
import time

DEV = re.compile(r'[\u0900-\u097F]')
LAT = re.compile(r'[A-Za-z]')


def get_hf_token():
    """Retrieve HF token from Kaggle secrets, Colab secrets, env var, or user input."""
    # 1. Check Kaggle secrets
    try:
        from kaggle_secrets import UserSecretsClient
        tok = UserSecretsClient().get_secret("HF_TOKEN")
        if tok:
            return tok
    except Exception:
        pass
    # 2. Check Google Colab secrets
    try:
        from google.colab import userdata
        tok = userdata.get("HF_TOKEN")
        if tok:
            return tok
    except Exception:
        pass
    # 3. Check environment variable
    tok = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")
    if tok:
        return tok
    return None


def stable_hash(s: str) -> int:
    """Stable 64-bit hash (PYTHONHASHSEED-independent). hash() is randomized per run,
    so dedup was unstable across Kaggle restarts -> duplicates. Use blake2b."""
    return int.from_bytes(hashlib.blake2b(s.encode("utf-8"), digest_size=8).digest(), "big")


def dev_ratio(s: str) -> float:
    dev = len(DEV.findall(s))
    lat = len(LAT.findall(s))
    tot = dev + lat
    return dev / max(tot, 1)


def is_valid_line(s: str, lang: str) -> bool:
    if len(s) < 10 or len(s.split()) < 3:
        return False
    if lang == "bilingual":
        return len(s) >= 15
    r = dev_ratio(s)
    if lang == "hinglish":
        # Must have Latin letters, Devanagari ratio <= 30%
        return bool(LAT.search(s)) and r <= 0.3
    elif lang == "hindi":
        # Must have Devanagari script
        return bool(DEV.search(s))
    elif lang == "english":
        # Must have Latin, zero Devanagari
        return bool(LAT.search(s)) and not bool(DEV.search(s))
    return True


def apply_skip(ds, skip_rows: int, source_name: str):
    """Applies fast skipping or falls back to itertools.islice."""
    if skip_rows <= 0:
        return ds
    try:
        ds = ds.skip(skip_rows)
        print(f"[resume] Fast-skipped {skip_rows:,} rows from {source_name}.")
        return ds
    except Exception as e:
        print(f"[resume] Fast-skip failed ({e}), using iterator slice for {source_name}...")
        return itertools.islice(ds, skip_rows, None)


def stream_hindi_data(skip_rows: int = 0):
    """Streams full IndicCorpV2 hin_Deva (152M rows) and Sangraha."""
    try:
        from datasets import load_dataset
        print("[source] Loading ai4bharat/IndicCorpV2 (indiccorp_v2, hin_Deva) streaming...")
        ds = load_dataset("ai4bharat/IndicCorpV2", "indiccorp_v2", split="hin_Deva", streaming=True)
        ds = apply_skip(ds, skip_rows, "IndicCorpV2")
        for r in ds:
            txt = (r.get("text") or "").strip()
            if txt:
                yield txt, "indiccorp_v2"
    except Exception as e:
        print(f"[warn] IndicCorpV2 stream error: {e}, falling back to sangraha...")
        try:
            from datasets import load_dataset
            ds = load_dataset("ai4bharat/sangraha", "verified", split="hin", streaming=True)
            ds = apply_skip(ds, skip_rows, "sangraha")
            for r in ds:
                txt = (r.get("text") or "").strip()
                if txt:
                    yield txt, "sangraha"
        except Exception as e2:
            print(f"[error] Sangraha stream error: {e2}")


def stream_english_data(skip_rows: int = 0):
    """Streams Fineweb-edu / TinyStories educational slices."""
    from datasets import load_dataset
    print("[source] Loading HuggingFaceFW/fineweb-edu (sample-10BT or train) streaming...")
    try:
        ds = load_dataset("HuggingFaceFW/fineweb-edu", split="train", streaming=True)
        ds = apply_skip(ds, skip_rows, "fineweb-edu")
        for r in ds:
            txt = (r.get("text") or "").strip()
            if txt:
                yield txt, "fineweb-edu"
    except Exception as e:
        print(f"[warn] fineweb-edu fallback to TinyStories: {e}")
        ds = load_dataset("roneneldan/TinyStories", split="train", streaming=True)
        ds = apply_skip(ds, skip_rows, "tinystories")
        for r in ds:
            txt = (r.get("text") or "").strip()
            if txt:
                yield txt, "tinystories"


def stream_hinglish_data(skip_rows: int = 0):
    """Streams sk-community/romanized_hindi and conversational corpora."""
    from datasets import load_dataset
    print("[source] Loading sk-community/romanized_hindi streaming...")
    try:
        ds = load_dataset("sk-community/romanized_hindi", split="train", streaming=True)
        ds = apply_skip(ds, skip_rows, "romanized_hindi")
        for r in ds:
            txt = (r.get("Transliterated Hindi") or r.get("hi_rom") or r.get("text") or "").strip()
            if txt:
                yield txt, "romanized_hindi"
    except Exception as e:
        print(f"[warn] sk-community/romanized_hindi stream error: {e}")


def stream_domain_data():
    """Streams curated 26-pillar knowledge from collect_domains."""
    try:
        from data.scripts.collect_domains import generate_domain_samples
        for rec in generate_domain_samples():
            yield rec["text"], rec["source"], rec["lang"], rec["domain"]
    except ImportError:
        try:
            from collect_domains import generate_domain_samples
            for rec in generate_domain_samples():
                yield rec["text"], rec["source"], rec["lang"], rec["domain"]
        except Exception as e:
            print(f"[warn] domain generator import failed: {e}")


def stream_translation_data(skip_rows: int = 0):
    """Streams parallel translation data from cfilt/iitb-english-hindi & Romanized Hindi."""
    from datasets import load_dataset
    print("[source] Loading cfilt/iitb-english-hindi parallel translation streaming...")
    try:
        ds = load_dataset("cfilt/iitb-english-hindi", split="train", streaming=True)
        ds = apply_skip(ds, skip_rows, "iitb-english-hindi")
        for r in ds:
            trans = r.get("translation") or {}
            en = (trans.get("en") or "").strip()
            hi = (trans.get("hi") or "").strip()
            if en and hi and len(en) >= 10 and len(hi) >= 10:
                text = f"English: {en}\nHindi: {hi}"
                yield text, "iitb-english-hindi", "bilingual", "translation"
    except Exception as e:
        print(f"[warn] iitb translation stream fallback: {e}")


def get_existing_shards_info(api, repo_id: str, stream_name: str):
    """
    Checks HF repo for existing shards in {stream_name}/.
    Returns (next_shard_id, max_shard_id, existing_count).
    """
    if not api:
        return 0, -1, 0
    try:
        files = api.list_repo_files(repo_id=repo_id, repo_type="dataset")
        prefix = f"{stream_name}/train-"
        shard_ids = []
        for f in files:
            if f.startswith(prefix) and (f.endswith(".parquet") or f.endswith(".jsonl.gz")):
                base = f[len(prefix):]
                num_str = base.split(".")[0]
                if num_str.isdigit():
                    shard_ids.append(int(num_str))
        if shard_ids:
            max_id = max(shard_ids)
            return max_id + 1, max_id, len(shard_ids)
    except Exception as e:
        print(f"[warn] Could not inspect existing shards for {stream_name}: {e}")
    return 0, -1, 0


def write_and_upload_shard(records, shard_id, stream_name, repo_id, api, dry_run=False, tmp_dir=None):
    """
    Writes records to a compressed parquet shard, uploads to HF,
    and immediately removes the local file and frees memory.
    """
    if tmp_dir is None:
        tmp_dir = Path(tempfile.gettempdir()) / "viu_shards"
    tmp_path = Path(tmp_dir) / stream_name
    tmp_path.mkdir(parents=True, exist_ok=True)

    file_name = f"train-{shard_id:05d}.parquet"
    out_file = tmp_path / file_name
    remote_path = f"{stream_name}/{file_name}"

    try:
        import pandas as pd
        df = pd.DataFrame(records)
        df.to_parquet(out_file, index=False, compression="snappy")
        del df
        gc.collect()
    except Exception as e:
        # Fallback to jsonl.gz if parquet/pandas is unavailable
        file_name = f"train-{shard_id:05d}.jsonl.gz"
        out_file = tmp_path / file_name
        remote_path = f"{stream_name}/{file_name}"
        with gzip.open(out_file, "wt", encoding="utf-8") as f:
            for r in records:
                f.write(json.dumps(r, ensure_ascii=False) + "\n")
        gc.collect()

    size_mb = out_file.stat().st_size / (1024 * 1024)
    print(f"[shard] Serialized {file_name} with {len(records):,} rows ({size_mb:.2f} MB)")

    if dry_run:
        print(f"[dry-run] Skipped upload of {out_file} to {repo_id}/{remote_path}")
        out_file.unlink(missing_ok=True)
        return True

    # Upload to HF with retry and exponential backoff + 429 rate limit handler (Overnight resilient)
    retries = 999
    success = False
    for attempt in range(1, retries + 1):
        try:
            api.upload_file(
                path_or_fileobj=str(out_file),
                path_in_repo=remote_path,
                repo_id=repo_id,
                repo_type="dataset",
                commit_message=f"Upload {remote_path} ({len(records):,} rows, {size_mb:.1f} MB)"
            )
            print(f"[hf] Uploaded successfully: {remote_path} ({size_mb:.2f} MB)")
            success = True
            break
        except Exception as err:
            err_str = str(err)
            if "429" in err_str or "rate limit" in err_str.lower():
                wait_time = 120
                m = re.search(r"Retry after (\d+) seconds", err_str)
                if m:
                    wait_time = int(m.group(1)) + 5
                print(f"[rate-limit] HF 128 commits/hr limit hit. Waiting {wait_time}s before retry ({attempt}/{retries})...")
                time.sleep(wait_time)
            else:
                print(f"[warn] Upload attempt {attempt}/{retries} failed: {err}")
                time.sleep(4 * attempt)

    # Immediately delete local temporary shard to keep Kaggle disk clean
    out_file.unlink(missing_ok=True)
    if not success:
        print(f"[error] Failed to upload {remote_path} after {retries} retries!")
    return success


def run_collector(stream_name, repo_id, shard_mb=150, shard_rows=100_000,
                  resume=True, skip_rows=None, limit=None, dry_run=False, token=None):
    from huggingface_hub import HfApi

    api = None
    if not dry_run:
        if not token:
            token = get_hf_token()
        if not token:
            print("[error] HF token missing! Set HF_TOKEN environment variable or Kaggle secret.")
            sys.exit(1)
        api = HfApi(token=token)
        try:
            api.create_repo(repo_id=repo_id, repo_type="dataset", exist_ok=True)
            print(f"[hf] Verified dataset repo: {repo_id}")
        except Exception as e:
            print(f"[warn] create_repo warning: {e}")

    # Determine starting shard_id and rows to skip if resuming
    # FIX: old code assumed 100k rows/shard even for 1.6GB english shards (~3-4M rows),
    # causing massive duplication on resume. Now shard_rows is enforced, so
    # skip = (max_id+1) * shard_rows is correct. For legacy 1GB shards, require
    # explicit --skip_rows or manual backfill (see docs/HF_REPAIR_RUNBOOK.md).
    shard_id = 0
    actual_skip = 0
    if resume and not dry_run:
        next_id, max_id, count = get_existing_shards_info(api, repo_id, stream_name)
        if count > 0:
            shard_id = next_id
            if skip_rows is None:
                actual_skip = (max_id + 1) * shard_rows
                print(f"[resume] Found {count} existing shards in '{stream_name}' (up to train-{max_id:05d}.parquet).")
                print(f"[resume] Resuming at shard_id={shard_id} (train-{shard_id:05d}.parquet), auto-skipping ~{actual_skip:,} rows.")
                print(f"[resume] NOTE: legacy 1GB shards (english/hindi-big) used variable rows. "
                      f"If resuming those, pass explicit --skip_rows. See docs/HF_REPAIR_RUNBOOK.md")
            else:
                actual_skip = skip_rows
                print(f"[resume] Resuming at shard_id={shard_id}, skipping specified {actual_skip:,} rows.")
        else:
            actual_skip = skip_rows or 0
    else:
        actual_skip = skip_rows or 0

    # Target byte calculation: Snappy parquet ~2.7x for hindi text, ~1.6x for english.
    # FIX: old default 1000MB caused 1.6GB english shards -> HF viewer 500 + OOM.
    # New default 150MB + hard row cap 100k keeps shards 30-150MB, viewer-stable.
    # Conservative ratio 2.0 to avoid overshoot on english.
    target_raw_bytes = int(shard_mb * 1024 * 1024 * 2.0)
    if stream_name == "domains":
        # Domain samples are curated; use smaller chunk threshold to flush promptly
        target_raw_bytes = int(min(shard_mb, 150) * 1024 * 1024 * 2.7)

    print(f"\n=== Starting Collector: stream={stream_name}, repo={repo_id} ===")
    print(f"    Target shard size: ~{shard_mb} MB compressed | Start shard: train-{shard_id:05d}.parquet | Skip: {actual_skip:,} rows")

    seen_hashes = set()
    current_records = []
    current_raw_bytes = 0
    total_kept = 0
    start_time = time.time()
    last_log_time = start_time

    # Determine generator
    if stream_name == "hindi":
        gen = ((txt, src, "hindi", "general_hindi") for txt, src in stream_hindi_data(skip_rows=actual_skip))
    elif stream_name == "english":
        gen = ((txt, src, "english", "general_english") for txt, src in stream_english_data(skip_rows=actual_skip))
    elif stream_name == "hinglish":
        gen = ((txt, src, "hinglish", "general_hinglish") for txt, src in stream_hinglish_data(skip_rows=actual_skip))
    elif stream_name == "domains":
        gen = stream_domain_data()
    elif stream_name == "translation":
        gen = stream_translation_data(skip_rows=actual_skip)
    else:
        raise ValueError(f"Unknown stream: {stream_name}")

    for item in gen:
        txt, src, lang, dom = item[0], item[1], item[2], item[3]
        s = txt.strip()
        if not is_valid_line(s, lang):
            continue

        # In-flight stable-hash deduplication (FIX: hash() was seed-randomized)
        # NOTE: seen_hashes grows with total rows. For >10M-row runs on Kaggle,
        # memory can blow up. We cap at 5M entries then clear oldest via sampling
        # reset to keep RAM <2GB (small duplicate leak acceptable vs OOM).
        h = stable_hash(s)
        if h in seen_hashes:
            continue
        seen_hashes.add(h)
        if len(seen_hashes) > 5_000_000:
            print("[dedup] seen_hashes >5M, clearing to bound RAM (minor dup leak accepted)")
            seen_hashes.clear()

        row_bytes = len(s.encode('utf-8'))
        current_records.append({
            "text": s,
            "lang": lang,
            "source": src,
            "domain": dom
        })
        current_raw_bytes += row_bytes
        total_kept += 1

        # Periodic status log every 100,000 rows or every 30 seconds
        now = time.time()
        if total_kept % 100_000 == 0 or (now - last_log_time) >= 30.0:
            elapsed = max(now - start_time, 1)
            rate = total_kept / elapsed
            est_parquet_mb = (current_raw_bytes / 2.7) / (1024 * 1024)
            print(f"[progress] {total_kept:,} rows | Shard buffer: {len(current_records):,} rows (~{est_parquet_mb:.1f} MB est. parquet) | {rate:.1f} rows/s")
            last_log_time = now

        # Flush condition: Target MB reached OR explicit row cap reached
        hit_bytes = current_raw_bytes >= target_raw_bytes
        hit_rows = (shard_rows is not None and len(current_records) >= shard_rows)
        if hit_bytes or hit_rows:
            write_and_upload_shard(current_records, shard_id, stream_name, repo_id, api, dry_run=dry_run)
            shard_id += 1
            current_records = []
            current_raw_bytes = 0
            gc.collect()

        if limit and total_kept >= limit:
            print(f"[limit] Reached limit of {limit:,} rows.")
            break

    # Flush remaining records
    if current_records:
        write_and_upload_shard(current_records, shard_id, stream_name, repo_id, api, dry_run=dry_run)
        shard_id += 1
        current_records = []
        gc.collect()

    elapsed = max(time.time() - start_time, 1)
    print(f"=== Collector Finished: {total_kept:,} total rows in this session ({elapsed:.1f}s, {total_kept/elapsed:.1f} rows/s) ===")


def main():
    parser = argparse.ArgumentParser(description="Kaggle High-Throughput Cloud Stream Collector")
    parser.add_argument("--stream", type=str, required=True,
                        choices=["hindi", "english", "hinglish", "domains", "translation", "all"],
                        help="Data stream to collect (or 'all' for sequential run)")
    parser.add_argument("--repo", type=str, default="ViuAI/viu-mini-raw-pretrain", help="HF dataset repo ID")
    parser.add_argument("--shard_mb", type=int, default=150, help="Target compressed Parquet shard size in MB (default 150 MB, viewer-stable; old 1000 MB caused 1.6GB shards + viewer 500)")
    parser.add_argument("--shard_rows", type=int, default=100000, help="Row cap per shard (default 100k; whichever hits first: bytes or rows)")
    parser.add_argument("--no-resume", action="store_true", help="Do not auto-resume; start from shard 0 and row 0")
    parser.add_argument("--skip_rows", type=int, default=None, help="Explicitly skip N rows before streaming")
    parser.add_argument("--limit", type=int, default=None, help="Optional max rows to collect")
    parser.add_argument("--dry-run", action="store_true", help="Run without uploading to HF")
    parser.add_argument("--token", type=str, default=None, help="HF API token")
    args = parser.parse_args()

    resume = not args.no_resume

    if args.stream == "all":
        # Prioritize Hinglish and English so they accumulate immediately alongside Hindi
        streams = ["domains", "translation", "hinglish", "english", "hindi"]
        for s in streams:
            print(f"\n==========================================")
            print(f"      STARTING STREAM: {s.upper()}")
            print(f"==========================================")
            # Smart check: if domains or translation are already fully uploaded, proceed immediately
            if resume and not args.dry_run:
                try:
                    from huggingface_hub import HfApi
                    check_api = HfApi(token=args.token or get_hf_token())
                    _, _, existing_cnt = get_existing_shards_info(check_api, args.repo, s)
                    if s == "domains" and existing_cnt >= 2:
                        print(f"[all] Stream 'domains' already fully uploaded ({existing_cnt} shards). Skipping to next...")
                        continue
                    if s == "translation" and existing_cnt >= 13:
                        print(f"[all] Stream 'translation' already fully uploaded ({existing_cnt} shards). Skipping to next...")
                        continue
                except Exception as e_chk:
                    print(f"[warn] Existing check error for {s}: {e_chk}")

            run_collector(
                stream_name=s,
                repo_id=args.repo,
                shard_mb=args.shard_mb,
                shard_rows=args.shard_rows,
                resume=resume,
                skip_rows=args.skip_rows,
                limit=args.limit,
                dry_run=args.dry_run,
                token=args.token
            )
    else:
        run_collector(
            stream_name=args.stream,
            repo_id=args.repo,
            shard_mb=args.shard_mb,
            shard_rows=args.shard_rows,
            resume=resume,
            skip_rows=args.skip_rows,
            limit=args.limit,
            dry_run=args.dry_run,
            token=args.token
        )


if __name__ == "__main__":
    main()

	"""
	Kaggle Cloud Data Collector for ViuMini-MoE-242M.
	Streams massive open datasets, cleans with stable blake2b hash dedup,
	chunks into ~150MB Parquet shards (configurable via --shard_mb, default 150),
	uploads to Hugging Face dataset repo `ViuAI/viu-mini-raw-pretrain`, and immediately deletes
	local temporary files to keep Kaggle disk usage < 2GB.

	Features:
	- Smart Auto-Resume: Automatically queries HF dataset repo, finds existing shards,
	starts numbering at (max_shard_id + 1), and fast-skips already processed rows.
	- ~150MB Shards: keeps HF viewer + DataLoader stable (old 1GB shards caused 1.6GB
	english files + viewer HTTP 500). Row cap 100k also enforced.
	- Zero Sub-Sampling: Ingests full corpora without artificial truncation.

	Usage on Kaggle:
	python data/scripts/kaggle_collector.py --stream all --repo ViuAI/viu-mini-raw-pretrain --shard_mb 150
	python data/scripts/kaggle_collector.py --stream hindi --repo ViuAI/viu-mini-raw-pretrain --shard_mb 150
	"""
	import argparse
	import gc
	import gzip
	import hashlib
	import itertools
	import json
	import os
	from pathlib import Path
	import re
	import sys
	import tempfile
	import time

	DEV = re.compile(r'[\u0900-\u097F]')
	LAT = re.compile(r'[A-Za-z]')


	def get_hf_token():
	"""Retrieve HF token from Kaggle secrets, Colab secrets, env var, or user input."""
	# 1. Check Kaggle secrets
	try:
	from kaggle_secrets import UserSecretsClient
	tok = UserSecretsClient().get_secret("HF_TOKEN")
	if tok:
	return tok
	except Exception:
	pass
	# 2. Check Google Colab secrets
	try:
	from google.colab import userdata
	tok = userdata.get("HF_TOKEN")
	if tok:
	return tok
	except Exception:
	pass
	# 3. Check environment variable
	tok = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")
	if tok:
	return tok
	return None


	def stable_hash(s: str) -> int:
	"""Stable 64-bit hash (PYTHONHASHSEED-independent). hash() is randomized per run,
	so dedup was unstable across Kaggle restarts -> duplicates. Use blake2b."""
	return int.from_bytes(hashlib.blake2b(s.encode("utf-8"), digest_size=8).digest(), "big")


	def dev_ratio(s: str) -> float:
	dev = len(DEV.findall(s))
	lat = len(LAT.findall(s))
	tot = dev + lat
	return dev / max(tot, 1)


	def is_valid_line(s: str, lang: str) -> bool:
	if len(s) < 10 or len(s.split()) < 3:
	return False
	if lang == "bilingual":
	return len(s) >= 15
	r = dev_ratio(s)
	if lang == "hinglish":
	# Must have Latin letters, Devanagari ratio <= 30%
	return bool(LAT.search(s)) and r <= 0.3
	elif lang == "hindi":
	# Must have Devanagari script
	return bool(DEV.search(s))
	elif lang == "english":
	# Must have Latin, zero Devanagari
	return bool(LAT.search(s)) and not bool(DEV.search(s))
	return True


	def apply_skip(ds, skip_rows: int, source_name: str):
	"""Applies fast skipping or falls back to itertools.islice."""
	if skip_rows <= 0:
	return ds
	try:
	ds = ds.skip(skip_rows)
	print(f"[resume] Fast-skipped {skip_rows:,} rows from {source_name}.")
	return ds
	except Exception as e:
	print(f"[resume] Fast-skip failed ({e}), using iterator slice for {source_name}...")
	return itertools.islice(ds, skip_rows, None)


	def stream_hindi_data(skip_rows: int = 0):
	"""Streams full IndicCorpV2 hin_Deva (152M rows) and Sangraha."""
	try:
	from datasets import load_dataset
	print("[source] Loading ai4bharat/IndicCorpV2 (indiccorp_v2, hin_Deva) streaming...")
	ds = load_dataset("ai4bharat/IndicCorpV2", "indiccorp_v2", split="hin_Deva", streaming=True)
	ds = apply_skip(ds, skip_rows, "IndicCorpV2")
	for r in ds:
	txt = (r.get("text") or "").strip()
	if txt:
	yield txt, "indiccorp_v2"
	except Exception as e:
	print(f"[warn] IndicCorpV2 stream error: {e}, falling back to sangraha...")
	try:
	from datasets import load_dataset
	ds = load_dataset("ai4bharat/sangraha", "verified", split="hin", streaming=True)
	ds = apply_skip(ds, skip_rows, "sangraha")
	for r in ds:
	txt = (r.get("text") or "").strip()
	if txt:
	yield txt, "sangraha"
	except Exception as e2:
	print(f"[error] Sangraha stream error: {e2}")


	def stream_english_data(skip_rows: int = 0):
	"""Streams Fineweb-edu / TinyStories educational slices."""
	from datasets import load_dataset
	print("[source] Loading HuggingFaceFW/fineweb-edu (sample-10BT or train) streaming...")
	try:
	ds = load_dataset("HuggingFaceFW/fineweb-edu", split="train", streaming=True)
	ds = apply_skip(ds, skip_rows, "fineweb-edu")
	for r in ds:
	txt = (r.get("text") or "").strip()
	if txt:
	yield txt, "fineweb-edu"
	except Exception as e:
	print(f"[warn] fineweb-edu fallback to TinyStories: {e}")
	ds = load_dataset("roneneldan/TinyStories", split="train", streaming=True)
	ds = apply_skip(ds, skip_rows, "tinystories")
	for r in ds:
	txt = (r.get("text") or "").strip()
	if txt:
	yield txt, "tinystories"


	def stream_hinglish_data(skip_rows: int = 0):
	"""Streams sk-community/romanized_hindi and conversational corpora."""
	from datasets import load_dataset
	print("[source] Loading sk-community/romanized_hindi streaming...")
	try:
	ds = load_dataset("sk-community/romanized_hindi", split="train", streaming=True)
	ds = apply_skip(ds, skip_rows, "romanized_hindi")
	for r in ds:
	txt = (r.get("Transliterated Hindi") or r.get("hi_rom") or r.get("text") or "").strip()
	if txt:
	yield txt, "romanized_hindi"
	except Exception as e:
	print(f"[warn] sk-community/romanized_hindi stream error: {e}")


	def stream_domain_data():
	"""Streams curated 26-pillar knowledge from collect_domains."""
	try:
	from data.scripts.collect_domains import generate_domain_samples
	for rec in generate_domain_samples():
	yield rec["text"], rec["source"], rec["lang"], rec["domain"]
	except ImportError:
	try:
	from collect_domains import generate_domain_samples
	for rec in generate_domain_samples():
	yield rec["text"], rec["source"], rec["lang"], rec["domain"]
	except Exception as e:
	print(f"[warn] domain generator import failed: {e}")


	def stream_translation_data(skip_rows: int = 0):
	"""Streams parallel translation data from cfilt/iitb-english-hindi & Romanized Hindi."""
	from datasets import load_dataset
	print("[source] Loading cfilt/iitb-english-hindi parallel translation streaming...")
	try:
	ds = load_dataset("cfilt/iitb-english-hindi", split="train", streaming=True)
	ds = apply_skip(ds, skip_rows, "iitb-english-hindi")
	for r in ds:
	trans = r.get("translation") or {}
	en = (trans.get("en") or "").strip()
	hi = (trans.get("hi") or "").strip()
	if en and hi and len(en) >= 10 and len(hi) >= 10:
	text = f"English: {en}\nHindi: {hi}"
	yield text, "iitb-english-hindi", "bilingual", "translation"
	except Exception as e:
	print(f"[warn] iitb translation stream fallback: {e}")


	def get_existing_shards_info(api, repo_id: str, stream_name: str):
	"""
	Checks HF repo for existing shards in {stream_name}/.
	Returns (next_shard_id, max_shard_id, existing_count).
	"""
	if not api:
	return 0, -1, 0
	try:
	files = api.list_repo_files(repo_id=repo_id, repo_type="dataset")
	prefix = f"{stream_name}/train-"
	shard_ids = []
	for f in files:
	if f.startswith(prefix) and (f.endswith(".parquet") or f.endswith(".jsonl.gz")):
	base = f[len(prefix):]
	num_str = base.split(".")[0]
	if num_str.isdigit():
	shard_ids.append(int(num_str))
	if shard_ids:
	max_id = max(shard_ids)
	return max_id + 1, max_id, len(shard_ids)
	except Exception as e:
	print(f"[warn] Could not inspect existing shards for {stream_name}: {e}")
	return 0, -1, 0


	def write_and_upload_shard(records, shard_id, stream_name, repo_id, api, dry_run=False, tmp_dir=None):
	"""
	Writes records to a compressed parquet shard, uploads to HF,
	and immediately removes the local file and frees memory.
	"""
	if tmp_dir is None:
	tmp_dir = Path(tempfile.gettempdir()) / "viu_shards"
	tmp_path = Path(tmp_dir) / stream_name
	tmp_path.mkdir(parents=True, exist_ok=True)

	file_name = f"train-{shard_id:05d}.parquet"
	out_file = tmp_path / file_name
	remote_path = f"{stream_name}/{file_name}"

	try:
	import pandas as pd
	df = pd.DataFrame(records)
	df.to_parquet(out_file, index=False, compression="snappy")
	del df
	gc.collect()
	except Exception as e:
	# Fallback to jsonl.gz if parquet/pandas is unavailable
	file_name = f"train-{shard_id:05d}.jsonl.gz"
	out_file = tmp_path / file_name
	remote_path = f"{stream_name}/{file_name}"
	with gzip.open(out_file, "wt", encoding="utf-8") as f:
	for r in records:
	f.write(json.dumps(r, ensure_ascii=False) + "\n")
	gc.collect()

	size_mb = out_file.stat().st_size / (1024 * 1024)
	print(f"[shard] Serialized {file_name} with {len(records):,} rows ({size_mb:.2f} MB)")

	if dry_run:
	print(f"[dry-run] Skipped upload of {out_file} to {repo_id}/{remote_path}")
	out_file.unlink(missing_ok=True)
	return True

	# Upload to HF with retry and exponential backoff + 429 rate limit handler (Overnight resilient)
	retries = 999
	success = False
	for attempt in range(1, retries + 1):
	try:
	api.upload_file(
	path_or_fileobj=str(out_file),
	path_in_repo=remote_path,
	repo_id=repo_id,
	repo_type="dataset",
	commit_message=f"Upload {remote_path} ({len(records):,} rows, {size_mb:.1f} MB)"
	)
	print(f"[hf] Uploaded successfully: {remote_path} ({size_mb:.2f} MB)")
	success = True
	break
	except Exception as err:
	err_str = str(err)
	if "429" in err_str or "rate limit" in err_str.lower():
	wait_time = 120
	m = re.search(r"Retry after (\d+) seconds", err_str)
	if m:
	wait_time = int(m.group(1)) + 5
	print(f"[rate-limit] HF 128 commits/hr limit hit. Waiting {wait_time}s before retry ({attempt}/{retries})...")
	time.sleep(wait_time)
	else:
	print(f"[warn] Upload attempt {attempt}/{retries} failed: {err}")
	time.sleep(4 * attempt)

	# Immediately delete local temporary shard to keep Kaggle disk clean
	out_file.unlink(missing_ok=True)
	if not success:
	print(f"[error] Failed to upload {remote_path} after {retries} retries!")
	return success


	def run_collector(stream_name, repo_id, shard_mb=150, shard_rows=100_000,
	resume=True, skip_rows=None, limit=None, dry_run=False, token=None):
	from huggingface_hub import HfApi

	api = None
	if not dry_run:
	if not token:
	token = get_hf_token()
	if not token:
	print("[error] HF token missing! Set HF_TOKEN environment variable or Kaggle secret.")
	sys.exit(1)
	api = HfApi(token=token)
	try:
	api.create_repo(repo_id=repo_id, repo_type="dataset", exist_ok=True)
	print(f"[hf] Verified dataset repo: {repo_id}")
	except Exception as e:
	print(f"[warn] create_repo warning: {e}")

	# Determine starting shard_id and rows to skip if resuming
	# FIX: old code assumed 100k rows/shard even for 1.6GB english shards (~3-4M rows),
	# causing massive duplication on resume. Now shard_rows is enforced, so
	# skip = (max_id+1) * shard_rows is correct. For legacy 1GB shards, require
	# explicit --skip_rows or manual backfill (see docs/HF_REPAIR_RUNBOOK.md).
	shard_id = 0
	actual_skip = 0
	if resume and not dry_run:
	next_id, max_id, count = get_existing_shards_info(api, repo_id, stream_name)
	if count > 0:
	shard_id = next_id
	if skip_rows is None:
	actual_skip = (max_id + 1) * shard_rows
	print(f"[resume] Found {count} existing shards in '{stream_name}' (up to train-{max_id:05d}.parquet).")
	print(f"[resume] Resuming at shard_id={shard_id} (train-{shard_id:05d}.parquet), auto-skipping ~{actual_skip:,} rows.")
	print(f"[resume] NOTE: legacy 1GB shards (english/hindi-big) used variable rows. "
	f"If resuming those, pass explicit --skip_rows. See docs/HF_REPAIR_RUNBOOK.md")
	else:
	actual_skip = skip_rows
	print(f"[resume] Resuming at shard_id={shard_id}, skipping specified {actual_skip:,} rows.")
	else:
	actual_skip = skip_rows or 0
	else:
	actual_skip = skip_rows or 0

	# Target byte calculation: Snappy parquet ~2.7x for hindi text, ~1.6x for english.
	# FIX: old default 1000MB caused 1.6GB english shards -> HF viewer 500 + OOM.
	# New default 150MB + hard row cap 100k keeps shards 30-150MB, viewer-stable.
	# Conservative ratio 2.0 to avoid overshoot on english.
	target_raw_bytes = int(shard_mb * 1024 * 1024 * 2.0)
	if stream_name == "domains":
	# Domain samples are curated; use smaller chunk threshold to flush promptly
	target_raw_bytes = int(min(shard_mb, 150) * 1024 * 1024 * 2.7)

	print(f"\n=== Starting Collector: stream={stream_name}, repo={repo_id} ===")
	print(f" Target shard size: ~{shard_mb} MB compressed \| Start shard: train-{shard_id:05d}.parquet \| Skip: {actual_skip:,} rows")

	seen_hashes = set()
	current_records = []
	current_raw_bytes = 0
	total_kept = 0
	start_time = time.time()
	last_log_time = start_time

	# Determine generator
	if stream_name == "hindi":
	gen = ((txt, src, "hindi", "general_hindi") for txt, src in stream_hindi_data(skip_rows=actual_skip))
	elif stream_name == "english":
	gen = ((txt, src, "english", "general_english") for txt, src in stream_english_data(skip_rows=actual_skip))
	elif stream_name == "hinglish":
	gen = ((txt, src, "hinglish", "general_hinglish") for txt, src in stream_hinglish_data(skip_rows=actual_skip))
	elif stream_name == "domains":
	gen = stream_domain_data()
	elif stream_name == "translation":
	gen = stream_translation_data(skip_rows=actual_skip)
	else:
	raise ValueError(f"Unknown stream: {stream_name}")

	for item in gen:
	txt, src, lang, dom = item[0], item[1], item[2], item[3]
	s = txt.strip()
	if not is_valid_line(s, lang):
	continue

	# In-flight stable-hash deduplication (FIX: hash() was seed-randomized)
	# NOTE: seen_hashes grows with total rows. For >10M-row runs on Kaggle,
	# memory can blow up. We cap at 5M entries then clear oldest via sampling
	# reset to keep RAM <2GB (small duplicate leak acceptable vs OOM).
	h = stable_hash(s)
	if h in seen_hashes:
	continue
	seen_hashes.add(h)
	if len(seen_hashes) > 5_000_000:
	print("[dedup] seen_hashes >5M, clearing to bound RAM (minor dup leak accepted)")
	seen_hashes.clear()

	row_bytes = len(s.encode('utf-8'))
	current_records.append({
	"text": s,
	"lang": lang,
	"source": src,
	"domain": dom
	})
	current_raw_bytes += row_bytes
	total_kept += 1

	# Periodic status log every 100,000 rows or every 30 seconds
	now = time.time()
	if total_kept % 100_000 == 0 or (now - last_log_time) >= 30.0:
	elapsed = max(now - start_time, 1)
	rate = total_kept / elapsed
	est_parquet_mb = (current_raw_bytes / 2.7) / (1024 * 1024)
	print(f"[progress] {total_kept:,} rows \| Shard buffer: {len(current_records):,} rows (~{est_parquet_mb:.1f} MB est. parquet) \| {rate:.1f} rows/s")
	last_log_time = now

	# Flush condition: Target MB reached OR explicit row cap reached
	hit_bytes = current_raw_bytes >= target_raw_bytes
	hit_rows = (shard_rows is not None and len(current_records) >= shard_rows)
	if hit_bytes or hit_rows:
	write_and_upload_shard(current_records, shard_id, stream_name, repo_id, api, dry_run=dry_run)
	shard_id += 1
	current_records = []
	current_raw_bytes = 0
	gc.collect()

	if limit and total_kept >= limit:
	print(f"[limit] Reached limit of {limit:,} rows.")
	break

	# Flush remaining records
	if current_records:
	write_and_upload_shard(current_records, shard_id, stream_name, repo_id, api, dry_run=dry_run)
	shard_id += 1
	current_records = []
	gc.collect()

	elapsed = max(time.time() - start_time, 1)
	print(f"=== Collector Finished: {total_kept:,} total rows in this session ({elapsed:.1f}s, {total_kept/elapsed:.1f} rows/s) ===")


	def main():
	parser = argparse.ArgumentParser(description="Kaggle High-Throughput Cloud Stream Collector")
	parser.add_argument("--stream", type=str, required=True,
	choices=["hindi", "english", "hinglish", "domains", "translation", "all"],
	help="Data stream to collect (or 'all' for sequential run)")
	parser.add_argument("--repo", type=str, default="ViuAI/viu-mini-raw-pretrain", help="HF dataset repo ID")
	parser.add_argument("--shard_mb", type=int, default=150, help="Target compressed Parquet shard size in MB (default 150 MB, viewer-stable; old 1000 MB caused 1.6GB shards + viewer 500)")
	parser.add_argument("--shard_rows", type=int, default=100000, help="Row cap per shard (default 100k; whichever hits first: bytes or rows)")
	parser.add_argument("--no-resume", action="store_true", help="Do not auto-resume; start from shard 0 and row 0")
	parser.add_argument("--skip_rows", type=int, default=None, help="Explicitly skip N rows before streaming")
	parser.add_argument("--limit", type=int, default=None, help="Optional max rows to collect")
	parser.add_argument("--dry-run", action="store_true", help="Run without uploading to HF")
	parser.add_argument("--token", type=str, default=None, help="HF API token")
	args = parser.parse_args()

	resume = not args.no_resume

	if args.stream == "all":
	# Prioritize Hinglish and English so they accumulate immediately alongside Hindi
	streams = ["domains", "translation", "hinglish", "english", "hindi"]
	for s in streams:
	print(f"\n==========================================")
	print(f" STARTING STREAM: {s.upper()}")
	print(f"==========================================")
	# Smart check: if domains or translation are already fully uploaded, proceed immediately
	if resume and not args.dry_run:
	try:
	from huggingface_hub import HfApi
	check_api = HfApi(token=args.token or get_hf_token())
	_, _, existing_cnt = get_existing_shards_info(check_api, args.repo, s)
	if s == "domains" and existing_cnt >= 2:
	print(f"[all] Stream 'domains' already fully uploaded ({existing_cnt} shards). Skipping to next...")
	continue
	if s == "translation" and existing_cnt >= 13:
	print(f"[all] Stream 'translation' already fully uploaded ({existing_cnt} shards). Skipping to next...")
	continue
	except Exception as e_chk:
	print(f"[warn] Existing check error for {s}: {e_chk}")

	run_collector(
	stream_name=s,
	repo_id=args.repo,
	shard_mb=args.shard_mb,
	shard_rows=args.shard_rows,
	resume=resume,
	skip_rows=args.skip_rows,
	limit=args.limit,
	dry_run=args.dry_run,
	token=args.token
	)
	else:
	run_collector(
	stream_name=args.stream,
	repo_id=args.repo,
	shard_mb=args.shard_mb,
	shard_rows=args.shard_rows,
	resume=resume,
	skip_rows=args.skip_rows,
	limit=args.limit,
	dry_run=args.dry_run,
	token=args.token
	)


	if __name__ == "__main__":
	main()