Allow specifying model name in build (octoml#161)

Previously when we introducing the support of "local model path" and
"HuggingFace model path", we removed the support of specifying model
name when build. This turns out preventing us from only specifying the
much shorter model name, as well as doing automatic search for models
existing on the disk.

Therefore, this PR brings back the argparse support for `--model`. Now
we will carefully handle the case where both the model name and one of
the model path / HF path are specified. We now also support model
searching on disk, and support specifying the model only by a short
model name.

README.md

@@ -90,6 +90,10 @@ If you have a local directory that has the model parameters, the tokenizer, and
 # Create the local build directory and compile the model
 mkdir build
 python build.py --model-path=/path/to/local/directory
+
+# If the model path is in the form of `dist/models/model_name`,
+# we can simplify the build command to
+# python build.py --model=model_name
 ```
 
 Similarly, the compiled model will be available at `dist/dolly-v2-3b-q3f16_0`, where the exact path will vary depending on your model type and specified quantization. Follow the platform specific instructions to build and run MLC LLM for [iOS](https://github.com/mlc-ai/mlc-llm/blob/main/ios/README.md), [Android](https://github.com/mlc-ai/mlc-llm/blob/main/android/README.md), and [CLI](https://github.com/mlc-ai/mlc-llm/tree/main/cpp/README.md).

android/README.md

@@ -38,11 +38,12 @@ We are excited to share that we have enabled the Android support for MLC-LLM. Ch
     git clone https://github.com/mlc-ai/mlc-llm.git --recursive
     cd mlc-llm
 
-    # From Hugging Face URL
-    python3 build.py --hf-path databricks/dolly-v2-3b --quantization q4f16_0 --target android --max-seq-len 768
-
     # From local directory
     python3 build.py --model-path path/to/vicuna-v1-7b --quantization q4f16_0 --target android --max-seq-len 768
+
+    # If the model path is `dist/models/vicuna-v1-7b`,
+    # we can simplify the build command to
+    # python build.py --model=vicuna-v1-7b --quantization q4f16_0 --target android --max-seq-len 768
     ```
 
 5. Build libraries for Android app.

build.py

@@ -16,6 +16,14 @@
 
 def _parse_args():
     args = argparse.ArgumentParser()
+    args.add_argument(
+        "--model",
+        type=str,
+        default="auto",
+        help='The name of the model to build. If it is "auto", we will automatically set the '
+        'model name according to "--model-path", "hf-path" or the model folders under '
+        '"--artifact-path/models"',
+    )
     args.add_argument(
         "--model-path",
         type=str,
@@ -81,17 +89,27 @@ def _parse_args():
 
 
 def _setup_model_path(args):
-    if args.model_path and args.hf_path:
-        assert (args.model_path and not args.hf_path) or (
-            args.hf_path and not args.model_path
-        ), "You cannot specify both a model path and a HF path. Please select one to specify."
+    assert not (
+        args.model_path and args.hf_path
+    ), "You cannot specify both a model path and a HF path. Please select one to specify."
+
     if args.model_path:
-        validate_config(args)
-        with open(os.path.join(args.model_path, "config.json")) as f:
-            config = json.load(f)
-            args.model = config["_name_or_path"].split("/")[-1]
+        if args.model != "auto":
+            assert args.model == os.path.basename(args.model_path), (
+                'When both "--model" and "--model-path" is specified, the '
+                'value of "--model" is required to match the basename of "--model-path"'
+            )
+        else:
+            args.model = os.path.basename(args.model_path)
+        validate_config(args.model_path)
     elif args.hf_path:
-        args.model = args.hf_path.split("/")[-1]
+        if args.model != "auto":
+            assert args.model == os.path.basename(args.hf_path), (
+                'When both "--model" and "--hf-path" is specified, the '
+                'value of "--model" is required to match the basename of "--hf-path"'
+            )
+        else:
+            args.model = os.path.basename(args.hf_path)
         args.model_path = os.path.join(args.artifact_path, "models", args.model)
         if os.path.exists(args.model_path):
             print(f"Weights exist at {args.model_path}, skipping download.")
@@ -102,18 +120,40 @@ def _setup_model_path(args):
                 f"git clone https://huggingface.co/{args.hf_path} {args.model_path}"
             )
             print(f"Downloaded weights to {args.model_path}")
-        validate_config(args)
+        validate_config(args.model_path)
+    elif args.model != "auto":
+        args.model_path = os.path.join(args.artifact_path, "models", args.model)
+        validate_config(args.model_path)
     else:
-        raise ValueError(f"Please specify either the model_path or the hf_path.")
+        lookup_path = os.path.join(args.artifact_path, "models")
+        print(
+            'None of "--model", "--model-path" and "--hf-path" is specified. Searching '
+            f"in {lookup_path} for existing models."
+        )
+        for dirname in os.listdir(lookup_path):
+            if os.path.isdir(os.path.join(lookup_path, dirname)) and os.path.isfile(
+                os.path.join(lookup_path, dirname, "config.json")
+            ):
+                try:
+                    validate_config(os.path.join(lookup_path, dirname))
+                except:
+                    pass
+                else:
+                    args.model_path = os.path.join(lookup_path, dirname)
+                    args.model = dirname
+                    break
+        if args.model == "auto":
+            raise ValueError(f"Please specify either the model_path or the hf_path.")
+
     print(f"Using model path {args.model_path}")
     return args
 
 
-def validate_config(args):
+def validate_config(model_path: str):
     assert os.path.exists(
-        os.path.join(args.model_path, "config.json")
+        os.path.join(model_path, "config.json")
     ), "Model path must contain valid config file."
-    with open(os.path.join(args.model_path, "config.json")) as f:
+    with open(os.path.join(model_path, "config.json")) as f:
         config = json.load(f)
         assert ("model_type" in config) and (
             "_name_or_path" in config

cpp/README.md

@@ -7,10 +7,14 @@
     cd mlc-llm
 
     # From Hugging Face URL
-    python3 build.py --hf-path databricks/dolly-v2-3b --quantization q3f16_0 --target iphone --max-seq-len 768
+    python3 build.py --hf-path databricks/dolly-v2-3b --quantization q3f16_0 --max-seq-len 768
 
     # From local directory
-    python3 build.py --model-path path/to/vicuna-v1-7b --quantization q3f16_0 --target iphone --max-seq-len 768
+    python3 build.py --model-path path/to/vicuna-v1-7b --quantization q3f16_0 --max-seq-len 768
+
+    # If the model path is in the form of `dist/models/model_name`,
+    # we can simplify the build command to
+    # python build.py --model model_name --quantization q3f16_0 --max-seq-len 768
     ```
 
 2. Build the CLI.

ios/README.md

@@ -26,11 +26,12 @@ Note: You will need Apple Developer Account to build iOS App locally.
     git clone https://github.com/mlc-ai/mlc-llm.git --recursive
     cd mlc-llm
 
-    # From Hugging Face URL
-    python3 build.py --hf-path databricks/dolly-v2-3b --quantization q3f16_0 --target iphone --max-seq-len 768
-
     # From local directory
     python3 build.py --model-path path/to/vicuna-v1-7b --quantization q3f16_0 --target iphone --max-seq-len 768
+
+    # If the model path is `dist/models/vicuna-v1-7b`,
+    # we can simplify the build command to
+    # python build.py --model vicuna-v1-7b --quantization q3f16_0 --target iphone --max-seq-len 768
     ```
 
 3. Prepare lib and params