File size: 2,832 Bytes
3190e1e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 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 |