import gradio as gr
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

# Load the model and tokenizer
model_name = "t5-small"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)

def summarize_text(text):
    # Prepend "summarize:" as required by T5
    input_text = "summarize: " + text.strip()
    inputs = tokenizer.encode(
        input_text, 
        return_tensors="pt", 
        max_length=512, 
        truncation=True
    )
    # Generate summary
    summary_ids = model.generate(
        inputs, 
        max_length=150, 
        min_length=30, 
        length_penalty=2.0, 
        num_beams=4, 
        early_stopping=True
    )
    summary = tokenizer.decode(summary_ids[0], skip_special_tokens=True)
    return summary

# Create the Gradio UI
iface = gr.Interface(
    fn=summarize_text,
    inputs=gr.Textbox(lines=10, placeholder="Enter your note here..."),
    outputs="text",
    title=None,
    description=None,
    flagging_mode="never"  # ✅ Updated from allow_flagging
)

if __name__ == "__main__":  # ✅ Fixed underscores
    iface.launch()