Update summarize.py

summarize.py

@@ -20,15 +20,23 @@ def load_model_and_tokenizer(model_name: str) -> tuple:
     :param str model_name: the model name/ID on the hub
     :return tuple: a tuple containing the model and tokenizer
     """
+    MODEL_OPTIONS = {
+        "Model 1": "pszemraj/long-t5-tglobal-base-16384-book-summary",
+        "Model 2": "pszemraj/long-t5-tglobal-base-sci-simplify",
+        "Model 3": "pszemraj/long-t5-tglobal-base-sci-simplify-elife",
+        "Model 4": "pszemraj/long-t5-tglobal-base-16384-booksci-summary-v1",
+        "Model 5": "pszemraj/pegasus-x-large-book-summary",
+    }
+    selected_model_identifier = MODEL_OPTIONS.get(user_selected_model_name)
     device = "cuda" if torch.cuda.is_available() else "cpu"
     model = AutoModelForSeq2SeqLM.from_pretrained(
-        model_name,
+        selected_model_identifier,
     ).to(device)
     model = model.eval()
 
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    tokenizer = AutoTokenizer.from_pretrained(selected_model_identifier)
 
-    logging.info(f"Loaded model {model_name} to {device}")
+    logging.info(f"Loaded model {selected_model_identifier} to {device}")
 
     if validate_pytorch2():
         try: