import gradio as gr
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
import torch

MODEL_NAME = "facebook/bart-large-cnn"

# Load model and tokenizer
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME)

device = "cuda" if torch.cuda.is_available() else "cpu"
model = model.to(device)


def summarize(text, max_length, min_length):
    if not text.strip():
        return "Please enter some text to summarize."

    inputs = tokenizer(
        text,
        return_tensors="pt",
        max_length=1024,
        truncation=True
    ).to(device)

    summary_ids = model.generate(
        inputs["input_ids"],
        num_beams=4,
        max_length=max_length,
        min_length=min_length,
        early_stopping=True
    )

    summary = tokenizer.decode(summary_ids[0], skip_special_tokens=True)
    return summary


with gr.Blocks(title="Advanced BART Summarizer") as demo:
    gr.Markdown("# Advanced BART Text Summarizer")
    gr.Markdown("Summarization using facebook/bart-large-cnn")

    input_text = gr.Textbox(
        lines=12,
        placeholder="Enter long article or paragraph here..."
    )

    with gr.Row():
        max_len = gr.Slider(50, 300, value=150, label="Max Summary Length")
        min_len = gr.Slider(20, 100, value=40, label="Min Summary Length")

    output_text = gr.Textbox(
        lines=8,
        label="Generated Summary"
    )

    summarize_btn = gr.Button("Summarize")

    summarize_btn.click(
        summarize,
        inputs=[input_text, max_len, min_len],
        outputs=output_text
    )

demo.launch()