import gradio as gr
import torch
from transformers import T5ForConditionalGeneration, T5TokenizerFast as T5Tokenizer

# Load the pre-trained model and tokenizer
MODEL_NAME = "t5-base"
tokenizer = T5Tokenizer.from_pretrained(MODEL_NAME)
model = T5ForConditionalGeneration.from_pretrained(MODEL_NAME, return_dict=True)

# Define a function to summarize news text
def summarize_news(news_text):
    # Preprocess the news text
    inputs = tokenizer(news_text, return_tensors="pt", max_length=512, truncation=True, padding="max_length")
    input_ids = inputs.input_ids
    attention_mask = inputs.attention_mask
    decoder_attention_mask = torch.ones_like(input_ids)

    # Generate the summary
    with torch.no_grad():
        output = model.generate(input_ids=input_ids,
                                attention_mask=attention_mask,
                                max_length=150,  # Adjust the maximum length as needed
                                num_beams=2,     # Adjust the number of beams for beam search
                                early_stopping=True)
    # Decode the summary tokens to text
    summary = tokenizer.decode(output[0], skip_special_tokens=True)
    return summary

# Create a Gradio interface for your function
input_text = gr.Textbox(lines=10, label="Input Text")
output_text = gr.Textbox(label="Summary")

gr.Interface(
    fn=summarize_news,
    inputs=input_text,
    outputs=output_text,
    title="News Summary App",
    description="Enter a news text and get its summary."
).launch()