Doc Summarizer version 1

README.md

@@ -1,14 +1,34 @@
 ---
-title: DOC Summarizer Studyonly
-emoji: 🐢
-colorFrom: indigo
-colorTo: pink
+title: Document Summarization
+emoji: 🌖
+colorFrom: gray
+colorTo: indigo
 sdk: gradio
-sdk_version: 5.36.2
+sdk_version: 5.12.0
 app_file: app.py
-pinned: false
-license: mit
-short_description: DOC Summarizer by studyonly
+pinned: true
+license: apache-2.0
+short_description: text2text models for document summarization
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+Check out the configuration reference at <https://huggingface.co/docs/hub/spaces-config-reference>
+
+# README - Document Summarization
+
+The original demo/what this repo was built for can be found [here](https://huggingface.co/spaces/pszemraj/document-summarization)
+
+## Usage
+
+If you are using this **not** as a gradio demo on hf spaces, you can run it locally with:
+
+```bash
+python app.py --share
+```
+
+To see all the available arguments, run `python app.py --help`.
+
+## Installation
+
+```bash
+pip install -r requirements.txt
+```

aggregate.py

@@ -0,0 +1,192 @@
+"""
+aggregate.py - module for 'reducing' multiple 'summary chunks' into one
+
+an overly complicated class for legacy compatibility reasons, for usage of the
+2024 map-reduce models see hf.co/pszemraj/bart-large-summary-map-reduce#usage
+"""
+
+import logging
+import pprint as pp
+import time
+
+import torch
+from transformers import GenerationConfig, pipeline
+
+# Setting up logging
+logging.basicConfig(
+    level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
+)
+
+
+class BatchAggregator:
+    """
+    BatchAggregator is a class for aggregating text from multiple sources.
+
+    Usage:
+        from aggregate import BatchAggregator
+        aggregator = BatchAggregator()
+        agg = aggregator.infer_aggregate(["This is a test", "This is another test"])
+        print(agg)
+    """
+
+    GENERIC_CONFIG = GenerationConfig(
+        max_new_tokens=512,
+        num_beams=4,
+        early_stopping=True,
+        do_sample=False,
+        truncation=True,
+    )
+
+    def __init__(
+        self,
+        model_name: str = "pszemraj/bart-large-summary-map-reduce",
+        force_cpu: bool = False,
+        **kwargs,
+    ):
+        """
+        __init__ initializes the BatchAggregator class.
+
+        :param str model_name: model name to use, default: "pszemraj/bart-large-summary-map-reduce"
+        :param bool force_cpu: force the model to run on CPU, default: False
+        """
+        self.device = None
+        self.is_compiled = False
+        self.model_name = None
+        self.aggregator = None
+        self.force_cpu = force_cpu
+        self.logger = logging.getLogger(__name__)
+        self.init_model(model_name)
+
+    def init_model(self, model_name: str) -> None:
+        """
+        Initialize the model.
+
+        :param model_name: The name of the model to use.
+        """
+        # Free up memory
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+
+        self.logger.info(f"Setting model to {model_name}")
+        self.model_name = model_name
+        self.aggregator = self._create_pipeline(model_name)
+        self._configure_model()
+
+    def _create_pipeline(
+        self, model_name: str = "pszemraj/bart-large-summary-map-reduce"
+    ) -> pipeline:
+        """
+        _create_pipeline creates a pipeline for the model.
+
+        :param str model_name: model name to use
+        :return pipeline: the pipeline for the model
+
+        :raises Exception: if the pipeline cannot be created
+        """
+        device_map = (
+            "auto" if torch.cuda.is_available() and not self.force_cpu else "cpu"
+        )
+        try:
+            self.logger.info(
+                f"Creating pipeline with model {model_name} on device {device_map}"
+            )
+            return pipeline(
+                "text2text-generation",
+                model=model_name,
+                device_map=device_map,
+                torch_dtype=torch.float32,
+            )
+        except Exception as e:
+            self.logger.error(f"Failed to create pipeline: {e}")
+            raise
+
+    def _configure_model(self):
+        """
+        Configure the model for generation.
+        """
+        try:
+            self.aggregator.model = torch.compile(self.aggregator.model)
+            self.is_compiled = True
+        except Exception as e:
+            self.logger.warning(f"Could not compile model with Torch 2.0: {e}")
+
+        self._set_default_generation_config()
+        self.logger.info(self.aggregator.model.generation_config.to_json_string())
+
+    def _set_default_generation_config(self):
+        """
+        Set the default generation configuration for the model.
+        """
+        self.aggregator.model.generation_config.update(
+            **self.GENERIC_CONFIG.to_diff_dict()
+        )
+
+    def update_generation_config(self, **kwargs):
+        """
+        Update the generation configuration with the specified parameters.
+
+        Args:
+            **kwargs: The parameters to update in the generation configuration.
+        """
+        self.logger.info(f"Updating generation config with {pp.pformat(kwargs)}")
+        self.aggregator.model.generation_config.update(**kwargs)
+
+    def get_generation_config(self) -> dict:
+        """
+        Get the current generation configuration.
+
+        Returns:
+            dict: The current generation configuration.
+        """
+        return self.aggregator.model.generation_config.to_dict()
+
+    def update_loglevel(self, level: str = "INFO"):
+        """
+        Update the log level.
+
+        Args:
+            level (str): The log level to set. Defaults to "INFO".
+        """
+        self.logger.setLevel(level)
+
+    def infer_aggregate(
+        self,
+        text_list: list,
+        instruction: str = None,  # Kept for backward compatibility but not used
+        **kwargs,
+    ) -> str:
+        """
+        infer_aggregate - infers a consolidated summary from a list of texts.
+
+        Args:
+            text_list (list): The texts to summarize.
+            instruction (str): Not used by this model, kept for compatibility.
+            **kwargs: Additional parameters to update in the generation configuration.
+
+        Returns:
+            The generated summary.
+        """
+        joined_text = "\n\n".join(text_list)
+        if kwargs:
+            self.update_generation_config(**kwargs)
+        st = time.perf_counter()
+        self.logger.info(f"inference on {len(text_list)} texts ...")
+        result = self.aggregator(
+            joined_text,
+            generation_config=self.aggregator.model.generation_config,
+        )[0]["generated_text"]
+        self.logger.info(f"Done. runtime:\t{round(time.perf_counter() - st, 2)}s")
+        self.logger.info(
+            f"Input tokens:\t{self.count_tokens(joined_text)}. Output tokens:\t{self.count_tokens(result)}"
+        )
+        self.logger.debug(f"Generated text:\n{result}")
+
+        return result
+
+    def count_tokens(self, text: str) -> int:
+        """count the number of tokens in a text"""
+        return (
+            len(self.aggregator.tokenizer.encode(text, truncation=False, padding=False))
+            if text
+            else 0
+        )

app.py

@@ -0,0 +1,666 @@
+"""
+app.py - the main module for the gradio app for summarization
+
+Usage:
+    app.py [-h] [--share] [-m MODEL] [-nb ADD_BEAM_OPTION] [-batch TOKEN_BATCH_OPTION]
+              [-level {DEBUG,INFO,WARNING,ERROR}]
+Details:
+    python app.py --help
+
+Environment Variables:
+    USE_TORCH (str): whether to use torch (1) or not (0)
+    TOKENIZERS_PARALLELISM (str): whether to use parallelism (true) or not (false)
+Optional Environment Variables:
+    APP_MAX_WORDS (int): the maximum number of words to use for summarization
+    APP_OCR_MAX_PAGES (int): the maximum number of pages to use for OCR
+"""
+
+import argparse
+import contextlib
+import gc
+import logging
+import os
+import pprint as pp
+import random
+import time
+from pathlib import Path
+
+os.environ["USE_TORCH"] = "1"
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s [%(levelname)s] %(name)s - %(message)s",
+    datefmt="%Y-%b-%d %H:%M:%S",
+)
+
+import gradio as gr
+import nltk
+import torch
+from cleantext import clean
+from doctr.models import ocr_predictor
+
+from aggregate import BatchAggregator
+from pdf2text import convert_PDF_to_Text
+from summarize import load_model_and_tokenizer, summarize_via_tokenbatches
+from utils import (
+    contraction_aware_tokenize,
+    extract_batches,
+    load_example_filenames,
+    remove_stagnant_files,
+    remove_stopwords,
+    saves_summary,
+    textlist2html,
+    truncate_word_count,
+)
+
+_here = Path(__file__).parent
+
+nltk.download("punkt", force=True, quiet=True)
+nltk.download("popular", force=True, quiet=True)
+
+# Constants & Globals
+MODEL_OPTIONS = [
+    "BEE-spoke-data/pegasus-x-base-synthsumm_open-16k",
+    "pszemraj/long-t5-tglobal-base-sci-simplify",
+    "pszemraj/long-t5-tglobal-base-16384-book-summary",
+    "pszemraj/long-t5-tglobal-base-summary-souffle-16384-loD",
+    "pszemraj/pegasus-x-large-book_synthsumm",
+    "pszemraj/pegasus-x-large-book-summary",
+]  # models users can choose from
+BEAM_OPTIONS = [2, 3, 4]  # beam sizes users can choose from
+TOKEN_BATCH_OPTIONS = [
+    1024,
+    1536,
+    2048,
+    2560,
+    3072,
+]  # token batch sizes users can choose from
+
+SUMMARY_PLACEHOLDER = "<p><em>Output will appear below:</em></p>"
+AGGREGATE_MODEL = "pszemraj/bart-large-summary-map-reduce"  # map-reduce model
+
+# if duplicating space: uncomment this line to adjust the max words
+# os.environ["APP_MAX_WORDS"] = str(2048)  # set the max words to 2048
+# os.environ["APP_OCR_MAX_PAGES"] = str(40)  # set the max pages to 40
+# os.environ["APP_AGG_FORCE_CPU"] = str(1)  # force cpu for aggregation
+
+aggregator = BatchAggregator(
+    AGGREGATE_MODEL, force_cpu=os.environ.get("APP_AGG_FORCE_CPU", False)
+)
+
+
+def aggregate_text(
+    summary_text: str,
+    text_file: gr.File = None,
+) -> str:
+    """
+    Aggregate the text from the batches.
+
+        NOTE: you should probably include the BatchAggregator object as a fn arg if using this code
+
+    :param batches_html: The batches to aggregate, in html format
+    :param text_file: The text file to append the aggregate summary to
+    :return: The aggregate summary in html format
+    """
+    if summary_text is None or summary_text == SUMMARY_PLACEHOLDER:
+        logging.error("No text provided. Make sure a summary has been generated first.")
+        return "Error: No text provided. Make sure a summary has been generated first."
+
+    try:
+        extracted_batches = extract_batches(summary_text)
+    except Exception as e:
+        logging.info(summary_text)
+        logging.info(f"the batches html is: {type(summary_text)}")
+        return f"Error: unable to extract batches - check input: {e}"
+    if not extracted_batches:
+        logging.error("unable to extract batches - check input")
+        return "Error: unable to extract batches - check input"
+
+    out_path = None
+    if text_file is not None:
+        out_path = text_file.name  # assuming name attribute stores the file path
+
+    content_batches = [batch["content"] for batch in extracted_batches]
+    full_summary = aggregator.infer_aggregate(content_batches)
+
+    # if a path that exists is provided, append the summary with markdown formatting
+    if out_path:
+        out_path = Path(out_path)
+
+        try:
+            with open(out_path, "a", encoding="utf-8") as f:
+                f.write("\n\n## Aggregate Summary\n\n")
+                f.write(
+                    "- This is an instruction-based LLM aggregation of the previous 'summary batches'.\n"
+                )
+                f.write(f"- Aggregation model: {aggregator.model_name}\n\n")
+                f.write(f"{full_summary}\n\n")
+            logging.info(f"Updated {out_path} with aggregate summary")
+        except Exception as e:
+            logging.error(f"unable to update {out_path} with aggregate summary: {e}")
+
+    full_summary_html = f"""
+        <div style="
+            margin-bottom: 20px;
+            font-size: 18px;
+            line-height: 1.5em;
+            color: #333;
+        ">
+            <h2 style="font-size: 22px; color: #555;">Aggregate Summary:</h2>
+            <p style="white-space: pre-line;">{full_summary}</p>
+        </div>
+        """
+    return full_summary_html
+
+
+def predict(
+    input_text: str,
+    model_name: str,
+    token_batch_length: int = 1024,
+    empty_cache: bool = True,
+    **settings,
+) -> list:
+    """
+    predict - helper fn to support multiple models for summarization at once
+
+    :param str input_text: the input text to summarize
+    :param str model_name: model name to use
+    :param int token_batch_length: the length of the token batches to use
+    :param bool empty_cache: whether to empty the cache before loading a new= model
+    :return: list of dicts with keys "summary" and "score"
+    """
+    if torch.cuda.is_available() and empty_cache:
+        torch.cuda.empty_cache()
+
+    model, tokenizer = load_model_and_tokenizer(model_name)
+    summaries = summarize_via_tokenbatches(
+        input_text,
+        model,
+        tokenizer,
+        batch_length=token_batch_length,
+        **settings,
+    )
+
+    del model
+    del tokenizer
+    gc.collect()
+
+    return summaries
+
+
+def proc_submission(
+    input_text: str,
+    model_name: str,
+    num_beams: int,
+    token_batch_length: int,
+    length_penalty: float,
+    repetition_penalty: float,
+    no_repeat_ngram_size: int,
+    predrop_stopwords: bool,
+    max_input_length: int = 6144,
+):
+    """
+    proc_submission - a helper function for the gradio module to process submissions
+
+    Args:
+        input_text (str): the input text to summarize
+        model_name (str): the hf model tag of the model to use
+        num_beams (int): the number of beams to use
+        token_batch_length (int): the length of the token batches to use
+        length_penalty (float): the length penalty to use
+        repetition_penalty (float): the repetition penalty to use
+        no_repeat_ngram_size (int): the no repeat ngram size to use
+        predrop_stopwords (bool): whether to pre-drop stopwords before truncating/summarizing
+        max_input_length (int, optional): the maximum input length to use. Defaults to 6144.
+
+    Note:
+        the max_input_length is set to 6144 by default, but can be changed by setting the
+        environment variable APP_MAX_WORDS to a different value.
+
+    Returns:
+        tuple (4): a tuple containing the following:
+    """
+
+    remove_stagnant_files()  # clean up old files
+    settings = {
+        "length_penalty": float(length_penalty),
+        "repetition_penalty": float(repetition_penalty),
+        "no_repeat_ngram_size": int(no_repeat_ngram_size),
+        "encoder_no_repeat_ngram_size": 4,
+        "num_beams": int(num_beams),
+        "min_length": 4,
+        "max_length": int(token_batch_length // 4),
+        "early_stopping": True,
+        "do_sample": False,
+    }
+    max_input_length = int(os.environ.get("APP_MAX_WORDS", max_input_length))
+    logging.info(
+        f"max_input_length set to: {max_input_length}. pre-drop stopwords: {predrop_stopwords}"
+    )
+
+    st = time.perf_counter()
+    history = {}
+    cln_text = clean(input_text, lower=False)
+    parsed_cln_text = remove_stopwords(cln_text) if predrop_stopwords else cln_text
+    logging.info(
+        f"pre-truncation word count: {len(contraction_aware_tokenize(parsed_cln_text))}"
+    )
+    truncation_validated = truncate_word_count(
+        parsed_cln_text, max_words=max_input_length
+    )
+
+    if truncation_validated["was_truncated"]:
+        model_input_text = truncation_validated["processed_text"]
+        # create elaborate HTML warning
+        input_wc = len(contraction_aware_tokenize(parsed_cln_text))
+        msg = f"""
+        <div style="background-color: #FFA500; color: white; padding: 20px;">
+        <h3>Warning</h3>
+        <p>Input text was truncated to {max_input_length} words. That's about {100*max_input_length/input_wc:.2f}% of the original text.</p>
+        <p>Dropping stopwords is set to {predrop_stopwords}. If this is not what you intended, please validate the advanced settings.</p>
+        </div>
+        """
+        logging.warning(msg)
+        history["WARNING"] = msg
+    else:
+        model_input_text = truncation_validated["processed_text"]
+        msg = None
+
+    if len(input_text) < 50:
+        # this is essentially a different case from the above
+        msg = f"""
+        <div style="background-color: #880808; color: white; padding: 20px;">
+        <br>
+        <img src="https://i.imgflip.com/7kadd9.jpg" alt="no text">
+        <br>
+        <h3>Error</h3>
+        <p>Input text is too short to summarize. Detected {len(input_text)} characters.
+        Please load text by selecting an example from the dropdown menu or by pasting text into the text box.</p>
+        </div>
+        """
+        logging.warning(msg)
+        logging.warning("RETURNING EMPTY STRING")
+        history["WARNING"] = msg
+
+        return msg, "<strong>No summary generated.</strong>", "", []
+
+    _summaries = predict(
+        input_text=model_input_text,
+        model_name=model_name,
+        token_batch_length=token_batch_length,
+        **settings,
+    )
+    sum_text = [s["summary"][0].strip() + "\n" for s in _summaries]
+    sum_scores = [
+        f" - Batch Summary {i}: {round(s['summary_score'],4)}"
+        for i, s in enumerate(_summaries)
+    ]
+
+    full_summary = textlist2html(sum_text)
+    history["Summary Scores"] = "<br><br>"
+    scores_out = "\n".join(sum_scores)
+    rt = round((time.perf_counter() - st) / 60, 2)
+    logging.info(f"Runtime: {rt} minutes")
+    html = ""
+    html += f"<p>Runtime: {rt} minutes with model: {model_name}</p>"
+    if msg is not None:
+        html += msg
+
+    html += ""
+
+    settings["remove_stopwords"] = predrop_stopwords
+    settings["model_name"] = model_name
+    saved_file = saves_summary(summarize_output=_summaries, outpath=None, **settings)
+    return html, full_summary, scores_out, saved_file
+
+
+def load_single_example_text(
+    example_path: str or Path,
+    max_pages: int = 20,
+) -> str:
+    """
+    load_single_example_text - loads a single example text file
+
+    :param strorPath example_path: name of the example to load
+    :param int max_pages: the maximum number of pages to load from a PDF
+    :return str: the text of the example
+    """
+    global name_to_path, ocr_model
+    full_ex_path = name_to_path[example_path]
+    full_ex_path = Path(full_ex_path)
+    if full_ex_path.suffix in [".txt", ".md"]:
+        with open(full_ex_path, "r", encoding="utf-8", errors="ignore") as f:
+            raw_text = f.read()
+        text = clean(raw_text, lower=False)
+    elif full_ex_path.suffix == ".pdf":
+        logging.info(f"Loading PDF file {full_ex_path}")
+        max_pages = int(os.environ.get("APP_OCR_MAX_PAGES", max_pages))
+        logging.info(f"max_pages set to: {max_pages}")
+        conversion_stats = convert_PDF_to_Text(
+            full_ex_path,
+            ocr_model=ocr_model,
+            max_pages=max_pages,
+        )
+        text = conversion_stats["converted_text"]
+    else:
+        logging.error(f"Unknown file type {full_ex_path.suffix}")
+        text = "ERROR - check example path"
+
+    return text
+
+
+def load_uploaded_file(file_obj, max_pages: int = 20, lower: bool = False) -> str:
+    """
+    load_uploaded_file - loads a file uploaded by the user
+
+    :param file_obj (POTENTIALLY list): Gradio file object inside a list
+    :param int max_pages: the maximum number of pages to load from a PDF
+    :param bool lower: whether to lowercase the text
+    :return str: the text of the file
+    """
+    global ocr_model
+    logger = logging.getLogger(__name__)
+    # check if mysterious file object is a list
+    if isinstance(file_obj, list):
+        file_obj = file_obj[0]
+    file_path = Path(file_obj.name)
+    try:
+        logger.info(f"Loading file:\t{file_path}")
+        if file_path.suffix in [".txt", ".md"]:
+            with open(file_path, "r", encoding="utf-8", errors="ignore") as f:
+                raw_text = f.read()
+            text = clean(raw_text, lower=lower)
+        elif file_path.suffix == ".pdf":
+            logger.info(f"loading a PDF file: {file_path.name}")
+            max_pages = int(os.environ.get("APP_OCR_MAX_PAGES", max_pages))
+            logger.info(f"max_pages is: {max_pages}. Starting conversion...")
+            conversion_stats = convert_PDF_to_Text(
+                file_path,
+                ocr_model=ocr_model,
+                max_pages=max_pages,
+            )
+            text = conversion_stats["converted_text"]
+        else:
+            logger.error(f"Unknown file type:\t{file_path.suffix}")
+            text = "ERROR - check file - unknown file type. PDF, TXT, and MD are supported."
+
+        return text
+    except Exception as e:
+        logger.error(f"Trying to load file:\t{file_path},\nerror:\t{e}")
+        return f"Error: Could not read file {file_path.name}. Make sure it is a PDF, TXT, or MD file."
+
+
+def parse_args():
+    """arguments for the command line interface"""
+    parser = argparse.ArgumentParser(
+        description="Document Summarization with Long-Document Transformers - Demo",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+        epilog="Runs a local-only web UI to summarize documents. pass --share for a public link to share.",
+    )
+
+    parser.add_argument(
+        "--share",
+        dest="share",
+        action="store_true",
+        help="Create a public link to share",
+    )
+    parser.add_argument(
+        "-m",
+        "--model",
+        type=str,
+        default=None,
+        help=f"Add a custom model to the list of models: {pp.pformat(MODEL_OPTIONS, compact=True)}",
+    )
+    parser.add_argument(
+        "-nb",
+        "--add_beam_option",
+        type=int,
+        default=None,
+        help=f"Add a beam search option to the demo UI options, default: {pp.pformat(BEAM_OPTIONS, compact=True)}",
+    )
+    parser.add_argument(
+        "-batch",
+        "--token_batch_option",
+        type=int,
+        default=None,
+        help=f"Add a token batch size to the demo UI options, default: {pp.pformat(TOKEN_BATCH_OPTIONS, compact=True)}",
+    )
+    parser.add_argument(
+        "-max_agg",
+        "-2x",
+        "--aggregator_beam_boost",
+        dest="aggregator_beam_boost",
+        action="store_true",
+        help="Double the number of beams for the aggregator during beam search",
+    )
+    parser.add_argument(
+        "-level",
+        "--log_level",
+        type=str,
+        default="INFO",
+        choices=["DEBUG", "INFO", "WARNING", "ERROR"],
+        help="Set the logging level",
+    )
+
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    """main - the main function of the app"""
+    logger = logging.getLogger(__name__)
+    args = parse_args()
+    logger.setLevel(args.log_level)
+    logger.info(f"args: {pp.pformat(args.__dict__, compact=True)}")
+
+    # add any custom options
+    if args.model is not None:
+        logger.info(f"Adding model {args.model} to the list of models")
+        MODEL_OPTIONS.append(args.model)
+    if args.add_beam_option is not None:
+        logger.info(f"Adding beam search option {args.add_beam_option} to the list")
+        BEAM_OPTIONS.append(args.add_beam_option)
+    if args.token_batch_option is not None:
+        logger.info(f"Adding token batch option {args.token_batch_option} to the list")
+        TOKEN_BATCH_OPTIONS.append(args.token_batch_option)
+
+    if args.aggregator_beam_boost:
+        logger.info("Doubling aggregator num_beams")
+        _agg_cfg = aggregator.get_generation_config()
+        _agg_cfg["num_beams"] = _agg_cfg["num_beams"] * 2
+        aggregator.update_generation_config(**_agg_cfg)
+
+    logger.info("Loading OCR model")
+    with contextlib.redirect_stdout(None):
+        ocr_model = ocr_predictor(
+            "db_resnet50",
+            "crnn_mobilenet_v3_large",
+            pretrained=True,
+            assume_straight_pages=True,
+        )
+
+    # load the examples
+    name_to_path = load_example_filenames(_here / "examples")
+    logger.info(f"Loaded {len(name_to_path)} examples")
+
+    demo = gr.Blocks(title="Document Summarization")
+    _examples = list(name_to_path.keys())
+    logger.info("Starting app instance")
+    with demo:
+        gr.Markdown(
+            """# Document Summarization with Long-Document Transformers
+
+            An example use case for fine-tuned long document transformers. Model(s) are trained on [book summaries](https://hf.co/datasets/kmfoda/booksum). Architectures [in this demo](https://hf.co/spaces/pszemraj/document-summarization) are [LongT5-base](https://hf.co/pszemraj/long-t5-tglobal-base-16384-book-summary) and [Pegasus-X-Large](https://hf.co/pszemraj/pegasus-x-large-book-summary).
+
+            **Want more performance?** Run this demo from a free [Google Colab GPU](https://colab.research.google.com/gist/pszemraj/52f67cf7326e780155812a6a1f9bb724/document-summarization-on-gpu.ipynb)
+            """
+        )
+        with gr.Column():
+            gr.Markdown(
+                """## Load Inputs & Select Parameters
+
+                Enter/paste text below, or upload a file. Pick a model & adjust params (_optional_), and press **Summarize!**
+
+                See [the guide doc](https://gist.github.com/pszemraj/722a7ba443aa3a671b02d87038375519) for details.
+                """
+            )
+            with gr.Row():
+                with gr.Column(variant="compact"):
+                    model_name = gr.Dropdown(
+                        choices=MODEL_OPTIONS,
+                        value=MODEL_OPTIONS[0],
+                        label="Model Name",
+                    )
+                    num_beams = gr.Radio(
+                        choices=BEAM_OPTIONS,
+                        value=BEAM_OPTIONS[len(BEAM_OPTIONS) // 2],
+                        label="Beam Search: # of Beams",
+                    )
+                    load_examples_button = gr.Button(
+                        "Load Example in Dropdown",
+                    )
+                    load_file_button = gr.Button("Upload & Process File")
+                with gr.Column(variant="compact"):
+                    example_name = gr.Dropdown(
+                        _examples,
+                        label="Examples",
+                        value=random.choice(_examples),
+                    )
+                    uploaded_file = gr.File(
+                        label="File Upload",
+                        file_count="single",
+                        file_types=[".txt", ".md", ".pdf"],
+                        type="filepath",
+                    )
+            with gr.Row():
+                input_text = gr.Textbox(
+                    lines=4,
+                    max_lines=8,
+                    label="Text to Summarize",
+                    placeholder="Enter text to summarize, the text will be cleaned and truncated on Spaces. Narrative, academic (both papers and lecture transcription), and article text work well. May take a bit to generate depending on the input text :)",
+                )
+        with gr.Column():
+            gr.Markdown("## Generate Summary")
+            with gr.Row():
+                summarize_button = gr.Button(
+                    "Summarize!",
+                    variant="primary",
+                )
+                gr.Markdown(
+                    "_Summarization should take ~1-2 minutes for most settings, but may extend up to 5-10 minutes in some scenarios._"
+                )
+            output_text = gr.HTML("<em>Output will appear below:</em>")
+            with gr.Column():
+                gr.Markdown("### Results & Scores")
+                with gr.Row():
+                    with gr.Column(variant="compact"):
+                        gr.Markdown(
+                            "Download the summary as a text file, with parameters and scores."
+                        )
+                        text_file = gr.File(
+                            label="Download as Text File",
+                            file_count="single",
+                            type="filepath",
+                            interactive=False,
+                        )
+                    with gr.Column(variant="compact"):
+                        gr.Markdown(
+                            "Scores **roughly** represent the summary quality as a measure of the model's 'confidence'. less-negative numbers (closer to 0) are better."
+                        )
+                        summary_scores = gr.Textbox(
+                            label="Summary Scores",
+                            placeholder="Summary scores will appear here",
+                        )
+            with gr.Column(variant="panel"):
+                gr.Markdown("### **Summary Output**")
+                summary_text = gr.HTML(
+                    label="Summary",
+                    value="<i>Summary will appear here!</i>",
+                )
+            with gr.Column():
+                gr.Markdown("### **Aggregate Summary Batches**")
+                with gr.Row():
+                    aggregate_button = gr.Button(
+                        "Aggregate!",
+                        variant="primary",
+                    )
+                    gr.Markdown(
+                        f"""Aggregate the above batches into a cohesive summary.
+                    - A secondary instruct-tuned LM consolidates info
+                    - Current model: [{AGGREGATE_MODEL}](https://hf.co/{AGGREGATE_MODEL})
+                                """
+                    )
+                with gr.Column(variant="panel"):
+                    aggregated_summary = gr.HTML(
+                        label="Aggregate Summary",
+                        value="<i>Aggregate summary will appear here!</i>",
+                    )
+
+        with gr.Column():
+            gr.Markdown(
+                """### Advanced Settings
+
+            Refer to [the guide doc](https://gist.github.com/pszemraj/722a7ba443aa3a671b02d87038375519) for what these are, and how they impact _quality_ and _speed_.
+            """
+            )
+            with gr.Row():
+                length_penalty = gr.Slider(
+                    minimum=0.3,
+                    maximum=1.1,
+                    label="length penalty",
+                    value=0.7,
+                    step=0.05,
+                )
+                token_batch_length = gr.Radio(
+                    choices=TOKEN_BATCH_OPTIONS,
+                    label="token batch length",
+                    # select median option
+                    value=TOKEN_BATCH_OPTIONS[len(TOKEN_BATCH_OPTIONS) // 2],
+                )
+
+            with gr.Row():
+                repetition_penalty = gr.Slider(
+                    minimum=1.0,
+                    maximum=5.0,
+                    label="repetition penalty",
+                    value=1.5,
+                    step=0.1,
+                )
+                no_repeat_ngram_size = gr.Radio(
+                    choices=[2, 3, 4, 5],
+                    label="no repeat ngram size",
+                    value=3,
+                )
+                predrop_stopwords = gr.Checkbox(
+                    label="Drop Stopwords (Pre-Truncation)",
+                    value=False,
+                )
+
+        load_examples_button.click(
+            fn=load_single_example_text, inputs=[example_name], outputs=[input_text]
+        )
+
+        load_file_button.click(
+            fn=load_uploaded_file, inputs=uploaded_file, outputs=[input_text]
+        )
+
+        summarize_button.click(
+            fn=proc_submission,
+            inputs=[
+                input_text,
+                model_name,
+                num_beams,
+                token_batch_length,
+                length_penalty,
+                repetition_penalty,
+                no_repeat_ngram_size,
+                predrop_stopwords,
+            ],
+            outputs=[output_text, summary_text, summary_scores, text_file],
+        )
+        aggregate_button.click(
+            fn=aggregate_text,
+            inputs=[summary_text, text_file],
+            outputs=[aggregated_summary],
+        )
+    demo.launch(share=args.share, debug=True)

gitattributes

@@ -0,0 +1,31 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

gitignore

@@ -0,0 +1,29 @@
+# logs
+*.log
+*LOGFILE*
+
+# output files need to be force-added
+*.csv
+*.png
+*.jpg
+*.jpeg
+*.pkl
+*.xlsx
+*.txt
+
+# cache
+*__pycache__/
+*.pyc
+
+# reports folder - need to be force-added
+*reports/
+
+# scratch files and folders
+
+*scratch*
+*scratch/
+
+# notebooks
+
+*notebooks/
+*.ipynb

pdf2text.py

@@ -0,0 +1,346 @@
+# -*- coding: utf-8 -*-
+"""
+pdf2text.py - convert pdf files to text files using OCR
+"""
+import logging
+import os
+import re
+import shutil
+import time
+from datetime import date
+from os.path import join
+from pathlib import Path
+
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s %(levelname)s %(message)s",
+    datefmt="%m/%d/%Y %I:%M:%S",
+)
+
+
+os.environ["USE_TORCH"] = "1"
+
+from cleantext import clean
+from doctr.io import DocumentFile
+from doctr.models import ocr_predictor
+from spellchecker import SpellChecker
+
+
+def simple_rename(filepath, target_ext=".txt"):
+    """simple_rename - get a new str to rename a file"""
+    _fp = Path(filepath)
+    basename = _fp.stem
+    return f"OCR_{basename}_{target_ext}"
+
+
+def rm_local_text_files(name_contains="RESULT_"):
+    """
+    rm_local_text_files - remove local text files
+    """
+    files = [
+        f
+        for f in Path.cwd().iterdir()
+        if f.is_file() and f.suffix == ".txt" and name_contains in f.name
+    ]
+    logging.info(f"removing {len(files)} text files")
+    for f in files:
+        os.remove(f)
+    logging.info("done")
+
+
+def corr(
+    s: str,
+    add_space_when_numerics=False,
+    exceptions=["e.g.", "i.e.", "etc.", "cf.", "vs.", "p."],
+) -> str:
+    """corrects spacing in a string
+
+    Args:
+        s (str): the string to correct
+        add_space_when_numerics (bool, optional): [add a space when a period is between two numbers, example 5.73]. Defaults to False.
+        exceptions (list, optional): [do not change these substrings]. Defaults to ['e.g.', 'i.e.', 'etc.', 'cf.', 'vs.', 'p.'].
+
+    Returns:
+        str: the corrected string
+    """
+    if add_space_when_numerics:
+        s = re.sub(r"(\d)\.(\d)", r"\1. \2", s)
+
+    s = re.sub(r"\s+", " ", s)
+    s = re.sub(r'\s([?.!"](?:\s|$))', r"\1", s)
+
+    # fix space before apostrophe
+    s = re.sub(r"\s\'", r"'", s)
+    # fix space after apostrophe
+    s = re.sub(r"'\s", r"'", s)
+    # fix space before comma
+    s = re.sub(r"\s,", r",", s)
+
+    for e in exceptions:
+        expected_sub = re.sub(r"\s", "", e)
+        s = s.replace(expected_sub, e)
+
+    return s
+
+
+def fix_punct_spaces(string: str) -> str:
+    """
+    fix_punct_spaces - fix spaces around punctuation
+
+    :param str string: input string
+    :return str: string with spaces fixed
+    """
+
+    fix_spaces = re.compile(r"\s*([?!.,]+(?:\s+[?!.,]+)*)\s*")
+    string = fix_spaces.sub(lambda x: "{} ".format(x.group(1).replace(" ", "")), string)
+    string = string.replace(" ' ", "'")
+    string = string.replace(' " ', '"')
+    return string.strip()
+
+
+def clean_OCR(ugly_text: str) -> str:
+    """
+    clean_OCR - clean up the OCR text
+
+    :param str ugly_text: input text to be cleaned
+    :return str: cleaned text
+    """
+    # Remove all the newlines.
+    cleaned_text = ugly_text.replace("\n", " ")
+    # Remove all the tabs.
+    cleaned_text = cleaned_text.replace("\t", " ")
+    # Remove all the double spaces.
+    cleaned_text = cleaned_text.replace("  ", " ")
+    # Remove all the spaces at the beginning of the text.
+    cleaned_text = cleaned_text.lstrip()
+    # remove all instances of "- " and " - "
+    cleaned_text = cleaned_text.replace("- ", "")
+    cleaned_text = cleaned_text.replace(" -", "")
+    return fix_punct_spaces(cleaned_text)
+
+
+def move2completed(
+    from_dir, filename, new_folder: str = "completed", verbose: bool = False
+):
+    """
+    move2completed - move a file to a new folder
+    """
+    old_filepath = join(from_dir, filename)
+
+    new_filedirectory = join(from_dir, new_folder)
+
+    if not os.path.isdir(new_filedirectory):
+        os.mkdir(new_filedirectory)
+        if verbose:
+            print("created new directory for files at: \n", new_filedirectory)
+    new_filepath = join(new_filedirectory, filename)
+
+    try:
+        shutil.move(old_filepath, new_filepath)
+        logging.info("successfully moved the file {} to */completed.".format(filename))
+    except:
+        logging.info(
+            "ERROR! unable to move file to \n{}. Please investigate".format(
+                new_filepath
+            )
+        )
+
+
+custom_replace_list = {
+    "t0": "to",
+    "'$": "'s",
+    ",,": ", ",
+    "_ ": " ",
+    " '": "'",
+}
+
+replace_corr_exceptions = {
+    "i. e.": "i.e.",
+    "e. g.": "e.g.",
+    "e. g": "e.g.",
+    " ,": ",",
+}
+
+
+spell = SpellChecker()
+
+
+def check_word_spelling(word: str) -> bool:
+    """
+    check_word_spelling - check the spelling of a word
+
+    Args:
+        word (str): word to check
+
+    Returns:
+        bool: True if word is spelled correctly, False if not
+    """
+
+    misspelled = spell.unknown([word])
+
+    return len(misspelled) == 0
+
+
+def eval_and_replace(text: str, match_token: str = "- ") -> str:
+    """
+    eval_and_replace  - conditionally replace all instances of a substring in a string based on whether the eliminated substring results in a valid word
+
+    Args:
+        text (str): text to evaluate
+        match_token (str, optional): token to replace. Defaults to "- ".
+
+    Returns:
+        str:  text with replaced tokens
+    """
+
+    if match_token not in text:
+        return text
+    else:
+        while True:
+            full_before_text = text.split(match_token, maxsplit=1)[0]
+            before_text = [
+                char for char in full_before_text.split()[-1] if char.isalpha()
+            ]
+            before_text = "".join(before_text)
+            full_after_text = text.split(match_token, maxsplit=1)[-1]
+            after_text = [char for char in full_after_text.split()[0] if char.isalpha()]
+            after_text = "".join(after_text)
+            full_text = before_text + after_text
+            if check_word_spelling(full_text):
+                text = full_before_text + full_after_text
+            else:
+                text = full_before_text + " " + full_after_text
+            if match_token not in text:
+                break
+        return text
+
+
+def cleantxt_ocr(ugly_text, lower=False, lang: str = "en") -> str:
+    """
+    cleantxt_ocr - clean text from OCR
+
+        https://pypi.org/project/clean-text/
+    Args:
+        ugly_text (str): text to clean
+        lower (bool, optional): lowercase text. Defaults to False.
+        lang (str, optional): language of text. Defaults to "en".
+
+    Returns:
+        str: cleaned text
+    """
+
+    cleaned_text = clean(
+        ugly_text,
+        fix_unicode=True,  # fix various unicode errors
+        to_ascii=True,  # transliterate to closest ASCII representation
+        lower=lower,  # lowercase text
+        no_line_breaks=True,  # fully strip line breaks as opposed to only normalizing them
+        no_urls=True,  # replace all URLs with a special token
+        no_emails=True,  # replace all email addresses with a special token
+        no_phone_numbers=True,  # replace all phone numbers with a special token
+        no_numbers=False,  # replace all numbers with a special token
+        no_digits=False,  # replace all digits with a special token
+        no_currency_symbols=False,  # replace all currency symbols with a special token
+        no_punct=False,  # remove punctuations
+        replace_with_punct="",  # instead of removing punctuations you may replace them
+        replace_with_url="this url",
+        replace_with_email="this email",
+        replace_with_phone_number="this phone number",
+        lang=lang,  # set to 'de' for German special handling
+    )
+
+    return cleaned_text
+
+
+def format_ocr_out(OCR_data):
+    """format OCR output to text"""
+    if isinstance(OCR_data, list):
+        text = " ".join(OCR_data)
+    else:
+        text = str(OCR_data)
+    _clean = cleantxt_ocr(text)
+    return corr(_clean)
+
+
+def postprocess(text: str) -> str:
+    """to be used after recombining the lines"""
+
+    proc = corr(cleantxt_ocr(text))
+
+    for k, v in custom_replace_list.items():
+        proc = proc.replace(str(k), str(v))
+
+    proc = corr(proc)
+
+    for k, v in replace_corr_exceptions.items():
+        proc = proc.replace(str(k), str(v))
+
+    return eval_and_replace(proc)
+
+
+def result2text(result, as_text=False) -> str or list:
+    """Convert OCR result to text"""
+
+    full_doc = []
+    for i, page in enumerate(result.pages, start=1):
+        text = ""
+        for block in page.blocks:
+            text += "\n\t"
+            for line in block.lines:
+                for word in line.words:
+                    # print(dir(word))
+                    text += word.value + " "
+        full_doc.append(text)
+
+    return "\n".join(full_doc) if as_text else full_doc
+
+
+def convert_PDF_to_Text(
+    PDF_file,
+    ocr_model=None,
+    max_pages: int = 20,
+) -> str:
+    """
+    convert_PDF_to_Text - convert a PDF file to text
+
+    :param str PDF_file: path to PDF file
+    :param ocr_model: model to use for OCR, defaults to None (uses the default model)
+    :param int max_pages: maximum number of pages to process, defaults to 20
+    :return str: text from PDF
+    """
+    st = time.perf_counter()
+    PDF_file = Path(PDF_file)
+    ocr_model = ocr_predictor(pretrained=True) if ocr_model is None else ocr_model
+    logging.info(f"starting OCR on {PDF_file.name}")
+    doc = DocumentFile.from_pdf(PDF_file)
+    truncated = False
+    if len(doc) > max_pages:
+        logging.warning(
+            f"PDF has {len(doc)} pages, which is more than {max_pages}.. truncating"
+        )
+        doc = doc[:max_pages]
+        truncated = True
+
+    # Analyze
+    logging.info(f"running OCR on {len(doc)} pages")
+    result = ocr_model(doc)
+    raw_text = result2text(result)
+    proc_text = [format_ocr_out(r) for r in raw_text]
+    fin_text = [postprocess(t) for t in proc_text]
+
+    ocr_results = "\n\n".join(fin_text)
+
+    fn_rt = time.perf_counter() - st
+
+    logging.info("OCR complete")
+
+    results_dict = {
+        "num_pages": len(doc),
+        "runtime": round(fn_rt, 2),
+        "date": str(date.today()),
+        "converted_text": ocr_results,
+        "truncated": truncated,
+        "length": len(ocr_results),
+    }
+
+    return results_dict

requirements.txt

@@ -0,0 +1,12 @@
+accelerate
+clean-text[gpl]
+gradio==5.5.0
+natsort
+nltk
+pyspellchecker
+python-doctr[torch]
+rapidfuzz==2.13.7
+sentencepiece
+torch
+tqdm
+transformers==4.46.2

summarize.py

@@ -0,0 +1,177 @@
+"""
+summarize - a module for summarizing text using a model from the Hugging Face model hub
+"""
+import logging
+import os
+import pprint as pp
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(message)s")
+
+import torch
+from tqdm.auto import tqdm
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
+
+from utils import validate_pytorch2
+
+
+def load_model_and_tokenizer(model_name: str) -> tuple:
+    """
+    load_model_and_tokenizer - load a model and tokenizer from a model name/ID on the hub
+
+    :param str model_name: the model name/ID on the hub
+    :return tuple: a tuple containing the model and tokenizer
+    """
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    model = AutoModelForSeq2SeqLM.from_pretrained(
+        model_name,
+        use_auth_token=os.environ.get("HF_TOKEN", None),
+    ).to(device)
+    model = model.eval()
+
+    tokenizer = AutoTokenizer.from_pretrained(
+        model_name,
+        use_auth_token=os.environ.get("HF_TOKEN", None),
+    )
+
+    logging.info(f"Loaded model {model_name} to {device}")
+
+    if validate_pytorch2():
+        try:
+            logging.info("Compiling model with Torch 2.0")
+            model = torch.compile(model)
+        except Exception as e:
+            logging.warning(f"Could not compile model with Torch 2.0: {e}")
+    else:
+        logging.info("Torch 2.0 not detected, skipping compilation")
+
+    return model, tokenizer
+
+
+def summarize_and_score(
+    ids, mask, model, tokenizer, is_general_attention_model=True, **kwargs
+) -> tuple:
+    """
+    summarize_and_score - given a batch of ids and a mask, return a summary and a score for the summary
+
+    Args:
+        ids (): the batch of ids
+        mask (): the attention mask for the batch
+        model   (): the model to use for summarization
+        tokenizer (): the tokenizer to use for summarization
+        is_general_attention_model (bool, optional): whether the model is a general attention model. Defaults to True.
+        **kwargs: any additional arguments to pass to the model
+    Returns:
+        tuple (str, float): the summary,  the score for the summary
+    """
+
+    ids = ids[None, :]
+    mask = mask[None, :]
+
+    input_ids = ids.to("cuda") if torch.cuda.is_available() else ids
+    attention_mask = mask.to("cuda") if torch.cuda.is_available() else mask
+
+    global_attention_mask = torch.zeros_like(attention_mask)
+    # put global attention on <s> token
+    global_attention_mask[:, 0] = 1
+
+    if is_general_attention_model:
+        summary_pred_ids = model.generate(
+            input_ids,
+            attention_mask=attention_mask,
+            output_scores=True,
+            return_dict_in_generate=True,
+            **kwargs,
+        )
+    else:
+        summary_pred_ids = model.generate(
+            input_ids,
+            attention_mask=attention_mask,
+            global_attention_mask=global_attention_mask,
+            output_scores=True,
+            return_dict_in_generate=True,
+            **kwargs,
+        )
+    summary = tokenizer.batch_decode(
+        summary_pred_ids.sequences,
+        skip_special_tokens=True,
+        remove_invalid_values=True,
+    )
+    score = round(summary_pred_ids.sequences_scores.cpu().numpy()[0], 4)
+
+    return summary, score
+
+
+def summarize_via_tokenbatches(
+    input_text: str,
+    model,
+    tokenizer,
+    batch_length=2048,
+    batch_stride=16,
+    min_batch_length=512,
+    **kwargs,
+) -> list:
+    """
+    summarize_via_tokenbatches - summarize a long string via batches of tokens
+
+    Args:
+        input_text (str): the text to summarize
+        model (): the model to use for summarization
+        tokenizer (): the tokenizer to use for summarization
+        batch_length (int, optional): the length of each batch. Defaults to 2048.
+        batch_stride (int, optional): the stride of each batch. Defaults to 16. The stride is the number of tokens that overlap between batches.
+        min_batch_length (int, optional): the minimum length of each batch. Defaults to 512.
+
+        **kwargs: any additional arguments to pass to the model for inference
+    Returns:
+        list: a list of dictionaries containing the input tokens, the summary, and the summary score
+    """
+
+    logger = logging.getLogger(__name__)
+    # log all input parameters
+    if batch_length < min_batch_length:
+        logger.warning(
+            f"batch_length must be at least {min_batch_length}. Setting batch_length to {min_batch_length}"
+        )
+        batch_length = min_batch_length
+
+    logger.info(f"input parameters:\n{pp.pformat(kwargs)}")
+    logger.info(f"batch_length: {batch_length}, batch_stride: {batch_stride}")
+
+    encoded_input = tokenizer(
+        input_text,
+        padding="max_length",
+        truncation=True,
+        max_length=batch_length,
+        stride=batch_stride,
+        return_overflowing_tokens=True,
+        add_special_tokens=False,
+        return_tensors="pt",
+    )
+
+    in_id_arr, att_arr = encoded_input.input_ids, encoded_input.attention_mask
+    gen_summaries = []
+
+    pbar = tqdm(total=len(in_id_arr))
+
+    for _id, _mask in zip(in_id_arr, att_arr):
+        result, score = summarize_and_score(
+            ids=_id,
+            mask=_mask,
+            model=model,
+            tokenizer=tokenizer,
+            **kwargs,
+        )
+        score = round(float(score), 4)
+        _sum = {
+            "input_tokens": _id,
+            "summary": result,
+            "summary_score": score,
+        }
+        gen_summaries.append(_sum)
+        logger.debug(f"Score for batch: {score}. num chars: {len(repr(result))}")
+        logger.debug(f"Summary:\n\t{result}")
+        pbar.update()
+
+    pbar.close()
+
+    return gen_summaries

utils.py

@@ -0,0 +1,450 @@
+"""
+    utils.py - Utility functions for the project.
+"""
+import logging
+import re
+import subprocess
+from collections import defaultdict, deque
+from datetime import datetime, timedelta
+from itertools import combinations
+from pathlib import Path
+from typing import List
+
+logging.basicConfig(
+    format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
+    level=logging.INFO,
+)
+
+import torch
+from natsort import natsorted
+from nltk.tokenize import WhitespaceTokenizer, sent_tokenize, word_tokenize
+from rapidfuzz import fuzz
+
+STOPWORDS = set(
+    "a about above after again all also am an and any are aren't as at back be because been before being below between both but by can't cannot could couldn't did didn't do does doesn't doing don't down during each few for from further had hadn't has hasn't have haven't having he'd he'll he's hence her here here's hers herself him himself his how how's however i'd i'll i'm i've if in into is isn't it's its itself just let's me more moreover most mustn't my myself new nor now of off on once only or other ought our ours ourselves out over own really same shan't she'd she'll she's should shouldn't so some such than that that's the their theirs them themselves then there there's therefore these they they'd they'll they're they've this those through thus to too under until up use used using very was wasn't we we'd we'll we're we've were weren't what what's when when's where where's which while who who's whom why why's with won't would wouldn't you'd you'll you're you've your yours yourself yourselves".split()
+)
+
+
+def contraction_aware_tokenize(text: str) -> List[str]:
+    """contraction_aware_tokenize - merges words containing apostrophes as one token."""
+
+    # Tokenize the text using the WhitespaceTokenizer
+    tokenizer = WhitespaceTokenizer()
+    tokens = tokenizer.tokenize(text)
+
+    merged_tokens = []
+    merged_token = ""
+
+    for token in tokens:
+        if re.search(r"\w+'\w+", token):
+            # Token contains an apostrophe, merge with previous token
+            merged_token += token
+        else:
+            # no apostrophe, add previous merged token (if any) and current
+            if merged_token:
+                merged_tokens.append(merged_token)
+                merged_token = ""
+            merged_tokens.append(token)
+
+    # Add the last merged token (if any)
+    if merged_token:
+        merged_tokens.append(merged_token)
+
+    return merged_tokens
+
+
+def remove_stopwords(
+    text: str, stopwords: List[str] = STOPWORDS, contraction_tokenize: bool = True
+) -> str:
+    """
+    remove_stopwords - Remove stopwords from text.
+
+    :param str text: input text
+    :param List[str] stopwords: list of stopwords, defaults to STOPWORDS
+    :param bool contraction_tokenize: use custom apostrophe tokenizer, defaults to True
+    :return str: text with stopwords removed
+    """
+    lines = text.split("\n")
+    filtered_lines = []
+
+    def fix_commas(text: str) -> str:
+        """fixes commas in text to have a space after them"""
+        spaced_text = text.replace(",", ", ")
+        return spaced_text.replace("  ", " ").strip()
+
+    for line in lines:
+        sentences = sent_tokenize(line)
+        filtered_sentences = []
+
+        for sentence in sentences:
+            # Add space around punctuations for the regex to work correctly, only if they are followed by a letter
+            sentence_with_spaces = re.sub(r"([.,!?])(\w)", r"\1 \2", sentence[:-1])
+
+            words = (
+                contraction_aware_tokenize(sentence_with_spaces)
+                if contraction_tokenize
+                else word_tokenize(sentence_with_spaces)
+            )
+
+            filtered_words = []
+            for word in words:
+                if word.lower() not in stopwords:
+                    filtered_words.append(word)
+
+            filtered_sentence = " ".join(filtered_words)
+            # Restore original spaces around punctuation marks
+            filtered_sentence = re.sub(r"([.,!?])\s*", r"\1", filtered_sentence)
+
+            filtered_sentences.append(filtered_sentence + sentence[-1])
+
+        filtered_line = " ".join(filtered_sentences)
+
+        # Replace multiple consecutive whitespaces with a single space
+        filtered_line = re.sub(r"\s+", " ", filtered_line)
+        filtered_line = fix_commas(filtered_line.strip())
+
+        filtered_lines.append(filtered_line)
+
+    filtered_text = "\n".join(filtered_lines)
+
+    return filtered_text
+
+
+def remove_stagnant_files(
+    freq: str = "hourly",
+    search_path: str = ".",
+    substring="DocSumm",
+    remove_suffix=".txt",
+):
+    """
+    remove_stagnant_files - Remove files that have not been modified in a certain amount of time.
+
+    :param str freq: frequency of file removal, defaults to "hourly"
+    :param str search_path: location to search for files, defaults to "."
+    :param str substring: substring to search for in file names, defaults to "DocSumm"
+    :param str remove_suffix: suffix of files to remove, defaults to ".txt"
+    :raises ValueError: if freq is not one of "hourly", "daily", or "weekly"
+    """
+    current_time = datetime.now()
+    search_path = Path(search_path)
+
+    if freq == "hourly":
+        time_threshold = current_time - timedelta(hours=1)
+    elif freq == "daily":
+        time_threshold = current_time - timedelta(days=1)
+    elif freq == "weekly":
+        time_threshold = current_time - timedelta(weeks=1)
+    else:
+        raise ValueError(
+            "Invalid frequency. Supported values are 'hourly', 'daily', and 'weekly'."
+        )
+
+    files_to_remove = []
+    potential_files = [
+        f for f in search_path.iterdir() if f.is_file() and f.suffix == remove_suffix
+    ]
+    logging.info(f"Found {len(potential_files)} files.")
+    for candidate in potential_files:
+        if (
+            candidate.is_file()
+            and substring in candidate.name
+            and candidate.stat().st_mtime < time_threshold.timestamp()
+        ):
+            files_to_remove.append(candidate)
+        logging.debug(f"File {candidate} last modified at {candidate.stat().st_mtime}")
+    logging.info(f"Removing {len(files_to_remove)} files.")
+    for file_path in files_to_remove:
+        file_path.unlink()
+    logging.debug(f"Removed files: {files_to_remove}")
+
+
+def compare_model_size(model_name: str, threshold: int = 500) -> bool:
+    """
+    compare_model_size - compare string representations of model size to a threshold
+
+    :param str model_name: the model name to compare
+    :param int threshold: the threshold to compare against in millions, defaults to 500
+    :return: True if the model size is greater than the threshold, False or None otherwise
+    """
+    pattern = r"(\d+)(M|G|k|b)?"  # param regex
+
+    matches = re.findall(pattern, model_name)
+    if not matches:
+        return None
+
+    # Extract the parameter count and unit
+    parameter_count, unit = matches[-1]
+    parameter_count = int(parameter_count)
+
+    # Convert to the standard form (M for million, G for billion, k for thousand)
+    if unit == "G" or unit == "b":
+        parameter_count *= 1000
+    elif unit == "M":
+        pass
+    elif unit == "k":
+        parameter_count /= 1000
+    else:
+        return None  # Unknown
+
+    return parameter_count > threshold
+
+
+def validate_pytorch2(torch_version: str = None) -> bool:
+    """
+    validate_pytorch2 - validate that the PyTorch version is 2.0 or greater
+
+    :param str torch_version: the PyTorch version to validate, defaults to None
+    :return: True if the PyTorch version is 2.0 or greater, False otherwise
+    """
+
+    torch_version = torch.__version__ if torch_version is None else torch_version
+
+    pattern = r"^2\.\d+(\.\d+)*"
+
+    return True if re.match(pattern, torch_version) else False
+
+
+def get_timestamp(detailed=False) -> str:
+    """
+    get_timestamp - get a timestamp for the current time
+    :param bool detailed: whether to include seconds and microseconds, defaults to False
+    :return: str, the timestamp
+    """
+    return (
+        datetime.now().strftime("%b%d%Y_%H%M%S%f")
+        if detailed
+        else datetime.now().strftime("%b%d%Y_%H")
+    )
+
+
+def truncate_word_count(text: str, max_words=1024) -> dict:
+    """
+    truncate_word_count - truncate a text to a maximum number of words
+    :param str text: the text to truncate
+    :param int max_words: the maximum number of words to keep, defaults to 1024
+    :return: dict, the processed text
+    """
+    words = contraction_aware_tokenize(str(text))
+    processed = {}
+    if len(words) > max_words:
+        processed["was_truncated"] = True
+        processed["processed_text"] = " ".join(words[:max_words])
+    else:
+        processed["was_truncated"] = False
+        processed["processed_text"] = text
+    return processed
+
+
+def load_examples(src, filetypes=[".txt", ".pdf"]):
+    """
+    load_examples - a helper function for the gradio module to load examples
+    :param str src: the path to the examples
+    """
+    src = Path(src)
+    src.mkdir(exist_ok=True)
+
+    pdf_url = (
+        "https://www.dropbox.com/s/y92xy7o5qb88yij/all_you_need_is_attention.pdf?dl=1"
+    )
+    subprocess.run(["wget", pdf_url, "-O", src / "all_you_need_is_attention.pdf"])
+    examples = [f for f in src.iterdir() if f.suffix in filetypes]
+    examples = natsorted(examples)
+    # load the examples into a list
+    text_examples = []
+    for example in examples:
+        with open(example, "r") as f:
+            text = f.read()
+            text_examples.append([text, "base", 2, 1024, 0.7, 3.5, 3])
+
+    return text_examples
+
+
+def load_example_filenames(example_path: str or Path):
+    """
+    load_example_filenames - a helper function for the gradio module to load examples
+    Returns:
+        dict, the examples (filename:full path)
+    """
+    example_path = Path(example_path)
+    # load the examples into a list
+    examples = {f.name: f for f in example_path.glob("*.txt")}
+    return examples
+
+
+def textlist2html(text_batches: List[str]) -> str:
+    """textlist2html - convert a list of text summaries into a single HTML string"""
+    # Step 1: Generate each summary batch as a string of HTML
+    formatted_batches = [
+        f"""
+        <div style="
+            margin-bottom: 20px;
+            font-size: 18px;
+            line-height: 1.5em;
+            color: #333;
+        ">
+            <h2 style="font-size: 22px; color: #555;">Batch {i}:</h2>
+            <p style="white-space: pre-line;">{s}</p>
+        </div>
+        """
+        for i, s in enumerate(text_batches, start=1)
+    ]
+
+    # Step 2: Join all the summary batches together into one string
+    joined_batches = "".join(formatted_batches)
+
+    # Step 3: Wrap the summary string in a larger div with background color, border, and padding
+    text_html_block = f"""
+    <div style="
+        border: 1px solid #ddd;
+        border-radius: 5px;
+        padding: 20px;
+    ">
+    {joined_batches}
+    </div>
+    """
+
+    return text_html_block
+
+
+def extract_batches(html_string: str, pattern=None, flags=None) -> list:
+    """
+    Extract batches of text from an HTML string.
+
+    Args:
+        html_string (str): The HTML string to extract batches from.
+        pattern (str, optional): The regular expression pattern to use. Defaults to a pattern that matches batches in the format provided.
+        flags (int, optional): The flags to use with the regular expression. Defaults to re.DOTALL.
+
+    Returns:
+        list: A list of dictionaries where each dictionary represents a batch and has 'title' and 'content' keys.
+    """
+    # Set default pattern if none provided
+    if pattern is None:
+        pattern = r'<h2 style="font-size: 22px; color: #555;">(.*?)</h2>\s*<p style="white-space: pre-line;">(.*?)</p>'
+
+    # Set default flags if none provided
+    if flags is None:
+        flags = re.DOTALL
+
+    try:
+        # Find all matches in the string
+        matches = re.findall(pattern, html_string, flags)
+
+        # Convert matches to a list of dictionaries
+        batches = [
+            {"title": title.strip(), "content": content.strip()}
+            for title, content in matches
+        ]
+
+        return batches
+    except re.error as e:
+        logging.error(f"An error occurred while trying to extract batches: {e}")
+        return []
+
+
+def extract_keywords(
+    text: str, num_keywords: int = 3, window_size: int = 5, kw_max_len: int = 20
+) -> List[str]:
+    """
+    Extracts keywords from a text using a simplified TextRank algorithm.
+
+    Args:
+        text: The text to extract keywords from.
+        num_keywords: The number of keywords to extract. Default: 3
+        window_size: The number of words considered for co-occurrence. Default: 5
+        kw_max_len: The maximum length of a keyword (truncate longer keywords to max). Default: 20
+    Returns:
+        A list of strings, where each string is a keyword extracted from the input text.
+    """
+    logger = logging.getLogger(__name__)
+    # Remove stopwords and tokenize the text into words
+    words = [
+        word
+        for word in re.findall(r"\b\w{3,}\b", text.lower())
+        if word not in STOPWORDS
+    ]
+
+    # Create a graph of word co-occurrences within a moving window of words
+    cooccur = defaultdict(lambda: defaultdict(int))
+    deque_words = deque(maxlen=window_size)
+    for word in words:
+        for w1, w2 in combinations(deque_words, 2):
+            cooccur[w1][w2] += 1
+            cooccur[w2][w1] += 1
+        deque_words.append(word)
+
+    # Assign scores to words using a simplified TextRank algorithm
+    scores = defaultdict(float)
+    for _ in range(10):
+        new_scores = defaultdict(float)
+        for word, co_words in cooccur.items():
+            new_scores[word] = 0.15 + 0.85 * sum(
+                cooccur[word][other] / sum(cooccur[other].values()) * scores[other]
+                for other in co_words
+            )
+        scores = new_scores
+
+    # Sort the words by score and return the top num_keywords keywords
+    keywords = sorted(scores, key=scores.get, reverse=True)[:num_keywords]
+    logger.debug(f"All keywords: {keywords}")
+    # Use fuzzy matching to remove similar keywords
+    final_keywords = []
+    for keyword in keywords:
+        if not any(fuzz.ratio(keyword, other) > 70 for other in final_keywords):
+            final_keywords.append(keyword[:kw_max_len])
+    logger.debug(f"Keywords (max len. {kw_max_len}):\t{final_keywords}")
+    return final_keywords
+
+
+def saves_summary(
+    summarize_output, outpath: str or Path = None, add_signature=True, **kwargs
+) -> Path:
+    """
+    saves_summary - save the summary generated from summarize_via_tokenbatches() to a text file
+
+    summarize_output: output from summarize_via_tokenbatches()
+    outpath: path to the output file
+    add_signature: whether to add a signature to the output file
+    kwargs: additional keyword arguments to include in the output file
+    """
+    logger = logging.getLogger(__name__)
+    sum_text = [f"{s['summary'][0]}\n" for s in summarize_output]
+    sum_scores = [f"\n - {round(s['summary_score'],4)}" for s in summarize_output]
+    scores_text = "\n".join(sum_scores)
+    full_summary = "\n".join(sum_text)
+
+    keywords = "_".join(extract_keywords(full_summary, kw_max_len=4))
+    logger.debug(f"kw:\t{keywords}")
+    outpath = (
+        Path.cwd() / f"DocSumm_{keywords}_{get_timestamp()}.txt"
+        if outpath is None
+        else Path(outpath)
+    )
+    logger.info(f"Saving summary to:\t{outpath.name}")
+    with open(
+        outpath,
+        "w",
+        encoding="utf-8",
+    ) as fo:
+        fo.writelines(full_summary)
+        fo.write("\n\n")
+        if add_signature:
+            fo.write("\n\n---\n\n")
+            fo.write("Generated with the Document Summarization space :)\n\n")
+            fo.write("https://hf.co/spaces/pszemraj/document-summarization\n\n")
+    with open(
+        outpath,
+        "a",
+        encoding="utf-8",
+    ) as fo:
+        fo.write("\n")
+        fo.write("## Section Scores:\n\n")
+        fo.writelines(scores_text)
+        fo.write("\n\n")
+        fo.write(f"Date: {get_timestamp()}\n\n")
+        if kwargs:
+            fo.write("---\n\n")
+            fo.write("## Parameters:\n\n")
+            for key, value in kwargs.items():
+                fo.write(f"{key}: {value}\n")
+    return str(outpath.resolve())