import gradio as gr
import wikipedia
from transformers import pipeline
import os

# Setting to use the 0th GPU
os.environ["CUDA_VISIBLE_DEVICES"] = "0"


def summarize(text):
    # Setting to use the bart-large-cnn model for summarization
    summarizer = pipeline("summarization")

    # To use the t5-base model for summarization:
    # summarizer = pipeline("summarization", model="t5-base", tokenizer="t5-base", framework="tf")

    summary_text = summarizer(text, max_length=100, min_length=5, do_sample=False)[0]['summary_text']
    print(f'Length of initial text: {len(text)}')
    print(f'Length of summary: {len(summary_text)}')
    print(summary_text)
    return summary_text


def greet(name):
    return "Hello " + name.orig_name + "!"


def get_ocr():
    return ''


def search_wiki(text):
    return wikipedia.search(text)


def get_wiki(search_term):
    # text = wikipedia.summary(search_term)
    orig_text_len = len(search_term)
    text = summarize(search_term)
    sum_length = len(text)
    return [text, orig_text_len, sum_length]


# def inference(file):
    # get_ocr()
    # model = AutoModelForSeq2SeqLM.from_pretrained("sgugger/my-awesome-model")

out_sum_text = gr.Textbox(label='Summarized Text', lines=15)
out_orig_test_len = gr.Number(label='Original Text Length')
out_sum_text_len = gr.Number(label='Summarized Text Length')

iface = gr.Interface(fn=get_wiki,
                     inputs=gr.Textbox(lines=50, placeholder="Paste article here....", label='Article to Summarize'),
                     outputs=[out_sum_text, out_orig_test_len, out_sum_text_len],
                     title='Article Summary',
                     description='Paste in an article and it will be summarized.'
                     )
iface.launch()  # To create a public link, set `share=True` in `launch()`.