File size: 943 Bytes
1bb47ca
 
600489b
1bb47ca
 
 
600489b
1bb47ca
600489b
1bb47ca
 
600489b
1bb47ca
 
 
 
 
 
600489b
1bb47ca
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
import nlpaug.augmenter.word as naw
import nlpaug.augmenter.translator as nat

# Initialize augmenters
synonym_aug = naw.SynonymAug(aug_src='en', lang='eng')
back_translate_en_to_hi = nat.BackTranslationAug(from_model_name='Helsinki-NLP/opus-mt-en-hi', to_model_name='Helsinki-NLP/opus-mt-hi-en')

def augment_text(text, augmentation_type='synonym'):
    """
    Augment text based on the specified type.
    augmentation_type: 'synonym' for synonym replacement, 'back_translation' for back translation
    """
    if augmentation_type == 'synonym':
        return synonym_aug.augment(text)
    elif augmentation_type == 'back_translation':
        return back_translate_en_to_hi.augment(text)
    else:
        return text

# Test augmentation functions
print(augment_text("What is your address?", augmentation_type='synonym'))  # Synonym
print(augment_text("What is your address?", augmentation_type='back_translation'))  # Back translation