import gradio as gr
import nlpaug.augmenter.word as naw
import nlpaug.augmenter.char as nac
import nlpaug.augmenter.sentence as nas

# Function for NLP augmentation
def augment_text(text, method):
    if method == "Synonym Replacement":
        # Use BERT model for synonym replacement (contextual embeddings)
        aug = naw.ContextualWordEmbsAug(model_path='bert-base-uncased', action='substitute')
    elif method == "Word Embedding Substitution":
        # Use BERT for word embedding substitution instead of Word2Vec
        aug = naw.ContextualWordEmbsAug(model_path='bert-base-uncased', action='substitute')
    elif method == "Contextual Word Insertion":
        aug = naw.ContextualWordEmbsAug(model_path="bert-base-uncased", action="insert")
    elif method == "Back Translation":
        aug = naw.BackTranslationAug(from_model_name='facebook/wmt19-en-de', to_model_name='facebook/wmt19-de-en')
    
    augmented_text = aug.augment(text)
    return augmented_text

# Gradio Interface
def nlp_augmentor_interface(text, method):
    augmented_text = augment_text(text, method)
    return augmented_text

iface = gr.Interface(
    fn=nlp_augmentor_interface,
    inputs=[
        gr.Textbox(lines=2, placeholder="Enter sentence to augment here..."),
        gr.Radio(["Synonym Replacement", "Word Embedding Substitution", "Contextual Word Insertion", "Back Translation"], label="Augmentation Method")
    ],
    outputs="text",
    title="NLP Text Augmentation with Gradio"
)

iface.launch(share=True)