fix: probe OOM reserve 4GB + expandable_segments + natural-mode stream shuffle (Hindi-only trap)

.gitattributes

@@ -4,3 +4,5 @@
 *.onnx filter=lfs diff=lfs merge=lfs -text
 tokenizer/outputs/*.json filter=lfs diff=lfs merge=lfs -text
 data/processed/* filter=lfs diff=lfs merge=lfs -text
+data/processed_spiritual/train-p1p2_boost_sanatan_dharma_epics_upanishads-00000.parquet filter=lfs diff=lfs merge=lfs -text
+data/processed_spiritual/train-p1p2_boost_sanatan_dharma_gita_ramayana-00000.parquet filter=lfs diff=lfs merge=lfs -text

data/processed_spiritual/train-p1p2_boost_sanatan_dharma_epics_upanishads-00000.parquet

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:69bfd314261e8d7020c52aecfb7ad396f183c7a8429465ac65d548e7c13c9436
+size 1853597

data/processed_spiritual/train-p1p2_boost_sanatan_dharma_gita_ramayana-00000.parquet

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b2205306ac4cbfde6b7b25086cc3c72c8646081325fd1f27269f0c4e994ec46f
+size 8458680

data/scripts/collect_spiritual_and_boost_corpus.py

@@ -0,0 +1,515 @@
+"""
+Sanatan Dharma, Epics, Hinglish & Reasoning Data Collector
+==========================================================
+Collects authentic open-source Hugging Face datasets:
+1. Shreemad Bhagavad Gita (verses with Sanskrit & Hindi anuvad)
+2. Shri Ramcharitmanas & Valmiki Ramayana (all Kandas & shlokas with Hindi translations)
+3. Mahabharata, Upanishads, Puranas, Chanakya Niti, and Charak Samhita
+4. Hinduism Philosophy & Scripture Q&A
+5. Chatbot Arena Hinglish Conversations
+6. GSM8k Hindi Math Word Problems & CoT
+7. Indian History & GK Hindi Q&A
+
+All data is standardized into the unified 5-column schema:
+['text', 'lang', 'source', 'domain', 'safety_tag']
+and exported as clean Parquet files for upload to ViuAI/viu-mini-raw-pretrain.
+"""
+
+import os
+import sys
+import json
+import time
+import argparse
+from pathlib import Path
+import pandas as pd
+import pyarrow as pa
+import pyarrow.parquet as pq
+from huggingface_hub import HfApi, hf_hub_download
+
+# Enforce pure English logging
+print("======================================================================")
+print("Starting Sanatan Dharma & Knowledge Corpus Collector (ViuMini-MoE-242M)")
+print("======================================================================")
+
+DATA_REPO = "ViuAI/viu-mini-raw-pretrain"
+OUTPUT_DIR = Path("data/processed_spiritual")
+OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
+
+SCHEMA = pa.schema([
+    ("text", pa.string()),
+    ("lang", pa.string()),
+    ("source", pa.string()),
+    ("domain", pa.string()),
+    ("safety_tag", pa.string()),
+])
+
+def get_hf_token():
+    """Retrieve Hugging Face token securely from environment or secret stores."""
+    token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")
+    if not token:
+        try:
+            from kaggle_secrets import UserSecretsClient
+            token = UserSecretsClient().get_secret("HF_TOKEN")
+        except Exception:
+            pass
+    if not token:
+        try:
+            from google.colab import userdata
+            token = userdata.get("HF_TOKEN")
+        except Exception:
+            pass
+    return token
+
+TOKEN = get_hf_token()
+api = HfApi(token=TOKEN)
+
+
+def collect_gita_and_ramayana():
+    """Ingest Bhagavad Gita, Ramcharitmanas, and Valmiki Ramayana."""
+    print("\n[Step 1/5] Collecting Bhagavad Gita & Ramayana...")
+    records = []
+
+    # 1. Bhagavad Gita verses (OEvortex/Bhagavad_Gita)
+    try:
+        url = "https://huggingface.co/datasets/OEvortex/Bhagavad_Gita/resolve/main/bhagavad-gita.csv"
+        df = pd.read_csv(url)
+        for _, row in df.iterrows():
+            title = str(row.get("Title", "")).strip()
+            ch = row.get("Chapter", "")
+            vs = row.get("Verse", "")
+            anuvad = str(row.get("Sanskrit Anuvad", "")).strip()
+            if anuvad and len(anuvad) > 20:
+                text = f"श्रीमद्भगवद्गीता - अध्याय {ch}, श्लोक {vs}\nशीर्षक: {title}\n\nसंस्कृत श्लोक एवं हिन्दी अनुवाद:\n{anuvad}"
+                records.append({
+                    "text": text,
+                    "lang": "hi",
+                    "source": "OEvortex/Bhagavad_Gita",
+                    "domain": "sanatan_dharma",
+                    "safety_tag": "safe"
+                })
+        print(f"  - Loaded {len(records)} Bhagavad Gita verses from OEvortex/Bhagavad_Gita")
+    except Exception as exc:
+        print(f"  - Warning loading OEvortex/Bhagavad_Gita: {exc}")
+
+    # 2. Ramcharitmanas complete Kandas (adiiva/ramcharitmanas)
+    try:
+        url = "https://huggingface.co/datasets/adiiva/ramcharitmanas/resolve/main/data/train-00000-of-00001.parquet"
+        df_ram = pd.read_parquet(url)
+        count_ram = 0
+        for _, row in df_ram.iterrows():
+            kanda = str(row.get("kanda", "")).strip()
+            pg = row.get("page_number", "")
+            content = str(row.get("content", "")).strip()
+            if content and len(content) > 30:
+                text = f"श्रीरामचरितमानस ({kanda}, पृष्ठ {pg}):\n\n{content}"
+                records.append({
+                    "text": text,
+                    "lang": "hi",
+                    "source": "adiiva/ramcharitmanas",
+                    "domain": "sanatan_dharma",
+                    "safety_tag": "safe"
+                })
+                count_ram += 1
+        print(f"  - Loaded {count_ram} pages from adiiva/ramcharitmanas (all 7 Kandas)")
+    except Exception as exc:
+        print(f"  - Warning loading adiiva/ramcharitmanas: {exc}")
+
+    # 3. Valmiki Ramayana Shlokas with Hindi translations (thebharatai/Indian_Epics_Mythology)
+    try:
+        vr_path = hf_hub_download(
+            repo_id="thebharatai/Indian_Epics_Mythology",
+            filename="Valmiki_Ramayan_Shlokas.jsonl",
+            repo_type="dataset",
+            token=TOKEN
+        )
+        count_vr = 0
+        with open(vr_path, "r", encoding="utf-8") as f:
+            for line in f:
+                if not line.strip():
+                    continue
+                item = json.loads(line)
+                shloka = str(item.get("input_text", "")).strip()
+                meaning = str(item.get("output_text", "")).strip()
+                if shloka and meaning:
+                    text = f"वाल्मीकि रामायण श्लोक:\n{shloka}\n\nहिन्दी भावार्थ एवं संदर्भ:\n{meaning}"
+                    records.append({
+                        "text": text,
+                        "lang": "hi",
+                        "source": "thebharatai/Indian_Epics_Mythology/Valmiki_Ramayan",
+                        "domain": "sanatan_dharma",
+                        "safety_tag": "safe"
+                    })
+                    count_vr += 1
+        print(f"  - Loaded {count_vr} Valmiki Ramayana shlokas with Hindi translations")
+    except Exception as exc:
+        print(f"  - Warning loading Valmiki_Ramayan_Shlokas: {exc}")
+
+    # 4. Gita Mantras (thebharatai/Indian_Epics_Mythology)
+    try:
+        url_gm = "https://huggingface.co/datasets/thebharatai/Indian_Epics_Mythology/resolve/main/MERGED_GITA%20MANTRAS.csv"
+        df_gm = pd.read_csv(url_gm)
+        count_gm = 0
+        for _, row in df_gm.iterrows():
+            title = str(row.get("title", "")).strip()
+            mantra = str(row.get("mantra", "")).strip()
+            num = row.get("number", "")
+            if mantra and len(mantra) > 20:
+                text = f"भगवद्गीता मंत्र ({title}, संख्या: {num}):\n\n{mantra}"
+                records.append({
+                    "text": text,
+                    "lang": "hi",
+                    "source": "thebharatai/Indian_Epics_Mythology/Gita_Mantras",
+                    "domain": "sanatan_dharma",
+                    "safety_tag": "safe"
+                })
+                count_gm += 1
+        print(f"  - Loaded {count_gm} Gita mantras")
+    except Exception as exc:
+        print(f"  - Warning loading Gita Mantras: {exc}")
+
+    out_file = OUTPUT_DIR / "train-p1p2_boost_sanatan_dharma_gita_ramayana-00000.parquet"
+    table = pa.Table.from_pylist(records, schema=SCHEMA)
+    pq.write_table(table, out_file, compression="snappy")
+    print(f"[Done Step 1] Saved {len(records):,} rows to {out_file} ({out_file.stat().st_size / (1024**2):.2f} MB)")
+    return out_file, len(records)
+
+
+def collect_epics_and_ancient_wisdom():
+    """Ingest Mahabharata, Upanishads, Puranas, Chanakya Niti, and Charak Samhita."""
+    print("\n[Step 2/5] Collecting Mahabharata, Upanishads & Ancient Wisdom...")
+    records = []
+
+    # 1. Upanishads (thebharatai/Indian_Epics_Mythology)
+    try:
+        url_up = "https://huggingface.co/datasets/thebharatai/Indian_Epics_Mythology/resolve/main/MERGED_upanishads.csv"
+        df_up = pd.read_csv(url_up)
+        for _, row in df_up.iterrows():
+            title = str(row.get("title", "")).strip()
+            mantra = str(row.get("mantra", "")).strip()
+            num = row.get("number", "")
+            if mantra and len(mantra) > 20:
+                text = f"उपनिषद् दर्शन ({title}, मंत्र संख्या {num}):\n\n{mantra}"
+                records.append({
+                    "text": text,
+                    "lang": "hi",
+                    "source": "thebharatai/Indian_Epics_Mythology/Upanishads",
+                    "domain": "sanatan_dharma",
+                    "safety_tag": "safe"
+                })
+        print(f"  - Loaded Upanishads mantras")
+    except Exception as exc:
+        print(f"  - Warning loading Upanishads: {exc}")
+
+    # 2. Chanakya Niti (thebharatai/Indian_Epics_Mythology)
+    try:
+        url_cn = "https://huggingface.co/datasets/thebharatai/Indian_Epics_Mythology/resolve/main/CHANKYA%20NITI_cleaned.csv"
+        df_cn = pd.read_csv(url_cn)
+        for _, row in df_cn.iterrows():
+            ch = row.get("chapter", "")
+            v_id = row.get("verse_id", "")
+            txt = str(row.get("text", "")).strip()
+            if txt and len(txt) > 20:
+                text = f"चाणक्य नीति (अध्याय {ch}, सूत्र {v_id}):\n\n{txt}"
+                records.append({
+                    "text": text,
+                    "lang": "hi",
+                    "source": "thebharatai/Indian_Epics_Mythology/Chanakya_Niti",
+                    "domain": "indian_wisdom",
+                    "safety_tag": "safe"
+                })
+        print(f"  - Loaded Chanakya Niti verses")
+    except Exception as exc:
+        print(f"  - Warning loading Chanakya Niti: {exc}")
+
+    # 3. Charak Samhita (Ancient Ayurvedic Science)
+    try:
+        url_cs = "https://huggingface.co/datasets/thebharatai/Indian_Epics_Mythology/resolve/main/CIHARAK%20SAMHITA_cleaned.csv"
+        df_cs = pd.read_csv(url_cs)
+        for _, row in df_cs.iterrows():
+            sec = row.get("section", "")
+            ch = row.get("chapter", "")
+            v_id = row.get("verse_id", "")
+            txt = str(row.get("text", "")).strip()
+            if txt and len(txt) > 20:
+                text = f"चरक संहिता आयुर्वेद दर्शन ({sec}, अध्याय {ch}, सूत्र {v_id}):\n\n{txt}"
+                records.append({
+                    "text": text,
+                    "lang": "hi",
+                    "source": "thebharatai/Indian_Epics_Mythology/Charak_Samhita",
+                    "domain": "ayurveda_wisdom",
+                    "safety_tag": "safe"
+                })
+        print(f"  - Loaded Charak Samhita Ayurvedic verses")
+    except Exception as exc:
+        print(f"  - Warning loading Charak Samhita: {exc}")
+
+    # 4. Mahabharat dialogues & episodes
+    try:
+        url_mb = "https://huggingface.co/datasets/thebharatai/Indian_Epics_Mythology/resolve/main/Mahabharat.csv"
+        df_mb = pd.read_csv(url_mb)
+        count_mb = 0
+        for _, row in df_mb.iterrows():
+            inp = str(row.get("input_text", "")).strip()
+            out = str(row.get("output_text", "")).strip()
+            if out and len(out) > 20:
+                text = f"महाभारत प्रसंग एवं संवाद:\n{inp}\n\nव्याख्या एवं विवरण:\n{out}" if inp else f"महाभारत प्रसंग:\n{out}"
+                records.append({
+                    "text": text,
+                    "lang": "hi",
+                    "source": "thebharatai/Indian_Epics_Mythology/Mahabharat",
+                    "domain": "sanatan_dharma",
+                    "safety_tag": "safe"
+                })
+                count_mb += 1
+        print(f"  - Loaded {count_mb} Mahabharat episodes from thebharatai")
+    except Exception as exc:
+        print(f"  - Warning loading Mahabharat from thebharatai: {exc}")
+
+    # 5. Puranas (venkatarajendra/puranas)
+    try:
+        url_pur = "https://huggingface.co/datasets/venkatarajendra/puranas/resolve/main/data/train-00000-of-00001.parquet"
+        df_pur = pd.read_parquet(url_pur)
+        count_pur = 0
+        for _, row in df_pur.iterrows():
+            msgs = row.get("messages", [])
+            if isinstance(msgs, list):
+                convo_parts = []
+                for m in msgs:
+                    role = m.get("role", "")
+                    content = m.get("content", "")
+                    convo_parts.append(f"<|{role}|> {content}")
+                text = "\n".join(convo_parts)
+            else:
+                text = str(msgs).strip()
+            if text and len(text) > 30:
+                records.append({
+                    "text": text,
+                    "lang": "en",
+                    "source": "venkatarajendra/puranas",
+                    "domain": "sanatan_dharma",
+                    "safety_tag": "safe"
+                })
+                count_pur += 1
+        print(f"  - Loaded {count_pur} Purana discourses from venkatarajendra/puranas")
+    except Exception as exc:
+        print(f"  - Warning loading Puranas: {exc}")
+
+    out_file = OUTPUT_DIR / "train-p1p2_boost_sanatan_dharma_epics_upanishads-00000.parquet"
+    table = pa.Table.from_pylist(records, schema=SCHEMA)
+    pq.write_table(table, out_file, compression="snappy")
+    print(f"[Done Step 2] Saved {len(records):,} rows to {out_file} ({out_file.stat().st_size / (1024**2):.2f} MB)")
+    return out_file, len(records)
+
+
+def collect_hinduism_philosophy_qa():
+    """Ingest Hinduism StackExchange Q&A (Scriptures, Deities, Temples, Rituals)."""
+    print("\n[Step 3/5] Collecting Hinduism Scripture & Philosophy Q&A...")
+    records = []
+    try:
+        qa_path = hf_hub_download(
+            repo_id="alberto98fx/hinduism-stackexchange",
+            filename="q_and_a.json",
+            repo_type="dataset",
+            token=TOKEN
+        )
+        with open(qa_path, "r", encoding="utf-8") as f:
+            data = json.load(f)
+        
+        for item in data:
+            title = str(item.get("Title", "")).strip()
+            q_body = str(item.get("q_body", "")).strip()
+            # Clean HTML tags simply
+            import re
+            clean_q = re.sub(r"<[^>]+>", " ", q_body).strip()
+            clean_q = " ".join(clean_q.split())
+
+            answers = item.get("Answers", [])
+            ans_text = ""
+            if isinstance(answers, list) and len(answers) > 0:
+                # Find accepted answer or highest score answer
+                best_ans = answers[0]
+                for ans in answers:
+                    if ans.get("Id") == item.get("AcceptedAnswerId"):
+                        best_ans = ans
+                        break
+                ans_body = str(best_ans.get("a_body", "")).strip()
+                ans_text = re.sub(r"<[^>]+>", " ", ans_body).strip()
+                ans_text = " ".join(ans_text.split())
+
+            if title and ans_text and len(ans_text) > 40:
+                text = f"<|user|> Question on Hindu Philosophy & Scriptures: {title}\n{clean_q}\n<|assistant|> <soch> Analyzing Hindu scriptures, context, and philosophical principles... </soch> {ans_text}"
+                records.append({
+                    "text": text,
+                    "lang": "en",
+                    "source": "alberto98fx/hinduism-stackexchange",
+                    "domain": "sanatan_dharma",
+                    "safety_tag": "safe"
+                })
+        print(f"  - Loaded {len(records):,} in-depth Hinduism Q&A entries")
+    except Exception as exc:
+        print(f"  - Warning loading Hinduism StackExchange: {exc}")
+
+    out_file = OUTPUT_DIR / "train-p1p2_boost_hinduism_philosophy_qa-00000.parquet"
+    table = pa.Table.from_pylist(records, schema=SCHEMA)
+    pq.write_table(table, out_file, compression="snappy")
+    print(f"[Done Step 3] Saved {len(records):,} rows to {out_file} ({out_file.stat().st_size / (1024**2):.2f} MB)")
+    return out_file, len(records)
+
+
+def collect_hinglish_and_reasoning_boost():
+    """Ingest Chatbot Arena Hinglish, Hindi GSM8k Math, and Indian History QA."""
+    print("\n[Step 4/5] Collecting Hinglish Arena, Hindi GSM8k Math & History QA...")
+    records_hinglish = []
+    records_math = []
+    records_hist = []
+
+    # 1. Chatbot Arena Hinglish Conversations (one-thing/chatbot_arena_conversations_hinglish)
+    try:
+        url_ch = "https://huggingface.co/datasets/one-thing/chatbot_arena_conversations_hinglish/resolve/main/data.csv"
+        df_ch = pd.read_csv(url_ch)
+        for _, row in df_ch.iterrows():
+            u_hing = str(row.get("user_hinglish", "")).strip()
+            a_hing = str(row.get("assistant_hinglish", "")).strip()
+            if u_hing and a_hing and len(u_hing) > 5 and len(a_hing) > 10:
+                text = f"<|user|> {u_hing}\n<|assistant|> {a_hing}"
+                records_hinglish.append({
+                    "text": text,
+                    "lang": "hinglish",
+                    "source": "one-thing/chatbot_arena_conversations_hinglish",
+                    "domain": "conversational",
+                    "safety_tag": "safe"
+                })
+        print(f"  - Loaded {len(records_hinglish):,} Hinglish conversation turns")
+    except Exception as exc:
+        print(f"  - Warning loading Chatbot Arena Hinglish: {exc}")
+
+    # 2. GSM8K Hindi Math Word Problems with Reasoning Steps
+    try:
+        url_gsm = "https://huggingface.co/datasets/bingbangboom/gsm8k-hindi/resolve/main/data/train_main-00000-of-00001.parquet"
+        df_gsm = pd.read_parquet(url_gsm)
+        for _, row in df_gsm.iterrows():
+            q = str(row.get("question", "")).strip()
+            a = str(row.get("answer", "")).strip()
+            if q and a:
+                text = f"<|user|> {q}\n<|assistant|> <soch> गणितीय समस्या का चरणबद्ध समाधान करते हैं... </soch> {a}"
+                records_math.append({
+                    "text": text,
+                    "lang": "hi",
+                    "source": "bingbangboom/gsm8k-hindi",
+                    "domain": "math_reasoning",
+                    "safety_tag": "safe"
+                })
+        print(f"  - Loaded {len(records_math):,} GSM8K Hindi math reasoning problems")
+    except Exception as exc:
+        print(f"  - Warning loading GSM8K Hindi: {exc}")
+
+    # 3. Indian History Hindi QA (kaifahmad/indian-history-hindi-QA-3.4k)
+    try:
+        url_hist = "https://huggingface.co/datasets/kaifahmad/indian-history-hindi-QA-3.4k/resolve/main/data/train-00000-of-00001.parquet"
+        df_hist = pd.read_parquet(url_hist)
+        for _, row in df_hist.iterrows():
+            q = str(row.get("question", "")).strip()
+            a = str(row.get("answer", "")).strip()
+            if q and a:
+                text = f"<|user|> भारतीय इतिहास प्रश्न: {q}\n<|assistant|> {a}"
+                records_hist.append({
+                    "text": text,
+                    "lang": "hi",
+                    "source": "kaifahmad/indian-history-hindi-QA-3.4k",
+                    "domain": "history_gk",
+                    "safety_tag": "safe"
+                })
+        print(f"  - Loaded {len(records_hist):,} Indian History QA rows")
+    except Exception as exc:
+        print(f"  - Warning loading Indian History QA: {exc}")
+
+    # Save Hinglish
+    f_hinglish = OUTPUT_DIR / "train-p1p2_boost_hinglish_arena_chat-00000.parquet"
+    t_hing = pa.Table.from_pylist(records_hinglish, schema=SCHEMA)
+    pq.write_table(t_hing, f_hinglish, compression="snappy")
+    print(f"  - Saved Hinglish: {len(records_hinglish):,} rows ({f_hinglish.stat().st_size / (1024**2):.2f} MB)")
+
+    # Save Math
+    f_math = OUTPUT_DIR / "train-p1p2_boost_hindi_math_gsm8k-00000.parquet"
+    t_math = pa.Table.from_pylist(records_math, schema=SCHEMA)
+    pq.write_table(t_math, f_math, compression="snappy")
+    print(f"  - Saved Hindi Math: {len(records_math):,} rows ({f_math.stat().st_size / (1024**2):.2f} MB)")
+
+    # Save History
+    f_hist = OUTPUT_DIR / "train-p1p2_boost_indian_history_qa-00000.parquet"
+    t_hist = pa.Table.from_pylist(records_hist, schema=SCHEMA)
+    pq.write_table(t_hist, f_hist, compression="snappy")
+    print(f"  - Saved Indian History: {len(records_hist):,} rows ({f_hist.stat().st_size / (1024**2):.2f} MB)")
+
+    return [
+        (f_hinglish, len(records_hinglish)),
+        (f_math, len(records_math)),
+        (f_hist, len(records_hist))
+    ]
+
+
+def upload_shards_to_hub(shards):
+    """Upload all newly created boost shards to Hugging Face dataset repository."""
+    print("\n[Step 5/5] Uploading Shards to Hugging Face Dataset Hub...")
+    if not TOKEN:
+        print("  - [Error] HF_TOKEN not found! Cannot push to Hub.")
+        return False
+
+    for file_path, row_count in shards:
+        hub_path = f"domains/{file_path.name}"
+        print(f"  - Uploading {file_path.name} ({row_count:,} rows) -> {DATA_REPO}/{hub_path}...")
+        try:
+            api.upload_file(
+                path_or_fileobj=str(file_path),
+                path_in_repo=hub_path,
+                repo_id=DATA_REPO,
+                repo_type="dataset",
+                commit_message=f"feat: add spiritual & reasoning boost {file_path.name} ({row_count:,} rows)"
+            )
+            print(f"    ✓ Uploaded successfully!")
+        except Exception as exc:
+            print(f"    ✗ Failed to upload {file_path.name}: {exc}")
+
+    print("\n[Complete] All Sanatan Dharma & Knowledge Boost Shards Processed!")
+    return True
+
+
+def main():
+    start_time = time.time()
+    all_shards = []
+
+    # 1. Gita & Ramayana
+    f1, count1 = collect_gita_and_ramayana()
+    all_shards.append((f1, count1))
+
+    # 2. Epics, Upanishads & Wisdom
+    f2, count2 = collect_epics_and_ancient_wisdom()
+    all_shards.append((f2, count2))
+
+    # 3. Hinduism Philosophy Q&A
+    f3, count3 = collect_hinduism_philosophy_qa()
+    all_shards.append((f3, count3))
+
+    # 4. Hinglish, Math & History Boosts
+    more_shards = collect_hinglish_and_reasoning_boost()
+    all_shards.extend(more_shards)
+
+    total_rows = sum(cnt for _, cnt in all_shards)
+    total_bytes = sum(f.stat().st_size for f, _ in all_shards)
+
+    print("\n" + "=" * 70)
+    print("Collection Summary:")
+    print(f"  Total Boost Shards: {len(all_shards)}")
+    print(f"  Total Clean Rows:   {total_rows:,}")
+    print(f"  Total Parquet Size: {total_bytes / (1024**2):.2f} MB")
+    print("=" * 70)
+
+    # 5. Upload to Hugging Face
+    upload_shards_to_hub(all_shards)
+
+    elapsed = time.time() - start_time
+    print(f"\nElapsed time: {elapsed:.2f} seconds ({elapsed / 60:.2f} minutes)")
+
+
+if __name__ == "__main__":
+    main()

docs/PROGRESS.md

@@ -343,6 +343,12 @@ Kaggle par manual file upload se version mismatch hota hai (v1 tokenizer wala ca
   - **MED**: `from data.scripts.filter_by_safety` import CWD-depend fail hota tha (silent fallback). FIX: repo root `sys.path` me dala.
   - verify: syntax OK + `train.py --smoke` 3-step PASS (loss 6.31, MoE usages OK).
   - NOTE: DDP probe+broadcast fix local CPU pe full-test nahi ho sakta — next Kaggle T4x2 run me `[auto] ... [dist broadcast r0->all]` line dikhni chahiye.
+- [x] 2026-09-24 Kaggle DDP verify RUN #1 (100 steps): **partial PASS + 2 bugs mile**:
+  - PASS: bootstrap, 2xT4, tok v2 mask_id=4, `rank 0/2 + 1/2`, `[dist broadcast r0->all]` (micro=4 accum=4), DDP wrapped synced, model 241.7M/166.2M, step 1 loss 12.85 shuru hua.
+  - **BUG OOM (exit 1)**: step-1 backward pe dono ranks CUDA OOM (13.3/14.5 GB). Root cause: probe me sirf ~2GB AdamW reserve — DDP grad buckets (~1GB) + fragmentation account nahi the. FIX: probe reserve 4→2→1GB阶梯 + `PYTORCH_ALLOC_CONF=expandable_segments:True`.
+  - **BUG data 100% Hindi**: `enforce_mix=false` + natural Hub order → `mix_rows={hinglish:0, hindi:109625, english:0}` (pehle hindi/* files). FIX: natural mode me `hf_ds.shuffle(buffer_size=100000)` (ratio enforce nahi, sirf order bias) + config `natural_shuffle_buffer: 100000`.
+  - Minor: cell HEAD check `next(list)` fail (cosmetic); `find_unused_parameters=True` warning (MoE unused experts possible — rakha).
+  - verify #2 pending after push.
 
 ## Kaise note likhna hai
 ```

model/configs/train_config.yaml

@@ -41,9 +41,12 @@ hf_max_blocks: 50000
 hf_safety: strict
 # Streaming scan limit (rows): null = unlimited (max_blocks tak). Test ke liye 20000 set karo.
 hf_max_rows: null
-# Language-mix rule OPTIONAL (default OFF = natural Hub order).
-# ON karne ke liye: enforce_mix: true (weights lang_mix se; rare lang quota-skip scan se).
+# Language-mix rule OPTIONAL (default OFF = shuffled natural order, ratio enforce NAHI).
+# OFF par Hub file-order bias hatane ke liye streaming buffer-shuffle hota hai
+# (warna pehle N rows 100% Hindi aa sakte the). Ratio chahiye to enforce_mix: true.
 enforce_mix: false
 lang_mix: {hinglish: 0.4, hindi: 0.3, english: 0.3}
 mix_scan_mult: 5
 mix_max_rows: 2000000
+# natural (enforce_mix=false) mode ka shuffle buffer (rows); 0 = off (purana buggy order)
+natural_shuffle_buffer: 100000

model/scripts/train.py

@@ -167,19 +167,21 @@ def resolve_amp(want, device, bf16_ok):
 
 
 def probe_micro_batch(model, seq_len, vocab_size, device, amp_dtype):
-    """Asli OOM-test: bada micro-batch try karo, jo fite wahi lo. Target eff-batch same rehta hai."""
+    """Asli OOM-test: bada micro-batch try karo, jo fite wahi lo. Target eff-batch same rehta hai.
+    NOTE: probe DDP wrap SE PEHLE chalta hai — isliye AdamW (~2x params) + DDP grad buckets (~1x params)
+    + fragmentation ka explicit reserve chahiye, warna wrap ke baad training step-1 pe OOM."""
     model.train()
     best = 1
-    # AdamW safety margin: reserve buffer on CUDA so probe accounts for ~2GB optimizer states
+    # Reserve: AdamW states + DDP buckets + headroom. 242M fp32 ~0.97GB/param-tensor =>
+    # ~2GB AdamW + ~1GB buckets + ~1GB frag => try 4GB, warna 2GB, warna 1GB.
     adamw_buf = None
     if device.type == "cuda":
-        try:
-            adamw_buf = torch.empty((500_000_000,), dtype=torch.float32, device=device)  # ~2.0 GB
-        except RuntimeError:
+        for n_elems in (1_000_000_000, 500_000_000, 250_000_000):  # ~4.0 / 2.0 / 1.0 GB
             try:
-                adamw_buf = torch.empty((250_000_000,), dtype=torch.float32, device=device)  # ~1.0 GB
+                adamw_buf = torch.empty((n_elems,), dtype=torch.float32, device=device)
+                break
             except RuntimeError:
-                pass
+                adamw_buf = None
     try:
         for mb in (1, 2, 4, 8, 16, 32):
             try:
@@ -316,6 +318,8 @@ def sample_text(model, tok, prompt, device, max_new=30, seq_len=1024, temperatur
 
 
 def main():
+    # Fragmentation kam karo (T4 pe large-alloc OOM ke baad reserved-but-unused dikhta tha)
+    os.environ.setdefault("PYTORCH_ALLOC_CONF", "expandable_segments:True")
     ap = argparse.ArgumentParser()
     ap.add_argument("--config", default="../configs/train_config.yaml")
     ap.add_argument("--model_config", default="../configs/viu1_moe_config.yaml")
@@ -442,7 +446,7 @@ def main():
             if R0:
                 print(f"[hf] streaming dataset {hf_dataset} split={args.hf_split}")
             hf_ds = load_dataset(hf_dataset, split=args.hf_split, streaming=True)
-            # Language-mix rule OPTIONAL (default OFF = natural Hub order).
+            # Language-mix rule OPTIONAL (default OFF = natural order).
             # enforce_mix=true par weighted block sampler chalta hai (lang_mix weights se).
             max_blocks_cfg = int(C("hf_max_blocks", 50000))
             safety_mode = str(C("hf_safety", "strict")).lower()
@@ -459,10 +463,18 @@ def main():
             mix_max_rows = int(C("mix_max_rows", 2000000))
             allow_toxic = safety_mode in ("research", "full")
             allow_uncensored = safety_mode in ("full",)
+            # NATURAL mode trap fix: Hub file-order me pehle hindi/* aati hai — bina shuffle
+            # ke pehle N rows 100% Hindi milte the (verified: hinglish=0, english=0).
+            # enforce_mix=false par streaming buffer-shuffle se languages interleave hoti hain
+            # (ratio enforce NAHI — sirf order bias hatata hai).
+            natural_shuffle = int(C("natural_shuffle_buffer", 100000))
+            if (not enforce_mix) and natural_shuffle > 0:
+                hf_ds = hf_ds.shuffle(seed=mix_seed, buffer_size=natural_shuffle)
             if R0:
                 print(f"[hf] safety={safety_mode} (toxic={allow_toxic}, uncensored={allow_uncensored}) "
                       f"max_blocks={max_blocks_cfg} max_rows={max_rows_cfg or 'unlimited'} "
-                      f"enforce_mix={enforce_mix} lang_mix={lang_mix if enforce_mix else 'OFF (natural order)'}")
+                      f"enforce_mix={enforce_mix} lang_mix={lang_mix if enforce_mix else 'OFF (shuffle-buffer order)'}"
+                      + (f" natural_shuffle={natural_shuffle}" if (not enforce_mix and natural_shuffle > 0) else ""))
             try:
                 from data.scripts.filter_by_safety import keep_row as _keep
                 def _safety_ok(r):