import gradio as gr
from transformers import pipeline
from happytransformer import HappyTextToText, TTSettings

def translate(language_1, language_2, message):
  translator = pipeline(f"translation", model="Helsinki-NLP/opus-mt-{}-{}".format(language_1, language_2))
  output = translator(message)
  return output[0]['translation_text']

def summarize(text):
  summerizer = pipeline("summarization", model="facebook/bart-large-cnn")
  output = summerizer(text)
  return output[0]['summary_text']

def question(question, context):
  question_model = pipeline("question-answering", model="deepset/roberta-base-squad2")
  output = question_model({
      'question': question,
      'context' : context
  })
  return output['answer']


def grammar(user_input):
    grammar_model = HappyTextToText("T5", "vennify/t5-base-grammar-correction")
    args = TTSettings(num_beams=5, min_length=1)
    output = grammar_model.generate_text('grammar: ' + user_input, args=args)
    return output.text


def mask(text):
  mask_model = pipeline("fill-mask", model="google-bert/bert-base-uncased")
  output = mask_model(text)
  return output[0]['sequence']

def multi_model(text):

  mask_model = mask(text)

  summerize_model = summarize(mask_model)

  question_model = question('Is this text about politics?', summerize_model)

  translation_model = translate('en', 'nl', question_model)

  grammar_model = grammar(translation_model)

  return grammar_model



with gr.Blocks() as demo:
  gr.Markdown("Technology 2 - Ai interfaces")

#tab single model
  with gr.Tab("Single-Models"):
    with gr.Row(equal_height=True):

      #column left (1)
      with gr.Column(scale=1):

        choice1 = gr.Dropdown(
          choices=["nl", "fr", "en"],
          label="Select an option")

        choice2 = gr.Dropdown(
          choices=["nl", "fr", "en"],
          label="Select an option")

        textbox = gr.Textbox(label="message")

        translate_btn = gr.Button("Translate")

        translated_box = gr.Textbox()

        translate_btn.click(fn=translate,
                            inputs=[choice1, choice2, textbox],
                            outputs=translated_box )

      #column right (2)
      with gr.Column(scale=1):
        textbox_summary = gr.Textbox(label="text you want to summarize")
        textbox_summary_output = gr.Textbox()
        summary_btn = gr.Button('summarize')

        summary_btn.click(fn=summarize,
                          inputs=textbox_summary,
                          outputs=textbox_summary_output)

      with gr.Column(scale=1):
        textbox_question = gr.Textbox(label = 'question')
        textbox_context = gr.Textbox(label = 'context')
        textbox_question_output = gr.Textbox()
        question_btn = gr.Button("Get your Answer")

        question_btn.click(fn=question,
                  inputs=[textbox_question, textbox_context],
                  outputs = textbox_question_output)

      with gr.Column(scale=1):
        textbox_mask = gr.Textbox(label="Give me a sentance and use the word [MASK] to predict what will go there")
        mask_button = gr.Button("Generate")
        textbox_mask_output = gr.Textbox()

        mask_button.click(fn=mask,
                          inputs = textbox_mask,
                          outputs = textbox_mask_output)

      with gr.Column(scale=1):
        textbox_grammar = gr.Textbox(label="I will correct your grammar mistakes.")
        grammar_button = gr.Button("Fix grammar")
        textbox_grammar_output = gr.Textbox()

        grammar_button.click(fn=grammar,
                             inputs=textbox_grammar,
                             outputs=textbox_grammar_output)






#tab mulit model
  with gr.Tab("Multi-Models"):
    textblock = gr.Textbox(label="give me a sentence and replace some words with [MASK] and i will summerize and translate it for your")
    btn1 = gr.Button("START")
    response = gr.Textbox()

    btn1.click(fn=multi_model, inputs = textblock, outputs=response)


demo.launch(debug=True)