import numpy as np
import gradio as gr
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 
def flip_text(x):
    return x[::-1]
 
def flip_image(x):
    return np.fliplr(x)

tokenizer = AutoTokenizer.from_pretrained("suriya7/bart-finetuned-text-summarization")
model = AutoModelForSeq2SeqLM.from_pretrained("suriya7/bart-finetuned-text-summarization")

def generate_summary(text):
    print(text)
    inputs = tokenizer([text], max_length=1024, return_tensors='pt', truncation=True)
    summary_ids = model.generate(inputs['input_ids'], max_new_tokens=100, do_sample=False)
    summary = tokenizer.decode(summary_ids[0], skip_special_tokens=True)
    return summary
 
with gr.Blocks() as main:
    gr.Markdown("My AI interface")
    with gr.Tab("Single models"):
        text_to_summarize = gr.Textbox(label="Text to summarize")
        summary_output = gr.Textbox(label="Summary")
        summarize_btn = gr.Button("Summarize")


    with gr.Tab("Multi models"):
        with gr.Row():
            image_input = gr.Image()
            image_output = gr.Image()
        image_button = gr.Button("Flip")
 
    # text_button.click(flip_text, inputs=text_input, outputs=text_output)
    image_button.click(flip_image, inputs=image_input, outputs=image_output)
    summarize_btn.click(generate_summary, inputs=text_to_summarize, outputs=summary_output)
 
main.launch()