app.py

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

# Load the model and tokenizer from Hugging Face hub
model_name = "iimran/SAM-TheSummariserV2"
tokenizer = BartTokenizer.from_pretrained(model_name)
model = BartForConditionalGeneration.from_pretrained(model_name)
model.eval()  # Set the model to evaluation mode

# Function to summarize the input text
def summarize(input_text):
    # Tokenize the input text with truncation (adjust max_length as needed)
    inputs = tokenizer(input_text, return_tensors="pt", truncation=True, max_length=1024)

    # Create global attention mask: assign global attention to the first token (required by LED)
    global_attention_mask = torch.zeros(inputs["input_ids"].shape, dtype=torch.long)
    global_attention_mask[:, 0] = 1

    # Generate the summary using beam search (you can adjust parameters as needed)
    summary_ids = model.generate(
        inputs["input_ids"],
        attention_mask=inputs["attention_mask"],
        global_attention_mask=global_attention_mask,
        max_length=512,
        num_beams=4,
        early_stopping=True,
    )

    # Decode the generated ids to a summary string, skipping special tokens
    summary = tokenizer.decode(summary_ids[0], skip_special_tokens=True)
    return summary

# Create a Gradio interface with a title, description, submit button, and larger input text area
iface = gr.Interface(
    fn=summarize,  # Function that handles the summarization
    inputs=gr.Textbox(
        label="Enter Text to Summarize",
        lines=10,  # Make the input area larger by increasing the number of lines
        placeholder="Paste or type the text you want to summarize here...",
    ),
    outputs=gr.Textbox(
        label="Summary", 
        lines=5,  # Adjust output area size (number of lines)
        placeholder="Summary will appear here..."
    ),
    live=False,  # Disable live updates, use the submit button instead
    allow_flagging="never",  # Optionally disable flagging
    title="SAM - The Summariser",  # Title of the page
    description="SAM is a model which will help summarize large knowledge base articles into small summaries.",  # Description of the model
)

# Launch the interface
iface.launch()