#!/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()