import gradio as gr
import os
from transformers import pipeline, AutoTokenizer

# Initialize summarizer and tokenizer
summarizer = pipeline("summarization", model="sshleifer/distilbart-cnn-12-6", tokenizer="sshleifer/distilbart-cnn-12-6")
tokenizer = AutoTokenizer.from_pretrained("sshleifer/distilbart-cnn-12-6")

def summarize_text(input_text):
    """Summarizes the given input text."""
    max_length = tokenizer.model_max_length
    inputs = tokenizer(input_text, truncation=True, max_length=max_length, return_tensors="pt")
    summary_ids = summarizer.model.generate(inputs.input_ids, max_length=50, min_length=10, do_sample=False)
    summary_text = tokenizer.decode(summary_ids[0], skip_special_tokens=True)
    return {"summary": summary_text}

def generate_summary(input):
    output = summarize_text(input)
    return output["summary"]  # Return the summary directly

gr.close_all()
demo = gr.Interface(
    fn=generate_summary,
    inputs=[gr.Textbox(label="Text to summarize", lines=6)],
    outputs=[gr.Textbox(label="Summary", lines=3)],
    title="Text Summarization",
    description="Summarize text using the 'shleifer/distilbart-cnn-12-6' language model.",
)

demo.launch()