import gradio as gr
from summarizer import TransformerSummarizer, Summarizer

title = "Summarizer"
description = """
This is a demo of a text summarization NN - based on GPT-2, XLNet, BERT, 
works with English, Ukrainian, and Russian (and a few other languages too, these are SOTA NN after all).
"""

NN_OPTIONS_LIST = ["mean", "max", "min", "median"]
NN_LIST = ["GPT-2", "XLNet", "BERT"]


def start_fn(article_input: str, reduce_option="mean", model_type='GPT-2') -> str:
    """
    GPT-2 based solution, input full text, output summarized text
    :param model_type:
    :param reduce_option:
    :param article_input:
    :return summarized article_output:
    """
    if model_type == "GPT-2":
        GPT2_model = TransformerSummarizer(transformer_type="GPT2", transformer_model_key="gpt2-medium",
                                           reduce_option=reduce_option)
        full = ''.join(GPT2_model(article_input, min_length=60))
        return full
    elif model_type == "XLNet":
        XLNet_model = TransformerSummarizer(transformer_type="XLNet", transformer_model_key="xlnet-base-cased",
                                            reduce_option=reduce_option)
        full = ''.join(XLNet_model(article_input, min_length=60))
        return full

    elif model_type == "BERT":
        BERT_model = Summarizer(reduce_option=reduce_option)
        full = ''.join(BERT_model(article_input, min_length=60))
        return full


face = gr.Interface(fn=start_fn,
                    inputs=[gr.inputs.Textbox(lines=2, placeholder="Paste article here.", label='Input Article'),
                            gr.inputs.Dropdown(NN_OPTIONS_LIST, label="Summarize mode"),
                            gr.inputs.Dropdown(NN_LIST, label="Selected NN")],
                    outputs=gr.inputs.Textbox(lines=2, placeholder="Summarized article here.", label='Summarized '
                                                                                                     'Article'),
                    title=title,
                    description=description, )
face.launch(server_name="0.0.0.0", share=True)