from util import load_model
from util import pipeline

import gradio as gr
cp_aug = 'minnehwg/finetune-newwiki-summarization-ver-augmented2'

def get_model(cp):
    checkpoint = cp
    tokenizer, model = load_model(checkpoint)
    return tokenizer, model

tokenizer, model = get_model(cp_aug)

# demo = gr.Interface(
#     fn=generate_summary,
#     inputs=gr.Textbox(lines=2, placeholder="Enter your URL..."),
#     outputs=gr.Textbox(label="Generated Text"),
#     title="Chào mừng đến với hệ thống tóm tắt của Minne >.<",
#     description="Enter the URL to summarize and click 'Submit' to generate the summary."
# )

def generate_summary(url):
    results = pipeline(url, model, tokenizer)
    summary = "\n".join(results)
    return summary

demo = gr.Interface(
    fn=generate_summary,
    inputs=gr.Textbox(lines=2, placeholder="Enter your URL..."),
    outputs=gr.Textbox(label="Generated Text"),
    title="Chào mừng đến với hệ thống tóm tắt của Minne >.<",
    description="Enter the URL to summarize and click 'Submit' to generate the summary."
)

demo.launch(share=True)