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

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)

def summarize(text):
    try:
        inputs = tokenizer(
            text,
            return_tensors="pt",
            max_length=1024,
            truncation=True
        ).to(device)

        summary_ids = model.generate(
            inputs['input_ids'],
            attention_mask=inputs['attention_mask'],
            max_length=150,
            min_length=30,
            num_beams=4,
            early_stopping=True
        )

        summary = tokenizer.decode(summary_ids[0], skip_special_tokens=True)
        return summary
    except Exception as e:
        return str(e)

interface = gr.Interface(
    fn=summarize,
    inputs="text",
    outputs="text",
    title="BART Summarization",
    live=True,
    description="Enter an article to generate a summary using a fine-tuned BART model."
)

interface.launch()