augmentation/back_translator.py

"""Back-translation augmentation for Myanmar text.

Translates text to another language and back to create
paraphrased versions for data augmentation.
"""

import logging
from pathlib import Path
from typing import Dict, List, Optional, Tuple

logger = logging.getLogger(__name__)


class BackTranslator:
    """Back-translation augmentation using translation APIs."""
    
    def __init__(
        self,
        translator_api: Optional[object] = None,
        target_lang: str = "en",
        source_lang: str = "my",
    ):
        """
        Args:
            translator_api: Translation API instance
            target_lang: Target language for translation
            source_lang: Source language
        """
        self.translator_api = translator_api
        self.target_lang = target_lang
        self.source_lang = source_lang
    
    def translate(
        self,
        text: str,
        direction: str = "forward",
    ) -> Optional[str]:
        """Translate text.
        
        Args:
            text: Text to translate
            direction: "forward" (src->tgt) or "backward" (tgt->src)
        
        Returns:
            Translated text or None if failed
        """
        if self.translator_api is None:
            # Simulate translation for testing
            return self._simulate_translation(text, direction)
        
        try:
            if direction == "forward":
                return self.translator_api.translate(
                    text,
                    src=self.source_lang,
                    tgt=self.target_lang,
                )
            else:
                return self.translator_api.translate(
                    text,
                    src=self.target_lang,
                    tgt=self.source_lang,
                )
        except Exception as e:
            logger.error(f"Translation failed: {e}")
            return None
    
    def _simulate_translation(
        self,
        text: str,
        direction: str,
    ) -> str:
        """Simulate translation for testing without API.
        
        In real use, this would call a translation service.
        """
        # This is a placeholder - real implementation would use
        # Google Translate, DeepL, or similar API
        
        # For testing, just return the original text
        # with a marker to indicate it was "translated"
        marker = "[EN]" if direction == "forward" else "[MY]"
        return f"{marker}{text}{marker}"
    
    def back_translate(
        self,
        text: str,
    ) -> Tuple[Optional[str], Optional[str], Optional[str]]:
        """Translate text to target language and back.
        
        Args:
            text: Myanmar text
        
        Returns:
            (forward_translation, back_translation, final_text)
        """
        # Forward translation
        forward = self.translate(text, "forward")
        if forward is None:
            return None, None, None
        
        # Back translation
        back = self.translate(forward, "backward")
        if back is None:
            return forward, None, None
        
        return forward, back, back
    
    def augment_dataset(
        self,
        samples: List[Dict],
        batch_size: int = 10,
    ) -> List[Dict]:
        """Augment dataset using back-translation.
        
        Args:
            samples: List of sample dictionaries
            batch_size: Batch size for API calls
        
        Returns:
            List of augmented samples
        """
        augmented = []
        
        for i, sample in enumerate(samples):
            text = sample.get("text", "")
            
            forward, back, final = self.back_translate(text)
            
            if final and final != text:
                aug_sample = sample.copy()
                aug_sample["text"] = final
                aug_sample["forward_translation"] = forward
                aug_sample["back_translation"] = back
                aug_sample["augmentation_type"] = "back_translation"
                aug_sample["is_augmented"] = True
                augmented.append(aug_sample)
            
            if (i + 1) % batch_size == 0:
                logger.info(f"Processed {i + 1}/{len(samples)} samples")
        
        return augmented


class TranslationAugmenter:
    """Advanced translation-based augmentation."""
    
    def __init__(
        self,
        translator_api: Optional[object] = None,
        languages: Optional[List[str]] = None,
    ):
        """
        Args:
            translator_api: Translation API instance
            languages: List of intermediate languages for multi-hop translation
        """
        self.translator_api = translator_api
        self.languages = languages or ["en", "zh", "ja", "ko"]
    
    def multi_hop_translate(
        self,
        text: str,
        intermediate_langs: Optional[List[str]] = None,
    ) -> str:
        """Translate through multiple intermediate languages.
        
        Args:
            text: Text to translate
            intermediate_langs: Languages to translate through
        
        Returns:
            Final translated text
        """
        if intermediate_langs is None:
            intermediate_langs = random.sample(
                self.languages,
                k=min(2, len(self.languages))
            )
        
        current_text = text
        
        for lang in intermediate_langs:
            # Translate to intermediate language
            if self.translator_api:
                current_text = self.translator_api.translate(
                    current_text,
                    src="my",
                    tgt=lang,
                )
            
            # Translate back to Myanmar
            if self.translator_api:
                current_text = self.translator_api.translate(
                    current_text,
                    src=lang,
                    tgt="my",
                )
        
        return current_text
    
    def paraphrase_with_context(
        self,
        text: str,
        context: str,
    ) -> str:
        """Paraphrase text while maintaining context.
        
        Args:
            text: Text to paraphrase
            context: Additional context to help translation
        
        Returns:
            Paraphrased text
        """
        # Combine text with context
        combined = f"{context}: {text}"
        
        # Translate and back-translate
        translator = BackTranslator(self.translator_api)
        _, _, paraphrased = translator.back_translate(combined)
        
        return paraphrased if paraphrased else text


def create_back_translator(
    translator_api: Optional[object] = None,
    target_lang: str = "en",
) -> BackTranslator:
    """Factory function to create back translator."""
    return BackTranslator(
        translator_api=translator_api,
        target_lang=target_lang,
    )


if __name__ == "__main__":
    print("BackTranslator loaded")
    print("For production use, integrate with translation APIs like:")
    print("  - Google Cloud Translation")
    print("  - DeepL API")
    print("  - transformers.TranslationPipeline")