""" Marathi BPE Tokenizer - Core Utilities This module provides core functions for Byte Pair Encoding (BPE) tokenization optimized for Marathi (Devanagari script) text processing. """ import regex as re import unicodedata from typing import List, Dict, Tuple # Devanagari Character Definitions VOWELS = "अआइईउऊएऐओऔ" CONSONANTS = "कखगघङचछजझञटठडढणतथदधनपफबभमयरलवशषसह" MATRAS = "ािीुूेैोौृॄॢॣं:ँ" VIRAMA = "्" NUMERALS = "०१२३४५६७८९" PUNCTUATION = "।॥" def graphemes(text: str) -> List[str]: r""" Extract extended grapheme clusters from text. Uses Unicode NFC normalization and regex pattern \X to properly handle Devanagari combining characters (matras, virama, etc.). Args: text: Input text string Returns: List of grapheme cluster strings """ normalized = unicodedata.normalize("NFC", text) return re.findall(r'\X', normalized) def get_stats(ids: List[int]) -> Dict[Tuple[int, int], int]: """ Count frequency of all consecutive token pairs. Args: ids: List of token IDs Returns: Dictionary mapping (token1_id, token2_id) -> frequency count """ counts = {} for pair in zip(ids, ids[1:]): counts[pair] = counts.get(pair, 0) + 1 return counts def merge(ids: List[int], pair: Tuple[int, int], idx: int) -> List[int]: """ Replace all occurrences of a token pair with a new merged token ID. Args: ids: List of token IDs pair: Tuple of (first_token_id, second_token_id) to merge idx: New token ID to use for the merged pair Returns: New list with merged tokens """ newids = [] i = 0 while i < len(ids): # Check if current position matches the pair to merge if i < len(ids) - 1 and ids[i] == pair[0] and ids[i + 1] == pair[1]: newids.append(idx) i += 2 # Skip both tokens in the pair else: newids.append(ids[i]) i += 1 return newids def runMerges( num_merges: int, tokens: List[int], startFromId: int ) -> Tuple[Dict[Tuple[int, int], int], List[int]]: """ Run the BPE training algorithm for specified number of merge iterations. Args: num_merges: Number of merge operations to perform tokens: Initial list of token IDs (from base vocabulary) startFromId: Starting ID for newly created merged tokens Returns: Tuple of (merges_dict, final_token_ids) - merges_dict: Maps (token_pair) -> merged_token_id - final_token_ids: Token sequence after all merges """ merges = {} idx = startFromId newTokens = tokens[:] # Create a copy for i in range(num_merges): stats = get_stats(newTokens) if not stats: break # No more pairs to merge # Find most frequent pair pair = max(stats, key=stats.get) # Merge the pair newTokens = merge(newTokens, pair, idx) merges[pair] = idx idx += 1 return (merges, newTokens) def tokenToIdMap(token_to_id: Dict[str, int], tokens: List[str]) -> List[int]: """ Convert list of token strings to token IDs. Args: token_to_id: Dictionary mapping token string -> ID tokens: List of token strings Returns: List of token IDs Raises: KeyError: If any token is not in the vocabulary """ result = [] for tok in tokens: if tok not in token_to_id: raise KeyError(f"Token '{tok}' not found in vocabulary") result.append(token_to_id[tok]) return result def idToTokenMap(id_to_token: Dict[int, str], ids: List[int]) -> List[str]: """ Convert list of token IDs to token strings. Args: id_to_token: Dictionary mapping ID -> token string ids: List of token IDs Returns: List of token strings """ return [id_to_token[idx] for idx in ids] def encode( merges: Dict[Tuple[int, int], int], token_to_id: Dict[str, int], text: str ) -> List[int]: """ Encode text to token IDs using learned BPE merges. Process: 1. Split text into graphemes 2. Convert graphemes to base token IDs 3. Apply BPE merges in priority order (learned during training) Args: merges: Dictionary of learned merges (token_pair -> merged_id) token_to_id: Base vocabulary mapping (token_string -> id) text: Input text to encode Returns: List of token IDs Raises: ValueError: If text contains characters not in vocabulary """ # Extract graphemes from text grapheme_list = graphemes(text) # Validate all graphemes are in vocabulary missing_chars = [g for g in grapheme_list if g not in token_to_id] if missing_chars: unique_missing = sorted(set(missing_chars)) missing_str = ', '.join([f"'{c}'" for c in unique_missing[:10]]) # Show first 10 if len(unique_missing) > 10: missing_str += f" ... ({len(unique_missing) - 10} more)" raise ValueError( f"Characters not in vocabulary: {missing_str}. " f"This tokenizer only supports Marathi (Devanagari) characters." ) # Convert text to base token IDs tokens = tokenToIdMap(token_to_id, grapheme_list) # Apply merges iteratively while len(tokens) >= 2: stats = get_stats(tokens) # Find the pair with highest merge priority (lowest merge ID) pair = min(stats, key=lambda p: merges.get(p, float("inf"))) # If no valid merge found, stop if pair not in merges: break idx = merges[pair] tokens = merge(tokens, pair, idx) return tokens def decode(vocab: Dict[int, str], ids: List[int]) -> str: """ Decode token IDs back to original text. Args: vocab: Dictionary mapping token_id -> token_string ids: List of token IDs Returns: Reconstructed text string Raises: KeyError: If any token ID is not in vocabulary """ result = [] for idx in ids: if idx not in vocab: raise KeyError(f"Token ID {idx} not found in vocabulary during decoding") result.append(vocab[idx]) return "".join(result) def calculate_compression_ratio(original_text: str, token_ids: List[int]) -> Dict[str, float]: """ Calculate compression statistics. Args: original_text: Original input text token_ids: Encoded token IDs Returns: Dictionary with compression metrics """ grapheme_count = len(graphemes(original_text)) byte_count = len(original_text.encode('utf-8')) token_count = len(token_ids) return { 'grapheme_compression': grapheme_count / token_count if token_count > 0 else 0, 'byte_compression': byte_count / token_count if token_count > 0 else 0, 'grapheme_count': grapheme_count, 'byte_count': byte_count, 'token_count': token_count }