Initial commit

back_translator.py

@@ -0,0 +1,56 @@
+from transformers import (
+    MarianMTModel,
+    MarianTokenizer,
+)
+
+class BackTranslator:
+    """
+    Perform Back-translation with pivot language. English -> German -> Spanish -> English
+    Args:
+        source_lang: Source language (default: 'en')
+        pivot_lang: Pivot language (default: 'de')
+        target_lang: Target language (default: 'es')
+    Examples:
+        back_translator = BackTranslator()
+        back_translator.back_translate("Hello, how are you?")
+    """
+    def __init__(self, source_lang='en', pivot_lang='de', target_lang='es'):
+        # Forward (English to German)
+        pivot_forward_model_name = f'Helsinki-NLP/opus-mt-{source_lang}-{pivot_lang}'
+        self.tokenizer_pivot_forward = MarianTokenizer.from_pretrained(pivot_forward_model_name)
+        self.model_pivot_forward = MarianMTModel.from_pretrained(pivot_forward_model_name)
+        
+        # Pivot translation model (German to Spanish) 
+        pivot_backward_model_name = f'Helsinki-NLP/opus-mt-{pivot_lang}-{target_lang}'
+        self.tokenizer_pivot_backward = MarianTokenizer.from_pretrained(pivot_backward_model_name)
+        self.model_pivot_backward = MarianMTModel.from_pretrained(pivot_backward_model_name)
+
+        # Backward (Spanish to English)
+        backward_model_name = f'Helsinki-NLP/opus-mt-{target_lang}-{source_lang}'
+        self.tokenizer_backward = MarianTokenizer.from_pretrained(backward_model_name)
+        self.model_backward = MarianMTModel.from_pretrained(backward_model_name)
+
+    def back_translate(self, text):
+        """
+        Perform back-translation through German and Spanish to generate text variations.
+        Args:
+            text (str): The input text to be back-translated
+            
+        Returns:
+            str: The back-translated text
+        """
+        # 1. English to German
+        encoded_pivot = self.tokenizer_pivot_forward([text], padding=True, truncation=True, return_tensors='pt')
+        generated_pivot = self.model_pivot_forward.generate(**encoded_pivot)
+        pivot_text = self.tokenizer_pivot_forward.batch_decode(generated_pivot, skip_special_tokens=True)[0]
+
+        # 2. German to Spanish
+        encoded_back_pivot = self.tokenizer_pivot_backward([pivot_text], padding=True, truncation=True, return_tensors='pt')
+        retranslated_pivot = self.model_pivot_backward.generate(**encoded_back_pivot)
+        tgt_text_back = self.tokenizer_pivot_backward.batch_decode(retranslated_pivot, skip_special_tokens=True)[0]
+
+        # 3. Spanish to English
+        encoded_back = self.tokenizer_backward([tgt_text_back], padding=True, truncation=True, return_tensors='pt')
+        retranslated = self.model_backward.generate(**encoded_back)
+        src_text = self.tokenizer_backward.batch_decode(retranslated, skip_special_tokens=True)[0]
+        return src_text