augmentation/perturbator.py

"""Text perturbation for adversarial data augmentation.

Applies various perturbations to text to create challenging
training examples that improve model robustness.
"""

import random
import re
from enum import Enum
from typing import Callable, Dict, List, Optional, Tuple


class PerturbationType(str, Enum):
    """Types of text perturbations."""
    CHAR_SWAP = "char_swap"
    CHAR_DELETE = "char_delete"
    CHAR_DUPLICATE = "char_duplicate"
    WORD_SWAP = "word_swap"
    WORD_DELETE = "word_delete"
    WORD_DUPLICATE = "word_duplicate"
    SENTENCE_SHUFFLE = "sentence_shuffle"
    KEYBOARD_typo = "keyboard_typo"
    RANDOM_CASE = "random_case"


class TextPerturbator:
    """Apply perturbations to Myanmar text."""
    
    # Myanmar keyboard layout (simplified)
    KEYBOARD_LAYOUT = {
        "က": ["ခ", "ဂ"],
        "ခ": ["က", "ဂ", "ဃ"],
        "ဂ": ["က", "ခ"],
        "ဃ": ["ခ"],
        "င": ["စ", "ဆ"],
        "စ": ["က", "င", "ဆ", "ဇ"],
        "ဆ": ["င", "စ", "ဇ"],
        "ဇ": ["စ", "ဆ", "ဈ"],
        "ဈ": ["ဇ"],
        "ဉ": ["ည"],
        "ည": ["ဉ", "ဋ"],
        "ဋ": ["ည", "ဌ"],
        "ဌ": ["ဋ", "ဍ"],
        "ဍ": ["ဌ", "ဎ"],
        "ဎ": ["ဍ", "ဏ"],
        "ဏ": ["ဎ", "တ"],
        "တ": ["ဏ", "ထ", "ဒ"],
        "ထ": ["တ", "ဓ"],
        "ဓ": ["ထ", "ဒ"],
        "ဒ": ["တ", "ဓ", "န"],
        "န": ["ဒ", "ပ", "ဖ"],
        "ပ": ["န", "ဖ", "ဗ"],
        "ဖ": ["န", "ပ", "ဗ"],
        "ဗ": ["ပ", "ဖ"],
    }
    
    def __init__(self, seed: int = 42):
        random.seed(seed)
        self.perturbation_count = 0
    
    def char_swap(self, text: str, prob: float = 0.1) -> str:
        """Swap adjacent characters."""
        chars = list(text)
        for i in range(len(chars) - 1):
            if random.random() < prob:
                chars[i], chars[i + 1] = chars[i + 1], chars[i]
        return "".join(chars)
    
    def char_delete(self, text: str, prob: float = 0.05) -> str:
        """Delete random characters."""
        chars = list(text)
        result = [c for c in chars if random.random() > prob]
        return "".join(result) if result else text
    
    def char_duplicate(self, text: str, prob: float = 0.05) -> str:
        """Duplicate random characters."""
        chars = list(text)
        result = []
        for c in chars:
            result.append(c)
            if random.random() < prob:
                result.append(c)
        return "".join(result)
    
    def word_swap(self, text: str, prob: float = 0.1) -> str:
        """Swap adjacent words."""
        words = text.split()
        if len(words) < 2:
            return text
        
        for i in range(len(words) - 1):
            if random.random() < prob:
                words[i], words[i + 1] = words[i + 1], words[i]
        
        return " ".join(words)
    
    def word_delete(self, text: str, prob: float = 0.1) -> str:
        """Delete random words."""
        words = text.split()
        if len(words) < 2:
            return text
        
        result = [w for w in words if random.random() > prob]
        return " ".join(result) if result else text
    
    def word_duplicate(self, text: str, prob: float = 0.1) -> str:
        """Duplicate random words."""
        words = text.split()
        result = []
        for w in words:
            result.append(w)
            if random.random() < prob:
                result.append(w)
        return " ".join(result)
    
    def keyboard_typo(self, text: str, prob: float = 0.1) -> str:
        """Introduce keyboard typos."""
        chars = list(text)
        result = []
        
        for c in chars:
            if random.random() < prob and c in self.KEYBOARD_LAYOUT:
                # Replace with keyboard neighbor
                neighbor = random.choice(self.KEYBOARD_LAYOUT[c])
                result.append(neighbor)
            else:
                result.append(c)
        
        return "".join(result)
    
    def random_case(self, text: str, prob: float = 0.1) -> str:
        """Randomly change case of characters (for mixed scripts)."""
        # Myanmar doesn't have case, but this can affect punctuation
        chars = list(text)
        for i in range(len(chars)):
            if chars[i].isupper() and random.random() < prob:
                chars[i] = chars[i].lower()
            elif chars[i].islower() and random.random() < prob:
                chars[i] = chars[i].upper()
        return "".join(chars)
    
    def sentence_shuffle(self, text: str) -> str:
        """Shuffle sentences in multi-sentence text."""
        sentences = re.split(r'[။၊।\.\!\?]+', text)
        sentences = [s.strip() for s in sentences if s.strip()]
        
        if len(sentences) < 2:
            return text
        
        random.shuffle(sentences)
        return " ".join(sentences)
    
    def apply_perturbation(
        self,
        text: str,
        perturbation_type: PerturbationType,
        prob: float = 0.1,
    ) -> str:
        """Apply a specific perturbation.
        
        Args:
            text: Myanmar text
            perturbation_type: Type of perturbation
            prob: Probability of perturbation
        
        Returns:
            Perturbed text
        """
        if perturbation_type == PerturbationType.CHAR_SWAP:
            return self.char_swap(text, prob)
        elif perturbation_type == PerturbationType.CHAR_DELETE:
            return self.char_delete(text, prob)
        elif perturbation_type == PerturbationType.CHAR_DUPLICATE:
            return self.char_duplicate(text, prob)
        elif perturbation_type == PerturbationType.WORD_SWAP:
            return self.word_swap(text, prob)
        elif perturbation_type == PerturbationType.WORD_DELETE:
            return self.word_delete(text, prob)
        elif perturbation_type == PerturbationType.WORD_DUPLICATE:
            return self.word_duplicate(text, prob)
        elif perturbation_type == PerturbationType.KEYBOARD_typo:
            return self.keyboard_typo(text, prob)
        elif perturbation_type == PerturbationType.RANDOM_CASE:
            return self.random_case(text, prob)
        elif perturbation_type == PerturbationType.SENTENCE_SHUFFLE:
            return self.sentence_shuffle(text)
        else:
            return text
    
    def apply_random_perturbations(
        self,
        text: str,
        n_perturbations: int = 2,
        prob: float = 0.1,
    ) -> Tuple[str, List[PerturbationType]]:
        """Apply random perturbations.
        
        Args:
            text: Myanmar text
            n_perturbations: Number of perturbations to apply
            prob: Probability for each perturbation
        
        Returns:
            (perturbed_text, list_of_applied_perturbations)
        """
        perturbations = list(PerturbationType)
        applied = []
        
        current_text = text
        
        for _ in range(n_perturbations):
            pert_type = random.choice(perturbations)
            current_text = self.apply_perturbation(current_text, pert_type, prob)
            applied.append(pert_type)
        
        self.perturbation_count += 1
        
        return current_text, applied
    
    def augment_dataset(
        self,
        samples: List[Dict],
        n_perturbations: int = 2,
        prob: float = 0.1,
        n_augmentations: int = 2,
    ) -> List[Dict]:
        """Augment dataset with perturbations.
        
        Args:
            samples: List of sample dictionaries
            n_perturbations: Number of perturbations per augmentation
            prob: Probability for each perturbation
            n_augmentations: Number of augmentations per sample
        
        Returns:
            List of augmented samples
        """
        augmented = []
        
        for sample in samples:
            text = sample.get("text", "")
            
            for i in range(n_augmentations):
                aug_text, applied = self.apply_random_perturbations(
                    text,
                    n_perturbations=n_perturbations,
                    prob=prob,
                )
                
                aug_sample = sample.copy()
                aug_sample["text"] = aug_text
                aug_sample["augmentation_id"] = i
                aug_sample["perturbations"] = [p.value for p in applied]
                aug_sample["is_augmented"] = True
                augmented.append(aug_sample)
        
        return augmented


class AdversarialPerturbator:
    """Advanced adversarial perturbations targeting specific weaknesses."""
    
    def __init__(self):
        self.base_perturbator = TextPerturbator()
    
    def confuse_sentiment_keywords(
        self,
        text: str,
        keyword_replacements: Dict[str, str],
    ) -> str:
        """Replace sentiment keywords to flip or confuse sentiment.
        
        Args:
            text: Myanmar text
            keyword_replacements: Dict of keyword -> replacement
        
        Returns:
            Text with keywords replaced
        """
        for keyword, replacement in keyword_replacements.items():
            if keyword in text:
                text = text.replace(keyword, replacement, 1)  # Replace first occurrence only
        return text
    
    def add_distractors(
        self,
        text: str,
        distractors: List[str] = None,
    ) -> str:
        """Add distractor phrases to text.
        
        Args:
            text: Myanmar text
            distractors: List of distractor phrases
        
        Returns:
            Text with distractors added
        """
        if distractors is None:
            distractors = [
                "အဲ့ဒါကို",
                "ဟုတ်ကဲ့",
                "နောက်တော့",
            ]
        
        distractor = random.choice(distractors)
        words = text.split()
        
        if len(words) >= 3:
            insert_pos = random.randint(1, len(words) - 1)
            words.insert(insert_pos, distractor)
        
        return " ".join(words)
    
    def paraphrase_style(self, text: str, style: str = "formal") -> str:
        """Change text style (formal/informal).
        
        Args:
            text: Myanmar text
            style: Target style ("formal" or "informal")
        
        Returns:
            Text with changed style
        """
        style_markers = {
            "formal": {
                "add": ["သည်", "မှာ", "ကို", "ဖြင့်"],
                "remove": ["နော်", "ဟုတ်"],
            },
            "informal": {
                "add": ["နော်", "ဟုတ်"],
                "remove": ["သည်", "မှာ", "ကို", "ဖြင့်"],
            },
        }
        
        markers = style_markers.get(style, style_markers["formal"])
        
        for marker in markers.get("add", []):
            if marker not in text and random.random() < 0.3:
                words = text.split()
                insert_pos = random.randint(0, len(words))
                words.insert(insert_pos, marker)
                text = " ".join(words)
        
        for marker in markers.get("remove", []):
            if marker in text and random.random() < 0.5:
                text = text.replace(marker, "")
        
        return text


def create_perturbator(seed: int = 42) -> TextPerturbator:
    """Factory function to create perturbator."""
    return TextPerturbator(seed=seed)


if __name__ == "__main__":
    perturbator = create_perturbator()
    
    test_text = "ကျေးဇူးပါ မင်္ဂလာပါ"
    
    print(f"Original: {test_text}")
    print(f"\nRandom perturbations:")
    
    for i in range(3):
        aug, applied = perturbator.apply_random_perturbations(test_text)
        print(f"  {i+1}. {aug}")
        print(f"     Applied: {[p.value for p in applied]}")