✨ add ability to add custom options in CLI

Signed-off-by: peter szemraj <peterszemraj@gmail.com>

app.py

@@ -2,6 +2,9 @@
 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:
@@ -18,6 +21,8 @@ import logging
 import os
 import random
 import re
+import pprint as pp
+import sys
 import time
 from pathlib import Path
 
@@ -52,7 +57,7 @@ _here = Path(__file__).parent
 nltk.download("punkt", force=True, quiet=True)
 nltk.download("popular", force=True, quiet=True)
 
-
+# Constants & Globals
 MODEL_OPTIONS = [
     "pszemraj/long-t5-tglobal-base-16384-book-summary",
     "pszemraj/long-t5-tglobal-base-sci-simplify",
@@ -60,6 +65,14 @@ MODEL_OPTIONS = [
     "pszemraj/long-t5-tglobal-base-16384-booksci-summary-v1",
     "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 = "MBZUAI/LaMini-Flan-T5-783M"  # model to use for aggregation
@@ -67,8 +80,11 @@ AGGREGATE_MODEL = "MBZUAI/LaMini-Flan-T5-783M"  # model to use for aggregation
 # 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)
+aggregator = BatchAggregator(
+    AGGREGATE_MODEL, force_cpu=os.environ.get("APP_AGG_FORCE_CPU", False)
+)
 
 
 def aggregate_text(
@@ -364,10 +380,11 @@ def load_uploaded_file(file_obj, max_pages: int = 20, lower: bool = False) -> st
 def parse_args():
     """arguments for the command line interface"""
     parser = argparse.ArgumentParser(
-        description="Document Summarization with Long-Document Transformers Demo",
+        description="Document Summarization with Long-Document Transformers - Demo",
         formatter_class=argparse.ArgumentDefaultsHelpFormatter,
-        epilog="Runs a local-only web app to summarize documents. use --share for a public link to share.",
+        epilog="Runs a local-only web UI to summarize documents. pass --share for a public link to share.",
     )
+
     parser.add_argument(
         "--share",
         dest="share",
@@ -379,16 +396,34 @@ def parse_args():
         "--model",
         type=str,
         default=None,
-        help=f"Add a custom model to the list of models: {', '.join(MODEL_OPTIONS)}",
+        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 list of beam search options: {pp.pformat(BEAM_OPTIONS, compact=True)}",
+    )
+    parser.add_argument(
+        "-batch",
+        "--token_batch_option",
+        type=int,
+        default=None,
+        help=f"Add a token batch option to the list of token batch options: {pp.pformat(TOKEN_BATCH_OPTIONS, compact=True)}",
     )
     parser.add_argument(
         "-level",
-        "--log-level",
+        "--log_level",
         type=str,
         default="INFO",
         choices=["DEBUG", "INFO", "WARNING", "ERROR"],
         help="Set the logging level",
     )
+    # if "--help" in sys.argv or "-h" in sys.argv:
+    #     parser.print_help()
+    #     sys.exit(0)
+
     return parser.parse_args()
 
 
@@ -397,11 +432,19 @@ if __name__ == "__main__":
     logger = logging.getLogger(__name__)
     args = parse_args()
     logger.setLevel(args.log_level)
-    logger.info(f"args: {args}")
+    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)
-    logger.info("Starting app instance")
+    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)
+
     logger.info("Loading OCR model")
     with contextlib.redirect_stdout(None):
         ocr_model = ocr_predictor(
@@ -410,11 +453,14 @@ if __name__ == "__main__":
             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 with Long-Document Transformers")
     _examples = list(name_to_path.keys())
+    logger.info("Starting app instance")
     with demo:
         gr.Markdown("# Document Summarization with Long-Document Transformers")
         gr.Markdown(
@@ -436,9 +482,9 @@ if __name__ == "__main__":
                         label="Model Name",
                     )
                     num_beams = gr.Radio(
-                        choices=[2, 3, 4],
+                        choices=BEAM_OPTIONS,
                         label="Beam Search: # of Beams",
-                        value=2,
+                        value=BEAM_OPTIONS[0],
                     )
                     load_examples_button = gr.Button(
                         "Load Example in Dropdown",
@@ -542,9 +588,10 @@ if __name__ == "__main__":
                     step=0.05,
                 )
                 token_batch_length = gr.Radio(
-                    choices=[1024, 1536, 2048, 2560, 3072],
+                    choices=TOKEN_BATCH_OPTIONS,
                     label="token batch length",
-                    value=2048,
+                    # select median option
+                    value=TOKEN_BATCH_OPTIONS[len(TOKEN_BATCH_OPTIONS) // 2],
                 )
 
             with gr.Row(variant="compact"):