modules/utils/text_processing.py

"""Text processing utilities for tokenization and n-gram generation."""

import re
from typing import List, Tuple, Dict
import nltk
from collections import Counter

# Download required NLTK data
try:
    nltk.data.find('tokenizers/punkt')
except LookupError:
    nltk.download('punkt', quiet=True)


def tokenize(text: str, lowercase: bool = True) -> List[str]:
    """Tokenize text into words.
    
    Args:
        text: Input text string
        lowercase: Whether to convert to lowercase
        
    Returns:
        List of tokens
    """
    if lowercase:
        text = text.lower()
    # Simple word tokenization
    tokens = re.findall(r'\b\w+\b', text)
    return tokens


def get_ngrams(tokens: List[str], n: int) -> List[Tuple[str, ...]]:
    """Generate n-grams from token list.
    
    Args:
        tokens: List of tokens
        n: N-gram size
        
    Returns:
        List of n-gram tuples
    """
    if n > len(tokens):
        return []
    return [tuple(tokens[i:i+n]) for i in range(len(tokens) - n + 1)]


def get_all_ngrams(tokens: List[str], max_n: int = 4) -> Dict[int, List[Tuple[str, ...]]]:
    """Generate all n-grams up to max_n.
    
    Args:
        tokens: List of tokens
        max_n: Maximum n-gram size
        
    Returns:
        Dictionary mapping n to list of n-grams
    """
    return {n: get_ngrams(tokens, n) for n in range(1, max_n + 1)}


def count_ngrams(tokens: List[str], n: int) -> Counter:
    """Count n-gram occurrences.
    
    Args:
        tokens: List of tokens
        n: N-gram size
        
    Returns:
        Counter of n-gram frequencies
    """
    ngrams = get_ngrams(tokens, n)
    return Counter(ngrams)


def find_matching_ngrams(ref_tokens: List[str], cand_tokens: List[str], n: int) -> List[Tuple[Tuple[str, ...], int, int]]:
    """Find matching n-grams between reference and candidate.
    
    Args:
        ref_tokens: Reference tokens
        cand_tokens: Candidate tokens
        n: N-gram size
        
    Returns:
        List of (ngram, ref_pos, cand_pos) tuples
    """
    ref_ngrams = get_ngrams(ref_tokens, n)
    cand_ngrams = get_ngrams(cand_tokens, n)
    
    matches = []
    for i, ref_ngram in enumerate(ref_ngrams):
        for j, cand_ngram in enumerate(cand_ngrams):
            if ref_ngram == cand_ngram:
                matches.append((ref_ngram, i, j))
    
    return matches


def get_skip_bigrams(tokens: List[str], max_skip: int = 2) -> List[Tuple[str, str]]:
    """Generate skip-bigrams (pairs with gaps).
    
    Args:
        tokens: List of tokens
        max_skip: Maximum number of tokens to skip
        
    Returns:
        List of skip-bigram tuples
    """
    skip_bigrams = []
    for i in range(len(tokens)):
        for j in range(i + 1, min(i + max_skip + 2, len(tokens))):
            skip_bigrams.append((tokens[i], tokens[j]))
    return skip_bigrams


def highlight_matching_ngrams(text: str, matches: List[Tuple[Tuple[str, ...], int, int]], 
                               color: str = "#90EE90") -> str:
    """Highlight matching n-grams in text with HTML.
    
    Args:
        text: Original text
        matches: List of (ngram, start_pos, end_pos) - end_pos is length in tokens
        color: Highlight color
        
    Returns:
        HTML string with highlights
    """
    tokens = tokenize(text, lowercase=False)
    if not matches:
        return text
    
    # Sort matches by position
    matches = sorted(matches, key=lambda x: x[1])
    
    # Build highlighted text
    result = []
    last_end = 0
    
    for ngram, start, _ in matches:
        # Add text before match
        if start > last_end:
            result.append(" ".join(tokens[last_end:start]))
        
        # Add highlighted match
        ngram_text = " ".join(ngram)
        result.append(f'<span style="background-color: {color}; padding: 2px; border-radius: 3px;">{ngram_text}</span>')
        last_end = start + len(ngram)
    
    # Add remaining text
    if last_end < len(tokens):
        result.append(" ".join(tokens[last_end:]))
    
    return " ".join(result)