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()