app.py

import gradio as gr
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
import torch
import fitz
from langdetect import detect
import matplotlib.pyplot as plt
from collections import Counter
import re
import os

theme = gr.themes.Ocean(
    primary_hue="sky",
    secondary_hue="indigo",
    neutral_hue="neutral",
)

MODEL_PATH = "hideosnes/Bart-T2T-Distill_GildaBot"

def load_model():
    tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
    model = AutoModelForSeq2SeqLM.from_pretrained(
        MODEL_PATH,
        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
    )
    return tokenizer, model

tokenizer, model = load_model()

def summarize(file, text, style, length):
    text_input = ""
    if file is not None:
        if file.name.endswith(".pdf"): # and hasattr(file, 'read'):
            try:
                with fitz.open(file.name) as doc: #with fitz.open(stream=file.read(), filetype="pdf") as doc:
                    text_input = " ".join([page.get_text() for page in doc])
            except:
                # Fallback if fitz can't read the file
                text_input = "Error: Could not read PDF file"
        elif file.name.endswith(".txt"):
            try:
                with open(file.name, 'r', encoding='utf-8') as f:
                    text_input = f.read()
            except:
                text_input = str(file)
    elif text:
        text_input = text
    # If the input text is empty or contains only whitespace, 
    # return early with a user message and placeholder values.
    if not text_input.strip():
    # Gradio expects the summarize() function to always return the same number of outputs,
    # so we return a message for the first output (the summary box) and None for the rest.
    # This ensures the UI remains consistent and doesn't break if the input is empty.
        return "Maybe try uploading a file or typing some proper text?", None, None, None, None, None
    
    # Language detection
    try:
        lang_code = detect(text_input)
    except:
        lang_code = "en"

    # Length
    max_token, min_token = ( 
        (100, 85) if length == "Short" else 
        (200, 185) if length == "Medium" else 
        (300, 285)
    )

    # System prompt based on language and style
    prompt_map = {
        "en": {
            "Precise": "Summarize concisely in a dry, scientific and academic style:",
            "Sloppy": "Summarize comprehensively for a child:",
            "Keywords": "Write only a list of important keywords of main ideas separated by a comma:",
        }#, <-- don't forget the comma!!!!!
        #"foo": { "precise": "another language or prompt map could go here"}
    }
    prompt = prompt_map.get(lang_code, prompt_map["en"])[style] + " " + text_input

    # Summarization
    # Custom tokenizer: create a class with encode/decode methods following the HuggingFace 
    # tokenizer interface, or use the PreTrainedTokenizerFast class with your own 
    # vocab and pre-tokenization rules.

    # Note: 1024 tokens typically correspond to about 750–800 English words, 
    # depending on the tokenizer and language. ---------------------------------------------- (!)
    # Make sure to display this token/word information to the user in the app UI for clarity.

    # Note: "pyTorchTensor" is not a valid TensorType, use one of ['pt', 'tf', 'np', 'jax', 'mlx']
    inputs = tokenizer.encode(prompt, return_tensors="pt", truncation=True, max_length=1024)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model.to(device)
    inputs = inputs.to(device)
    # the generated summary is not text yet but a tensor-array of token IDs
    summary_ids = model.generate(
        inputs,
        max_length=max_token,
        min_length=min_token,
        length_penalty=2.0,
        num_beams=4,
        early_stopping=True,
        no_repeat_ngram_size=3
    )
    summary = tokenizer.decode(summary_ids[0], skip_special_tokens=True)
    # removing the prompt text from the summary
    prompt_text = prompt_map.get(lang_code, prompt_map["en"])[style]
    if summary.startswith(prompt_text):
        summary = summary[len(prompt_text):].strip()
    if text_input in summary:
        summary = summary.replace(text_input, "").strip()

    # These lines calculate and store the word count of the original text, 
    # the word count of the summary, and the percentage reduction in length after summarization.
    # Note: len() is a built-in Python function that returns the number of items in an object.
    original_len = len(text_input.split())
    summary_len = len(summary.split())
    reduction = 100 - (summary_len / original_len * 100)

    # Extracting the 5 most frequent words (longer than 3 characters)
    # from the summary, treating them as keywords.
    words = re.findall(r'\w+', summary.lower())
    keyword_counts = Counter(words).most_common(5)
    keywords = [kw for kw, _ in keyword_counts if len(kw) > 5]  # Changed from > 3 to > 5

    # Plot
    fig, ax = plt.subplots()
    ax.bar(
        ["Original", "Summary"], 
        [original_len, summary_len], 
        color=["indigo", "sky"]
    )
    ax.set_ylabel("Word Count")

    return summary, ", ".join(keywords), original_len, summary_len, f"{reduction:.2f}%", fig

with gr.Blocks(theme=theme) as demo:
    gr.Markdown("Summarizer (T2T)")

    with gr.Row(): #left column
        with gr.Column():
            file = gr.File(label="Upload a PDF file")
            text = gr.Textbox(label="Paste text from clipboard.", lines=10)

            with gr.Row(): # for inline horizontal layout
                style = gr.Dropdown(["Precise", "Sloppy", "Keywords"], label="Style (experimental)")
                length = gr.Dropdown(["Short", "Middle", "Long"], label="Length")
                token_info = gr.Text(label="Speed", value="750-800 words (1024 tokens) will take 4 minutes (~4t/s).")
                btn = gr.Button("Transform", variant="primary")
        
        with gr.Column(): #right column
            summary = gr.Textbox(label="Summary")
            keywords = gr.Textbox(label="Keywords (experimental)")
            original_len = gr.Number(label="Original Text Length")
            summary_len = gr.Number(label="Summary Length")
            reduction = gr.Textbox(label="Summary Efficiency")
            plot = gr.Plot(label="Summary Statistics")

    btn.click(
        summarize,
        inputs=[file, text, style, length],
        outputs=[summary, keywords, original_len, summary_len, reduction, plot]
    )

demo.launch(share=True)