prepare_cpt_data.py

#!/usr/bin/env python3
"""
Kannada CPT Data Preparation Script
====================================
Merges all available Kannada monolingual datasets into a single streaming
corpus, mixes with English at a 3:1 (KN:EN) ratio following the SambaLingo
methodology, and saves in HuggingFace datasets format (parquet shards).

Datasets merged:
  Kannada sources (~36 GB total):
    1. ai4bharat/sangraha (synthetic, kan_Knda)     — 17.7 GB
    2. Kannada-LLM-Labs/CulturaX-Kn                  — 3.8 GB
    3. Kannada-LLM-Labs/C4-Kn                         — 2.9 GB
    4. ai4bharat/IndicCorpV2 (kan_Knda)              — large
    5. pavan-naik/kannada_corpus_1m                   — 514 MB
    6. Kannada-LLM-Labs/Wikipedia-Kn                  — 140 MB
    7. ai4bharat/sangraha (kan_Latn transliterated)  — 11.6 GB (optional)

  English sources (for 1:3 EN:KN ratio):
    8. HuggingFaceFW/fineweb (sample, sample-10BT)    — 10 GB sample
    9. wikimedia/wikipedia (20231101.en)              — for quality

Output: ./data/cpt_kannada/  (parquet shards, streaming-compatible)

Usage:
  python prepare_cpt_data.py [--max_kannada_gb 30] [--english_ratio 0.25]
  python prepare_cpt_data.py --streaming  # for low-RAM machines (no full download)
"""

import argparse
import os
import sys
import random
import logging
from itertools import chain

logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
log = logging.getLogger(__name__)


# =============================================================================
# Dataset source definitions
# =============================================================================

KANNADA_SOURCES = [
    {
        "name": "sangraha_synthetic_kan",
        "load_spec": ("ai4bharat/sangraha", "synthetic", "kan_Knda"),
        "text_col": "text",
        "size_gb": 17.7,
        "gated": False,
    },
    {
        "name": "culturax_kn",
        "load_spec": ("Kannada-LLM-Labs/CulturaX-Kn", None, None),
        "text_col": "text",
        "size_gb": 3.8,
        "gated": False,
    },
    {
        "name": "c4_kn",
        "load_spec": ("Kannada-LLM-Labs/C4-Kn", None, None),
        "text_col": "text",
        "size_gb": 2.9,
        "gated": False,
    },
    {
        "name": "indicorpv2_kan",
        "load_spec": ("ai4bharat/IndicCorpV2", None, "kan_Knda"),
        "text_col": "text",
        "size_gb": 5.0,
        "gated": False,
    },
    {
        "name": "kannada_corpus_1m",
        "load_spec": ("pavan-naik/kannada_corpus_1m", None, None),
        "text_col": "text",
        "size_gb": 0.5,
        "gated": False,
    },
    {
        "name": "wikipedia_kn",
        "load_spec": ("Kannada-LLM-Labs/Wikipedia-Kn", None, None),
        "text_col": "text",
        "size_gb": 0.14,
        "gated": False,
    },
    # Optional: transliterated Kannada (Latin script)
    {
        "name": "sangraha_synthetic_kan_latn",
        "load_spec": ("ai4bharat/sangraha", "synthetic", "kan_Latn"),
        "text_col": "text",
        "size_gb": 11.6,
        "gated": False,
        "optional": True,
    },
]

ENGLISH_SOURCES = [
    {
        "name": "fineweb_sample",
        "load_spec": ("HuggingFaceFW/fineweb", "sample-10BT", None),
        "text_col": "text",
        "size_gb": 10.0,
    },
    {
        "name": "wikipedia_en",
        "load_spec": ("wikimedia/wikipedia", "20231101.en", None),
        "text_col": "text",
        "size_gb": 2.0,
    },
]


# =============================================================================
# Text cleaning
# =============================================================================

import re

# Remove very short lines, excessive whitespace, navigation artifacts
MIN_TEXT_LENGTH = 50  # characters
MAX_TEXT_LENGTH = 100000  # truncate extremely long docs

_URL_PATTERN = re.compile(r"https?://\S+|www\.\S+")
_MULTI_NEWLINE = re.compile(r"\n{3,}")
_MULTI_SPACE = re.compile(r"[ \t]{3,}")


def clean_text(text: str) -> str:
    """Clean a single text document. Returns None if text is too short/garbage."""
    if text is None:
        return None
    text = text.strip()
    if len(text) < MIN_TEXT_LENGTH:
        return None
    if len(text) > MAX_TEXT_LENGTH:
        text = text[:MAX_TEXT_LENGTH]
    # Normalize whitespace
    text = _URL_PATTERN.sub("", text)
    text = _MULTI_NEWLINE.sub("\n\n", text)
    text = _MULTI_SPACE.sub(" ", text)
    text = text.strip()
    if len(text) < MIN_TEXT_LENGTH:
        return None
    return text


# =============================================================================
# Streaming loaders
# =============================================================================

def load_source(source, streaming=True):
    """Load a dataset source and yield (text, source_name) tuples."""
    name = source["name"]
    ds_name, config, split = source["load_spec"]
    text_col = source["text_col"]

    log.info(f"Loading {name}: dataset={ds_name}, config={config}, split={split}")

    from datasets import load_dataset, DatasetDict, IterableDatasetDict

    load_kwargs = {"path": ds_name, "streaming": streaming}
    if config:
        load_kwargs["name"] = config
    if split:
        load_kwargs["split"] = split

    try:
        ds = load_dataset(**load_kwargs)
        if isinstance(ds, (DatasetDict, IterableDatasetDict)):
            if "train" in ds:
                ds = ds["train"]
            else:
                first_split = list(ds.keys())[0]
                ds = ds[first_split]
    except Exception as e:
        log.error(f"Failed to load {name}: {e}")
        return

    count = 0
    for example in ds:
        text = example.get(text_col)
        cleaned = clean_text(text)
        if cleaned is not None:
            yield {"text": cleaned, "source": name, "lang": "kn" if source in KANNADA_SOURCES else "en"}
            count += 1
            if count % 100000 == 0:
                log.info(f"  {name}: {count:,} examples processed")

    log.info(f"  {name}: done, {count:,} valid examples")


def interleave_streams(streams, ratios, seed=42):
    """
    Interleave multiple streams according to ratios.
    streams: list of iterables
    ratios: list of floats (same length as streams)
    Yields items from all streams proportionally.
    """
    rng = random.Random(seed)
    n = len(streams)

    # Convert ratios to a probability distribution
    total = sum(ratios)
    probs = [r / total for r in ratios]

    # Use a reservoir-style round-robin with probability weighting
    iterators = [iter(s) for s in streams]
    exhausted = [False] * n
    exhausted_count = 0

    while exhausted_count < n:
        # Pick which stream to draw from
        active_indices = [i for i in range(n) if not exhausted[i]]
        if not active_indices:
            break

        active_probs = [probs[i] for i in active_indices]
        active_total = sum(active_probs)
        active_probs = [p / active_total for p in active_probs]

        choice = rng.choices(active_indices, weights=active_probs, k=1)[0]

        try:
            item = next(iterators[choice])
            yield item
        except StopIteration:
            exhausted[choice] = True
            exhausted_count += 1
            log.info(f"Stream {choice} exhausted ({exhausted_count}/{n} done)")


# =============================================================================
# Main
# =============================================================================

def main():
    parser = argparse.ArgumentParser(description="Prepare Kannada CPT corpus")
    parser.add_argument("--output_dir", default="./data/cpt_kannada",
                        help="Output directory for parquet shards")
    parser.add_argument("--max_kannada_gb", type=float, default=None,
                        help="Max Kannada data to use in GB (None = all)")
    parser.add_argument("--english_ratio", type=float, default=0.25,
                        help="Fraction of English data (0.25 = 1:3 EN:KN ratio)")
    parser.add_argument("--include_latn", action="store_true",
                        help="Include transliterated Kannada (Latin script)")
    parser.add_argument("--streaming", action="store_true", default=True,
                        help="Use streaming mode (low RAM)")
    parser.add_argument("--no_streaming", dest="streaming", action="store_false",
                        help="Download full datasets (needs more RAM)")
    parser.add_argument("--shard_size_mb", type=int, default=500,
                        help="Approximate size of each output shard in MB")
    parser.add_argument("--seed", type=int, default=42)
    args = parser.parse_args()

    os.makedirs(args.output_dir, exist_ok=True)

    # ---- Select Kannada sources ----
    kn_sources = [s for s in KANNADA_SOURCES if not s.get("optional", False)]
    if args.include_latn:
        kn_sources = KANNADA_SOURCES  # include transliterated

    log.info("=" * 70)
    log.info("Kannada CPT Data Preparation")
    log.info("=" * 70)
    log.info(f"Kannada sources ({len(kn_sources)}):")
    for s in kn_sources:
        log.info(f"  - {s['name']}: ~{s['size_gb']} GB")
    log.info(f"English ratio: {args.english_ratio:.0%}")
    log.info(f"Streaming: {args.streaming}")
    log.info(f"Output: {args.output_dir}")
    log.info("=" * 70)

    # ---- Build streams ----
    kn_streams = []
    for src in kn_sources:
        kn_streams.append(load_source(src, streaming=args.streaming))

    # For English, use FineWeb sample + Wikipedia
    en_streams = []
    for src in ENGLISH_SOURCES:
        en_streams.append(load_source(src, streaming=args.streaming))

    # ---- Interleave ----
    # SambaLingo: 1:3 EN:KN ratio → 25% English, 75% Kannada
    # But we have multiple streams per language, so distribute within each language
    kn_ratio = 1.0 - args.english_ratio
    en_ratio = args.english_ratio

    # Split the language ratio equally among its streams
    kn_ratios = [kn_ratio / len(kn_streams)] * len(kn_streams) if kn_streams else []
    en_ratios = [en_ratio / len(en_streams)] * len(en_streams) if en_streams else []

    all_streams = kn_streams + en_streams
    all_ratios = kn_ratios + en_ratios

    log.info(f"Interleaving {len(all_streams)} streams with ratios: {[f'{r:.3f}' for r in all_ratios]}")

    # ---- Write output ----
    from datasets import Dataset
    import pyarrow as pa
    import pyarrow.parquet as pq

    shard_size_bytes = args.shard_size_mb * 1024 * 1024
    current_shard = []
    current_size = 0
    shard_num = 0
    total_examples = 0

    def write_shard(shard_data, shard_idx):
        if not shard_data:
            return
        out_path = os.path.join(args.output_dir, f"shard_{shard_idx:05d}.parquet")
        table = pa.Table.from_pylist(shard_data)
        pq.write_table(table, out_path)
        log.info(f"Wrote {out_path}: {len(shard_data):,} examples")

    for item in interleave_streams(all_streams, all_ratios, seed=args.seed):
        current_shard.append(item)
        current_size += len(item["text"].encode("utf-8"))
        total_examples += 1

        if current_size >= shard_size_bytes:
            write_shard(current_shard, shard_num)
            shard_num += 1
            current_shard = []
            current_size = 0

        if total_examples % 100000 == 0:
            log.info(f"Total examples written: {total_examples:,} ({shard_num} shards)")

    # Write remaining
    write_shard(current_shard, shard_num)

    log.info("=" * 70)
    log.info(f"DONE: {total_examples:,} examples in {shard_num + 1} shards")
    log.info(f"Output: {args.output_dir}")
    log.info("=" * 70)


if __name__ == "__main__":
    main()