Refactor common download logic into cli.py

cli.py

@@ -8,6 +8,7 @@
 from pathlib import Path
 
 from build.utils import allowable_dtype_names, allowable_params_table
+from download import download_and_convert, is_model_downloaded
 
 import torch
 
@@ -19,6 +20,18 @@ def check_args(args, name: str) -> None:
     pass
 
 
+# Handle CLI arguments that are common to a majority of subcommands.
+def handle_common_args(args) -> None:
+    # Handle model download. Skip this for download, since it has slightly
+    # different semantics.
+    if (
+        args.command != "download"
+        and args.model
+        and not is_model_downloaded(args.model, args.model_directory)
+    ):
+        download_and_convert(args.model, args.model_directory, args.hf_token)
+
+
 def add_arguments_for_chat(parser):
     # Only chat specific options should be here
     _add_arguments_common(parser)

eval.py

@@ -21,7 +21,6 @@
 from build.model import Transformer
 from build.utils import set_precision
 from cli import add_arguments, add_arguments_for_eval, arg_init
-from download import download_and_convert, is_model_downloaded
 from generate import encode_tokens, model_forward
 
 torch._dynamo.config.automatic_dynamic_shapes = True
@@ -221,10 +220,6 @@ def main(args) -> None:
 
     """
 
-    # If a named model was provided and not downloaded, download it.
-    if args.model and not is_model_downloaded(args.model, args.model_directory):
-        download_and_convert(args.model, args.model_directory, args.hf_token)
-
     builder_args = BuilderArgs.from_args(args)
     tokenizer_args = TokenizerArgs.from_args(args)
     quantize = args.quantize

export.py

@@ -20,7 +20,6 @@
 
 from build.utils import set_backend, set_precision, use_aoti_backend, use_et_backend
 from cli import add_arguments, add_arguments_for_export, arg_init, check_args
-from download import download_and_convert, is_model_downloaded
 from export_aoti import export_model as export_model_aoti
 
 try:
@@ -35,11 +34,6 @@
 
 
 def main(args):
-    # THIS BELONGS INTO CLI
-    # If a named model was provided and not downloaded, download it.
-    # if args.model and not is_model_downloaded(args.model, args.model_directory):
-    #    download_and_convert(args.model, args.model_directory, args.hf_token)
-
     builder_args = BuilderArgs.from_args(args)
     quantize = args.quantize
 

generate.py

@@ -26,7 +26,6 @@
 from build.model import Transformer
 from build.utils import device_sync, set_precision
 from cli import add_arguments, add_arguments_for_generate, arg_init, check_args
-from download import download_and_convert, is_model_downloaded
 
 logger = logging.getLogger(__name__)
 
@@ -586,10 +585,6 @@ def callback(x):
 
 
 def main(args):
-    # If a named model was provided and not downloaded, download it.
-    if args.model and not is_model_downloaded(args.model, args.model_directory):
-        download_and_convert(args.model, args.model_directory, args.hf_token)
-
     builder_args = BuilderArgs.from_args(args)
     speculative_builder_args = BuilderArgs.from_speculative_args(args)
     tokenizer_args = TokenizerArgs.from_args(args)

torchchat.py

@@ -19,6 +19,7 @@
     add_arguments_for_generate,
     arg_init,
     check_args,
+    handle_common_args,
 )
 
 default_device = "cpu"
@@ -97,6 +98,8 @@
         format="%(message)s", level=logging.DEBUG if args.verbose else logging.INFO
     )
 
+    handle_common_args(args)
+
     if args.command == "chat":
         # enable "chat"
         args.chat = True