import gradio as gr
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

# Load Model
MODEL_NAME = "google/flan-t5-small"

tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME)


def summarize(text):
    if not text.strip():
        return "Please enter some text."

    # FLAN-T5 يحتاج Instruction
    input_text = "Summarize: " + text

    inputs = tokenizer(
        input_text,
        return_tensors="pt",
        truncation=True,
        max_length=512
    )

    summary_ids = model.generate(
        inputs["input_ids"],
        max_length=80,
        min_length=20,
        num_beams=4,
        early_stopping=True
    )

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

    return summary


demo = gr.Interface(
    fn=summarize,
    inputs=gr.Textbox(
        lines=12,
        placeholder="Paste your English article here..."
    ),
    outputs=gr.Textbox(
        lines=6,
        label="Generated Summary"
    ),
    title="Text Summarization using FLAN-T5 Small",
    description="Summarize English documents using Google's FLAN-T5 Small model."
)

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