import gradio as gr
from transformers import pipeline

summarizer = pipeline(task="summarization", model="t5-small")
def summarize_text(x):
    return summarizer(x)[0]["summary_text"]

translator = pipeline(task="translation", model="t5-small")
def translate_text(x):
    return translator(x)[0]["translation_text"]

classifier = pipeline(task="sentiment-analysis")
def text_classification(x):
  preds = classifier(x)
  return [pred["label"] for pred in preds] 

demo = gr.Blocks()

with demo:
    gr.Markdown("Language Modeling Tasks")
    with gr.Tabs():
        with gr.TabItem("Text Classification"):
            with gr.Row():
                tc_input = gr.Textbox(label = "Input Text")
                tc_output = gr.Textbox(label = "Output")
            tc_button = gr.Button("Classify")
        with gr.TabItem("Summarization"):
            with gr.Row():
                s_input = gr.Textbox(label = "Input Text")
                s_output = gr.Textbox(label = "Output Text")
            s_button = gr.Button("Summarize")
        with gr.TabItem("Translator"):
            with gr.Row():
                translate_input = gr.Textbox(label = "Input Text")
                translate_output = gr.Textbox(label = "Output Text")
            translate_button = gr.Button("Translate")

    tc_button.click(text_classification, inputs=tc_input, outputs=tc_output)
    s_button.click(summarize_text, inputs=s_input, outputs=s_output)
    translate_button.click(translate_text, inputs=translate_input, outputs=translate_output)
    

demo.launch()