scripts/prepare_dataset.py

import re
import os
import json
from datasets import load_dataset, Features, Value, DatasetDict, Dataset, concatenate_datasets

def clean_hinglish(text: str) -> str:
    if not text:
        return ""
    
    # 1. Nasalization adjustments (Anusvara 'M' mapping on uppercase M only!)
    s = re.sub(r'M([bpfvBPFV])', r'm\1', text)
    s = re.sub(r'M(\s|[^bpfvBPFV]|$)', r'n\1', s)
    
    # 2. Scientific marker and case normalization
    s = s.replace('E', 'e').replace('O', 'o')
    s = s.lower()
    
    # 3. Colloquial texting Hinglish normalization
    s = re.sub(r'aa\b', 'a', s)
    s = re.sub(r'ee\b', 'i', s)
    s = re.sub(r'oo\b', 'u', s)
    
    s = s.replace('chchh', 'ch')
    s = s.replace('chch', 'ch')
    s = s.replace('chh', 'ch')
    s = s.replace('shsh', 'sh')
    s = s.replace('khkh', 'kh')
    s = s.replace('ghgh', 'gh')
    
    s = re.sub(r'c(?![h])', 'k', s)
    
    corrections = {
        "achchha": "acha",
        "achchhe": "ache",
        "achchhi": "achi",
        "rahaa": "raha",
        "kamm": "kaam",    # "काम" -> "kaam"
        "cur": "kar",      # "कर" -> "kar" (standard texting)
        "jie": "ji",       # "जी" -> "ji"
        "jee": "ji",       # "जी" -> "ji"
        "bhee": "bhi",     # "भी" -> "bhi"
        "kee": "ki",       # "की" -> "ki"
        "see": "se",       # "सी" -> "se"
        "hein": "hain",    # "हैं" -> "hain"
        "huun": "hun",    # "हूँ" -> "hoon"
    }
    
    words = s.split()
    cleaned_words = [corrections.get(w, w) for w in words]
    return " ".join(cleaned_words)

# Load user-vetted texting corrections as ground-truth overrides
unique_pairs_path = "/opt/vox/vox-hinglish-rnn/corpus/unique_word_pairs.json"
if os.path.exists(unique_pairs_path):
    print(f"Loading {unique_pairs_path} as ground-truth overrides...")
    with open(unique_pairs_path, "r", encoding="utf-8") as f:
        USER_CORRECTIONS = json.load(f)
else:
    USER_CORRECTIONS = {}

def is_valid_pair(native: str, english: str) -> bool:
    if not native or not english:
        return False
    # Limit max word length to 25 characters to filter run-on outliers and prevent CUDA OOM
    if len(native) > 25 or len(english) > 25:
        return False
    # Check if native word contains Devanagari characters
    # Devanagari range is U+0900 to U+097F
    has_devanagari = any('\u0900' <= char <= '\u097f' for char in native)
    if not has_devanagari:
        return False
    # Check if english word contains only lowercase latin characters, spaces or apostrophes
    is_latin = all(char.isalpha() or char.isspace() or char == "'" for char in english)
    return is_latin

def main():
    print("Loading ONLY 'hin.zip' from Hugging Face Aksharantar...")
    my_features = Features({
        'unique_identifier': Value('string'),
        'native word': Value('string'),
        'english word': Value('string'),
        'source': Value('string'),
        'score': Value('double')
    })
    
    aksharantar_raw = load_dataset(
        "ai4bharat/Aksharantar", 
        data_files="hin.zip",
        features=my_features,
        split="train"
    )
    print(f"Total raw Aksharantar examples loaded: {len(aksharantar_raw)}")
    
    # Extract word pairs from Aksharantar
    aksharantar_inputs = []
    aksharantar_targets = []
    
    for row in aksharantar_raw:
        native = row['native word']
        english = row['english word']
        if not native or not english:
            continue
        if native in USER_CORRECTIONS:
            target = USER_CORRECTIONS[native]
        else:
            target = clean_hinglish(english)
        
        if is_valid_pair(native, target):
            aksharantar_inputs.append(native)
            aksharantar_targets.append(target)
            
    print(f"Valid Aksharantar word pairs: {len(aksharantar_inputs)}")
    
    # Now load and process sk-community/romanized_hindi
    print("\nLoading sk-community/romanized_hindi...")
    romanized_raw = load_dataset(
        "sk-community/romanized_hindi",
        split="train"
    )
    print(f"Total raw sentences loaded from sk-community/romanized_hindi: {len(romanized_raw)}")
    
    romanized_inputs = []
    romanized_targets = []
    
    for row in romanized_raw:
        hindi_val = row.get('Hindi')
        roman_val = row.get('Transliterated Hindi')
        if not hindi_val or not roman_val:
            continue
            
        h_words = hindi_val.split()
        r_words = roman_val.split()
        
        # Word-level alignment filter (only when word count matches)
        if len(h_words) == len(r_words):
            for hw, rw in zip(h_words, r_words):
                # Clean punctuation from words
                clean_hw = ''.join(c for c in hw if '\u0900' <= c <= '\u097f')
                clean_rw = ''.join(c for c in rw if c.isalpha())
                
                if clean_hw in USER_CORRECTIONS:
                    target = USER_CORRECTIONS[clean_hw]
                else:
                    target = clean_hinglish(clean_rw)
                    
                if is_valid_pair(clean_hw, target):
                    romanized_inputs.append(clean_hw)
                    romanized_targets.append(target)
                    
    print(f"Valid aligned word pairs from sk-community/romanized_hindi: {len(romanized_inputs)}")
    
    # Combine datasets
    all_inputs = aksharantar_inputs + romanized_inputs
    all_targets = aksharantar_targets + romanized_targets
    print(f"\nTotal combined raw word pairs: {len(all_inputs)}")
    
    # Filter duplicate pairs to keep unique (native -> target) mappings
    unique_pairs = {}
    for n, e in zip(all_inputs, all_targets):
        if n in USER_CORRECTIONS:
            unique_pairs[n] = USER_CORRECTIONS[n]
        elif n not in unique_pairs:
            unique_pairs[n] = e
            
    print(f"Total unique word pairs: {len(unique_pairs)}")
    
    # Create HF dataset from unique pairs
    final_inputs = list(unique_pairs.keys())
    final_targets = list(unique_pairs.values())
    
    dataset_obj = Dataset.from_dict({
        'input_text': final_inputs,
        'target_text': final_targets
    })
    
    print("\nSplitting into Train (90%), Validation (5%), and Test (5%)...")
    # Split 10% out for test+val
    train_test = dataset_obj.train_test_split(test_size=0.10, seed=42)
    # Split the 10% test into 5% val and 5% test
    val_test = train_test['test'].train_test_split(test_size=0.5, seed=42)
    
    # Oversampling Target Boost: Inject all unique corrections oversampled 10 times into the training split!
    if USER_CORRECTIONS:
        print(f"Injecting and 10x oversampling {len(USER_CORRECTIONS)} user corrections into the train split...")
        extra_inputs = []
        extra_targets = []
        for _ in range(10):
            for n, e in USER_CORRECTIONS.items():
                extra_inputs.append(n)
                extra_targets.append(e)
                
        extra_dataset = Dataset.from_dict({
            'input_text': extra_inputs,
            'target_text': extra_targets
        })
        boosted_train = concatenate_datasets([train_test['train'], extra_dataset])
        print(f"Train split size boosted from {len(train_test['train'])} to {len(boosted_train)}!")
    else:
        boosted_train = train_test['train']
        
    final_dataset = DatasetDict({
        'train': boosted_train,
        'validation': val_test['train'],
        'test': val_test['test']
    })
    
    print("\nFinal Dataset splits:")
    print(final_dataset)
    
    output_dir = "/opt/vox/vox-hinglish-rnn/corpus/dataset"
    print(f"\nSaving processed dataset to disk at: {output_dir}")
    final_dataset.save_to_disk(output_dir)
    print("✓ Dataset preparation completed successfully!")

if __name__ == "__main__":
    main()