ketanchaudhary88's picture
Update app.py
1bb47ca verified
raw
history blame
943 Bytes
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