from transformers import pipeline
import gradio as gr

# Load models
summarizer = pipeline("summarization", model="sshleifer/distilbart-cnn-12-6")
answerer = pipeline("question-answering", model="valhalla/bart-large-finetuned-squadv1")
translator = pipeline("translation", model="facebook/nllb-200-distilled-600M")
filler = pipeline("fill-mask", model="FacebookAI/roberta-large")
paraphraser = pipeline("text2text-generation", model="humarin/chatgpt_paraphraser_on_T5_base")

# NLP functions
def summarize(text, min_length, max_length):
    summary = summarizer(text, min_length=min_length, max_length=max_length)
    return summary[0]['summary_text']

def answer(context, question):
    answers = answerer(context=context, question=question, top_k=3)
    return [" ".join(answer['answer'].split("\n")) for answer in answers]

languages = ["zho_Hans (Chinese)", "spa_Latn (Spanish)", "eng_Latn (English)", "hin_Deva (Hindi)", "por_Latn (Portuguese)", "rus_Cyrl (Russian)",
             "jpn_Jpan (Japanese)", "deu_Latn (German)", "yue_Hant (Yue Chinese)", "kor_Hang (Korean)", "fra_Latn (French)", "ita_Latn (Italian)"]
def translate(text, src_lang, tgt_lang):
    src_lang = src_lang.split()[0]
    tgt_lang = tgt_lang.split()[0]
    translation = translator(text, src_lang=src_lang, tgt_lang=tgt_lang, max_length=translator.tokenizer.model_max_length)
    return translation[0]['translation_text']

def fill(text, to_fill):
    if not to_fill:
        text = text.replace("_", filler.tokenizer.mask_token)
    else:
        text = text.replace(to_fill, filler.tokenizer.mask_token)
    words = filler(text, top_k=3)
    return [word['token_str'].strip() for word in words]

def paraphrase(text):
    paraphrases = paraphraser(text, num_beams=3, num_beam_groups=3, num_return_sequences=3, diversity_penalty=3.0, max_length=paraphraser.tokenizer.model_max_length)
    return [paraphrase['generated_text'] for paraphrase in paraphrases]

# Build demo
with gr.Blocks() as demo:
    gr.HTML("<center><h1>NLP Toolbox</h1></center>")
    with gr.Tabs():
        with gr.TabItem("Summarization"):
            text = gr.Textbox(label="text", placeholder="Enter text here...", lines=8)
            with gr.Accordion("set summary length", open=False):
                with gr.Row():
                    with gr.Column():
                        min_length = gr.Slider(label="minimum length", minimum=50, maximum=1000, step=10, value=100)
                    with gr.Column():
                        max_length = gr.Slider(label="maximum length", minimum=50, maximum=1000, step=10, value=800)
            output = gr.Textbox(label="summary", lines=3)
            submit = gr.Button("Summarize")
            submit.click(summarize, inputs=[text, min_length, max_length], outputs=output)
        with gr.TabItem("Question Answering"):
            context = gr.Textbox(label="context", placeholder="Enter text here...", lines=8)
            question = gr.Textbox(label="question", placeholder="Enter question here...", lines=1)
            output1 = gr.Textbox(label="answer no.1", lines=1)
            output2 = gr.Textbox(label="answer no.2", lines=1)
            output3 = gr.Textbox(label="answer no.3", lines=1)
            submit = gr.Button("Answer")
            submit.click(answer, inputs=[context, question], outputs=[output1, output2, output3])
        with gr.TabItem("Translation"):
            text = gr.Textbox(label="text", placeholder="Enter text here...", lines=8)
            with gr.Row():
                with gr.Column():
                    src_lang = gr.Dropdown(languages, label="source language")
                with gr.Column():
                    tgt_lang = gr.Dropdown(languages, label="target language")
            output = gr.Textbox(label="translation", lines=8)
            submit = gr.Button("Translate")
            submit.click(translate, inputs=[text, src_lang, tgt_lang], outputs=output)
        with gr.TabItem("Fill-Mask"):
            text = gr.Textbox(label="text", placeholder="Enter text here...", lines=6)
            gr.Markdown("Please use the \"_\" symbol to represent the blank.")
            to_fill = gr.Textbox(label="word to replace", placeholder="Enter word here...", lines=1)
            gr.Markdown("If you are filling a blank, please leave the cell above blank.")
            output1 = gr.Textbox(label="1st option", lines=1)
            output2 = gr.Textbox(label="2nd option", lines=1)
            output3 = gr.Textbox(label="3rd option", lines=1)
            submit = gr.Button("Fill/Replace")
            submit.click(fill, inputs=[text, to_fill], outputs=[output1, output2, output3])
        with gr.TabItem("Paraphrase"):
            text = gr.Textbox(label="text", placeholder="Enter text here...", lines=8)
            output1 = gr.Textbox(label="1st option", lines=8)
            output2 = gr.Textbox(label="2nd option", lines=8)
            output3 = gr.Textbox(label="3rd option", lines=8)
            submit = gr.Button("Paraphrase")
            submit.click(paraphrase, inputs=text, outputs=[output1, output2, output3])

demo.launch(share=True)