import gradio as gr
from transformers import PegasusTokenizer, PegasusForConditionalGeneration

# Load tokenizer and model
tokenizer = PegasusTokenizer.from_pretrained("google/pegasus-cnn_dailymail")
model = PegasusForConditionalGeneration.from_pretrained("google/pegasus-cnn_dailymail")

def summarize(text):
    inputs = tokenizer(text, return_tensors="pt", max_length=1024, truncation=True)
    summary_ids = model.generate(
        inputs.input_ids,
        max_length=1500,
        min_length=100,
        length_penalty=0.9,
        num_beams=6,
        early_stopping=True
    )
    summary = tokenizer.decode(summary_ids[0], skip_special_tokens=True)
    return summary


# Gradio interface
interface = gr.Interface(
    fn=summarize,
    inputs=[
        gr.Textbox(lines=10, label="Input Text"),
        gr.Textbox(lines=2, label="Prompt", placeholder="Enter a prompt for the summary here...")
    ],
    outputs=gr.Textbox(label="Summary"),
    examples=[
        ["The quick brown fox jumps over the lazy dog. This sentence is an English-language pangram—a sentence that contains all the letters of the English alphabet. It is often used for testing typewriters, keyboards, and computer fonts.", "Summarize this text:"],
        ["Artificial intelligence (AI) is intelligence demonstrated by machines, in contrast to the natural intelligence displayed by humans and animals. Leading AI textbooks define the field as the study of intelligent agents: any device that perceives its environment and takes actions that maximize its chance of successfully achieving its goals.", "Summarize this text:"],
        ["In data science, an anomaly is an observation that deviates significantly from other observations. Anomalies are also known as outliers, novelties, noise, deviations, and exceptions. Anomaly detection finds application in many domains, including fraud detection, network security, and fault detection.", "Summarize this text:"]
    ],
    title="Text Summarization",
    description="Generate a summary from the input text using a pre-trained PEGASUS model. Optionally, provide a prompt to guide the summarization.",
    theme="freddyaboulton/dracula_revamped"
)

interface.launch()