[BE] Replace copypasted with loops for arg parsing (#787)

* [BE] Replace copypasted with loops for arg parsing

At least for defining subparsers

* Fix loose reference

Author: malfet

cli.py

@@ -27,63 +27,21 @@
 ).expanduser()
 
 
+KNOWN_VERBS = ["chat", "browser", "download", "generate", "eval", "export", "list", "remove", "where"]
+
 # Handle CLI arguments that are common to a majority of subcommands.
-def check_args(args, name: str) -> None:
+def check_args(args, verb: str) -> None:
     # Handle model download. Skip this for download, since it has slightly
     # different semantics.
     if (
-        name not in ["download", "list", "remove"]
+        verb not in ["download", "list", "remove"]
         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)
-
-
-def add_arguments_for_browser(parser):
-    # Only browser specific options should be here
-    _add_arguments_common(parser)
-
-
-def add_arguments_for_download(parser):
-    # Only download specific options should be here
-    _add_arguments_common(parser)
-
-
-def add_arguments_for_generate(parser):
-    # Only generate specific options should be here
-    _add_arguments_common(parser)
-
-
-def add_arguments_for_eval(parser):
-    # Only eval specific options should be here
-    _add_arguments_common(parser)
-
-
-def add_arguments_for_export(parser):
-    # Only export specific options should be here
-    _add_arguments_common(parser)
-
-
-def add_arguments_for_list(parser):
-    # Only list specific options should be here
-    _add_arguments_common(parser)
-
-
-def add_arguments_for_remove(parser):
-    # Only remove specific options should be here
-    _add_arguments_common(parser)
-
-def add_arguments_for_where(parser):
-    # Only remove specific options should be here
-    _add_arguments_common(parser)
-
-
-def _add_arguments_common(parser):
+def add_arguments_for_verb(parser, verb: str):
     # Model specification. TODO Simplify this.
     # A model can be specified using a positional model name or HuggingFace
     # path. Alternatively, the model can be specified via --gguf-path or via
@@ -316,7 +274,7 @@ def arg_init(args):
         )
 
     if sys.version_info.major != 3 or sys.version_info.minor < 10:
-        raise RuntimeError("Please use Python 3.10 or later.")
+       raise RuntimeError("Please use Python 3.10 or later.")
 
     if hasattr(args, "quantize") and Path(args.quantize).is_file():
         with open(args.quantize, "r") as f:

eval.py

@@ -20,7 +20,7 @@
 
 from build.model import Transformer
 from build.utils import set_precision
-from cli import add_arguments_for_eval, arg_init
+from cli import add_arguments_for_verb, arg_init
 from generate import encode_tokens, model_forward
 
 torch._dynamo.config.automatic_dynamic_shapes = True
@@ -278,7 +278,7 @@ def main(args) -> None:
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="torchchat eval CLI")
-    add_arguments_for_eval(parser)
+    add_arguments_for_verb(parser, "eval")
     args = parser.parse_args()
     args = arg_init(args)
     main(args)

export.py

@@ -19,7 +19,7 @@
 )
 
 from build.utils import set_backend, set_precision
-from cli import add_arguments_for_export, arg_init, check_args
+from cli import add_arguments_for_verb, arg_init, check_args
 from export_aoti import export_model as export_model_aoti
 
 try:
@@ -112,7 +112,7 @@ def main(args):
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="torchchat export CLI")
-    add_arguments_for_export(parser)
+    add_arguments_for_verb(parser, "export")
     args = parser.parse_args()
     check_args(args, "export")
     args = arg_init(args)

generate.py

@@ -24,7 +24,7 @@
 )
 from build.model import Transformer
 from build.utils import device_sync, set_precision
-from cli import add_arguments_for_generate, arg_init, check_args, logger
+from cli import add_arguments_for_verb, arg_init, check_args, logger
 
 B_INST, E_INST = "[INST]", "[/INST]"
 B_SYS, E_SYS = "<<SYS>>", "<</SYS>>"
@@ -801,8 +801,9 @@ def main(args):
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="torchchat generate CLI")
-    add_arguments_for_generate(parser)
+    verb = "generate"
+    add_arguments_for_verb(parser, verb)
     args = parser.parse_args()
-    check_args(args, "generate")
+    check_args(args, verb)
     args = arg_init(args)
     main(args)

torchchat.py

@@ -10,15 +10,8 @@
 import sys
 
 from cli import (
-    add_arguments_for_browser,
-    add_arguments_for_chat,
-    add_arguments_for_download,
-    add_arguments_for_eval,
-    add_arguments_for_export,
-    add_arguments_for_generate,
-    add_arguments_for_list,
-    add_arguments_for_remove,
-    add_arguments_for_where,
+    add_arguments_for_verb,
+    KNOWN_VERBS,
     arg_init,
     check_args,
 )
@@ -39,59 +32,20 @@
     )
     subparsers.required = True
 
-    parser_chat = subparsers.add_parser(
-        "chat",
-        help="Chat interactively with a model",
-    )
-    add_arguments_for_chat(parser_chat)
-
-    parser_browser = subparsers.add_parser(
-        "browser",
-        help="Chat interactively in a browser",
-    )
-    add_arguments_for_browser(parser_browser)
-
-    parser_download = subparsers.add_parser(
-        "download",
-        help="Download a model from Hugging Face or others",
-    )
-    add_arguments_for_download(parser_download)
-
-    parser_generate = subparsers.add_parser(
-        "generate",
-        help="Generate responses from a model given a prompt",
-    )
-    add_arguments_for_generate(parser_generate)
-
-    parser_eval = subparsers.add_parser(
-        "eval",
-        help="Evaluate a model given a prompt",
-    )
-    add_arguments_for_eval(parser_eval)
-
-    parser_export = subparsers.add_parser(
-        "export",
-        help="Export a model for AOT Inductor or ExecuTorch",
-    )
-    add_arguments_for_export(parser_export)
-
-    parser_list = subparsers.add_parser(
-        "list",
-        help="List supported models",
-    )
-    add_arguments_for_list(parser_list)
-
-    parser_remove = subparsers.add_parser(
-        "remove",
-        help="Remove downloaded model artifacts",
-    )
-    add_arguments_for_remove(parser_remove)
-
-    parser_where = subparsers.add_parser(
-        "where",
-        help="Return directory containing downloaded model artifacts",
-    )
-    add_arguments_for_where(parser_where)
+    VERB_HELP = {
+       "chat": "Chat interactively with a model",
+       "browser": "Chat interactively in a browser",
+       "download": "Download a model from Hugging Face or others",
+       "generate": "Generate responses from a model given a prompt",
+       "eval": "Evaluate a model given a prompt",
+       "export": "Export a model for AOT Inductor or ExecuTorch",
+       "list": "List supported models",
+       "remove": "Remove downloaded model artifacts",
+       "where": "Return directory containing downloaded model artifacts",
+    }
+    for verb in KNOWN_VERBS:
+       subparser = subparsers.add_parser(verb, help=VERB_HELP[verb])
+       add_arguments_for_verb(subparser, verb)
 
     # Now parse the arguments
     args = parser.parse_args()