# =========================================
# Advanced BART-based Summarizer (Gradio)
# =========================================

import gradio as gr
import torch
from transformers import BartTokenizer, BartForConditionalGeneration

# -----------------------------
# 1. Load Model & Tokenizer
# -----------------------------
MODEL_NAME = "facebook/bart-large-cnn"

tokenizer = BartTokenizer.from_pretrained(MODEL_NAME)
model = BartForConditionalGeneration.from_pretrained(MODEL_NAME)

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


# -----------------------------
# 2. Summarization Function
# -----------------------------
def summarize_text(text, max_len, min_len, num_beams):
    if not text.strip():
        return "Please enter some text."

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

    summary_ids = model.generate(
        inputs["input_ids"],
        max_length=max_len,
        min_length=min_len,
        num_beams=num_beams,
        length_penalty=2.0,
        early_stopping=True,
        no_repeat_ngram_size=3
    )

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


# -----------------------------
# 3. Gradio UI
# -----------------------------
with gr.Blocks(title="Advanced BART Summarizer") as app:

    gr.Markdown("## 🧠 Advanced BART Text Summarizer")
    gr.Markdown("Summarize long documents using Facebook BART model")

    with gr.Row():
        input_text = gr.Textbox(
            lines=15,
            placeholder="Enter your text here...",
            label="Input Text"
        )

    with gr.Row():
        max_len = gr.Slider(50, 300, value=130, step=10, label="Max Length")
        min_len = gr.Slider(10, 100, value=30, step=5, label="Min Length")
        num_beams = gr.Slider(1, 8, value=4, step=1, label="Beam Size")

    summarize_btn = gr.Button("Summarize")

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

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

# -----------------------------
# 4. Launch
# -----------------------------
app.launch()