import gradio as gr
from transformers import BartTokenizer, BartForConditionalGeneration
import torch

# Load the fine-tuned BART model and tokenizer from the local directory
MODEL_DIR = './BART model small/model'
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

tokenizer = BartTokenizer.from_pretrained(MODEL_DIR)
model = BartForConditionalGeneration.from_pretrained(MODEL_DIR).to(device)

# Define the summarization function
def summarize(text):
    try:
        # Tokenize the input article
        inputs = tokenizer(
            text,
            return_tensors="pt",
            max_length=1024,
            truncation=True
        ).to(device)

        # Generate the summary
        summary_ids = model.generate(
            inputs['input_ids'],
            attention_mask=inputs['attention_mask'],
            max_length=150,  # Set maximum length for the summary
            min_length=30,   # Set minimum length for the summary
            num_beams=4,     # Use beam search to generate the summary
            early_stopping=True
        )

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

        return summary
    
    except Exception as e:
        return str(e)

# Create Gradio interface
# Textbox input for the article and output for the summary
interface = gr.Interface(
    fn=summarize,          # The function to summarize the article
    inputs="text",         # Input is a text box where users can input the article text
    outputs="text",        # Output is a text box displaying the summary
    title="BART Summarization",  # The title of the app
    description="Enter an article to generate a summary using a fine-tuned BART model."
)

# Launch the Gradio app
interface.launch()