tokenizer_utils.py

"""
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
    }