app.py

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

# --- App Configuration ---
TITLE = "📄 AI Summarization Tool"
DESCRIPTION = """
Enter text to compare summaries from three different state-of-the-art models.
This app uses:
1.  **DistilBART-CNN:** A fast and efficient summarizer.
2.  **BART-Large-CNN:** A large, high-performance model.
3.  **Pegasus-XSUM:** A model known for more abstract, human-like summaries.
"""

# --- Example Texts ---
example_1 = """
Topic identification, interpretation, summary generation, and evaluation of the generated summary are the key challenges in text summarization. Critical tasks in extraction-based summarization are identifying key phrases in the document and using them to select sentences in the document for inclusion in the summary. In contrast, abstraction-based methods paraphrase sections of the source document. Extraction-based summarizers perform capturing key aspects of text and storing as an intermediate representation. Scoring words, utterances and sentences in text is based on that representation. Composing a summary by selecting across a range of words improves accuracy.
"""
example_2 = """
Billions of people can live much longer and much healthier lives. As death drifts farther into the distance how will our search for meaning change as we reduce the health effects of aging as a disease? Does meaning require death or does it merely require struggle of reprogramming our biology? It will require us to delve deeper into understanding the human mind and the AI mind. Do your best to make humanity the best it can be. That is who I am and that is how I live. It is what i get up in the morning to do. I believe love is leaving the world a better place and helping others be the best they can be. Is it possible to bring back people that mean something to us personally? Not just brilliant scientists like Einstein and Johnny Von Neumann but also people that we've lost. Is there a way to achieve a kind of small artificial immortality? Where you are against others in your age group as in terms of health and longevity by your age is called your inner age. You are biologically based on what we call the epigenetic clock. We know that smoking increases the speed of that clock. We also know that fasting and people who eat the right foods have a slower clock. Without that knowledge you're flying blind.
"""
example_3 = """
Nine trends explain innovations in machine learning technologies benefit you and your business in 2022. These include: No-Code Machine Learning, Cognitive AI, NLP Match Language Modeling, TinyML, AutoML, Machine Learning Operationalization Management, Full-stack Deep Learning, Generative Adversarial Networks, Unsupervised ML, and finally Reinforcement Learning.
"""
sample_texts = [[example_1], [example_2], [example_3]]


# --- Model Initialization ---
# This section loads the three summarization models.
summarizer_1 = None
summarizer_2 = None
summarizer_3 = None
model_error = None

try:
    print("Initializing models... This may take a moment.")
    device = "cuda" if torch.cuda.is_available() else "cpu"
    
    summarizer_1 = pipeline("summarization", model="sshleifer/distilbart-cnn-12-6", device=device)
    print("✅ Model 1/3 (DistilBART) loaded.")
    
    summarizer_2 = pipeline("summarization", model="facebook/bart-large-cnn", device=device)
    print("✅ Model 2/3 (BART-Large) loaded.")
    
    summarizer_3 = pipeline("summarization", model="google/pegasus-xsum", device=device)
    print("✅ Model 3/3 (Pegasus-XSUM) loaded.")

except Exception as e:
    model_error = e
    print(f"--- 🚨 Error loading models ---")
    print(f"Error: {model_error}")


# --- App Logic ---
def generate_summaries(text_to_summarize: str) -> tuple[str, str, str]:
    """
    Takes input text and returns summaries from all three models.
    """
    print("--- Button clicked. Attempting to generate summaries... ---")

    # If models failed to load during startup, display an error.
    if model_error:
        error_message = f"**A model failed to load during startup.**\n\nPlease check the console logs for details.\n\n**Error:**\n`{str(model_error)}`"
        return error_message, error_message, error_message

    if not text_to_summarize:
        return "", "", ""

    try:
        # Generate summaries from all three models.
        # The pipeline returns a list of dictionaries; we extract the summary text.
        summary_1 = summarizer_1(text_to_summarize, max_length=150, min_length=30, do_sample=False)[0]['summary_text']
        summary_2 = summarizer_2(text_to_summarize, max_length=150, min_length=30, do_sample=False)[0]['summary_text']
        summary_3 = summarizer_3(text_to_summarize, max_length=150, min_length=30, do_sample=False)[0]['summary_text']
        
        print("✅ Summaries generated successfully.")
        return summary_1, summary_2, summary_3

    except Exception as e:
        print(f"--- 🚨 Error during summarization ---")
        print(f"Error: {e}")
        runtime_error_message = f"**An error occurred during summarization.**\n\n**Error:**\n`{str(e)}`"
        return runtime_error_message, runtime_error_message, runtime_error_message


# --- Gradio Interface ---
with gr.Blocks(theme=gr.themes.Soft(), css=".gradio-container {max-width: 95% !important;}") as demo:
    gr.Markdown(f"<h1 style='text-align: center;'>{TITLE}</h1>")
    gr.Markdown(DESCRIPTION)

    with gr.Row():
        # Input Column
        with gr.Column(scale=1):
            input_textbox = gr.TextArea(
                lines=15,
                label="Enter Text to Summarize",
                placeholder="Paste your article, notes, or any other text here..."
            )
            summarize_button = gr.Button("Generate Summaries", variant="primary")

        # Output Column with Tabs
        with gr.Column(scale=1):
            with gr.Tabs():
                with gr.TabItem("DistilBART (Fast)"):
                    output_1 = gr.TextArea(label="Summary from DistilBART", interactive=False, lines=15)
                with gr.TabItem("BART-Large (Balanced)"):
                    output_2 = gr.TextArea(label="Summary from BART-Large", interactive=False, lines=15)
                with gr.TabItem("Pegasus (Abstractive)"):
                    output_3 = gr.TextArea(label="Summary from Pegasus", interactive=False, lines=15)

    gr.Examples(
        examples=sample_texts,
        inputs=input_textbox,
        label="Example Texts (Click to use)"
    )

    # Connect the button to the function
    summarize_button.click(
        fn=generate_summaries,
        inputs=input_textbox,
        outputs=[output_1, output_2, output_3],
        api_name="summarize"
    )

if __name__ == "__main__":
    demo.launch()