from transformers import PegasusTokenizer, PegasusForConditionalGeneration
import gradio as gr

# Load model and tokenizer based on the provided configuration
tokenizer = PegasusTokenizer.from_pretrained("google/pegasus-cnn_dailymail")
model = PegasusForConditionalGeneration.from_pretrained("google/pegasus-cnn_dailymail")

def summarize(text, prompt):
    # Prepend the prompt to the input text
    inputs = tokenizer(prompt + " " + text, return_tensors="pt", max_length=1024, truncation=True)
    # Generate the summary
    summary_ids = model.generate(inputs.input_ids, max_length=128, min_length=32, length_penalty=0.8, num_beams=8, 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()