Support lib_path override in C++. Improvements on docs and error messages (mlc-ai#1086)

rickzx · web-flow · commit b0373d172e2e · 2023-10-19T12:09:12.000-07:00
* Support lib_path option in C++ CLI. Disable ChatConfig.model_lib override in Python API. Improvements on helper messages and error messages

* Update docs

* Rename lib_path -&gt; model_lib_path
diff --git a/cpp/cli_main.cc b/cpp/cli_main.cc
@@ -163,7 +163,7 @@ struct ModelPaths {
    */
   std::filesystem::path lib;
 
-  static ModelPaths Find(const std::string& device_name, const std::string& local_id);
+  static ModelPaths Find(const std::string& device_name, const std::string& local_id, const std::string& user_lib_path);
 };
 
 /*!
@@ -337,7 +337,7 @@ std::string ReadStringFromJSONFile(const std::filesystem::path& config_path,
   return config[key].get<std::string>();
 }
 
-ModelPaths ModelPaths::Find(const std::string& device_name, const std::string& local_id) {
+ModelPaths ModelPaths::Find(const std::string& device_name, const std::string& local_id, const std::string &user_lib_path) {
   // Step 1. Find config path
   std::filesystem::path config_path;
   if (auto path = TryInferMLCChatConfig(local_id)) {
@@ -368,26 +368,36 @@ ModelPaths ModelPaths::Find(const std::string& device_name, const std::string& l
   }
   std::cout << "Use model weights: " << params_json << std::endl;
   // Step 3. Find model lib path
-  std::string lib_local_id = ReadStringFromJSONFile(config_path, "model_lib");
-  std::string lib_name = lib_local_id + "-" + device_name;
   std::filesystem::path lib_path;
-  if (auto path = FindFile({lib_local_id,
+  if (!user_lib_path.empty()) {
+    lib_path = user_lib_path;
+    if (!std::filesystem::exists(lib_path) || !std::filesystem::is_regular_file(lib_path)) {
+        LOG(FATAL) << "The `lib_path` you passed in is not a file: " << user_lib_path << "\n";
+        exit(1);
+    }
+  } else {
+    std::string lib_local_id = ReadStringFromJSONFile(config_path, "model_lib");
+    std::string lib_name = lib_local_id + "-" + device_name;
+    if (auto path = FindFile({lib_local_id,
                             "dist/prebuilt/lib",  // Using prebuilt workflow
                             "dist/" + local_id, "dist/prebuilt/" + lib_local_id},
                            {
                                lib_name + GetArchSuffix(),
                                lib_name,
                            },
                            GetLibSuffixes())) {
-    lib_path = path.value();
-  } else {
-    LOG(FATAL) << "Cannot find the model library that corresponds to `" << lib_local_id << "`.\n"
-               << "We searched over the following possible paths: \n"
-               << "- " + lib_local_id << "\n"
-               << "- dist/prebuilt/lib \n"
-               << "- dist/" + local_id << "\n"
-               << "- dist/prebuilt/" + lib_local_id;
-    exit(1);
+        lib_path = path.value();
+    } else {
+        LOG(FATAL) << "Cannot find the model library that corresponds to `" << lib_local_id << "`.\n"
+                   << "We searched over the following possible paths: \n"
+                   << "- " + lib_local_id << "\n"
+                   << "- dist/prebuilt/lib \n"
+                   << "- dist/" + local_id << "\n"
+                   << "- dist/prebuilt/" + lib_local_id << "\n"
+                   << "If you would like to directly specify the full model library path, you may "
+                   << "consider passing in the `--model-lib-path` argument.\n";
+        exit(1);
+    }
   }
   std::cout << "Use model library: " << lib_path << std::endl;
   return ModelPaths{config_path, params_json, lib_path};
@@ -427,8 +437,8 @@ void Converse(ChatModule* chat, const std::string& input, int stream_interval,
  * \param stream_interval The interval that should be used for streaming the response.
  */
 void Chat(ChatModule* chat, const std::string& device_name, std::string local_id,
-          int stream_interval = 2) {
-  ModelPaths model = ModelPaths::Find(device_name, local_id);
+          std::string lib_path, int stream_interval = 2) {
+  ModelPaths model = ModelPaths::Find(device_name, local_id, lib_path);
   PrintSpecialCommands();
   chat->Reload(model);
   chat->ProcessSystemPrompts();
@@ -456,7 +466,7 @@ void Chat(ChatModule* chat, const std::string& device_name, std::string local_id
       if (new_local_id.empty()) {
         new_local_id = local_id;
       }
-      model = ModelPaths::Find(device_name, new_local_id);
+      model = ModelPaths::Find(device_name, new_local_id, lib_path);
       chat->Reload(model);
       local_id = new_local_id;
     } else if (input.substr(0, 5) == "/help") {
@@ -470,7 +480,17 @@ void Chat(ChatModule* chat, const std::string& device_name, std::string local_id
 int main(int argc, char* argv[]) {
   argparse::ArgumentParser args("mlc_chat");
 
-  args.add_argument("--model");
+  args.add_description("MLCChat CLI is the command line tool to run MLC-compiled LLMs out of the box.\n"
+                       "Note: the --model argument is required. It can either be the model name with its "
+                       "quantization scheme or a full path to the model folder. In the former case, the "
+                       "provided name will be used to search for the model folder over possible paths. "
+                       "--model-lib-path argument is optional. If unspecified, the --model argument will be used "
+                       "to search for the library file over possible paths.");
+
+  args.add_argument("--model")
+      .help("[required] the model to use");
+  args.add_argument("--model-lib-path")
+      .help("[optional] the full path to the model library file to use");
   args.add_argument("--device").default_value("auto");
   args.add_argument("--evaluate").default_value(false).implicit_value(true);
   args.add_argument("--eval-prompt-len").default_value(128).scan<'i', int>();
@@ -485,6 +505,10 @@ int main(int argc, char* argv[]) {
   }
 
   std::string local_id = args.get<std::string>("--model");
+  std::string lib_path;
+  if (args.present("--model-lib-path")) {
+    lib_path = args.get<std::string>("--model-lib-path");
+  }
   auto [device_name, device_id] = DetectDevice(args.get<std::string>("--device"));
 
   try {
@@ -494,14 +518,14 @@ int main(int argc, char* argv[]) {
       // that are not supposed to be used in chat app setting
       int prompt_len = args.get<int>("--eval-prompt-len");
       int gen_len = args.get<int>("--eval-gen-len");
-      ModelPaths model = ModelPaths::Find(device_name, local_id);
+      ModelPaths model = ModelPaths::Find(device_name, local_id, lib_path);
       tvm::runtime::Module chat_mod = mlc::llm::CreateChatModule(GetDevice(device_name, device_id));
       std::string model_path = model.config.parent_path().string();
       tvm::runtime::Module lib = tvm::runtime::Module::LoadFromFile(model.lib.string());
       chat_mod.GetFunction("reload")(lib, tvm::String(model_path));
       chat_mod.GetFunction("evaluate")(prompt_len, gen_len);
     } else {
-      Chat(&chat, device_name, local_id);
+      Chat(&chat, device_name, local_id, lib_path);
     }
   } catch (const std::runtime_error& err) {
     std::cerr << err.what() << std::endl;
diff --git a/docs/deploy/cli.rst b/docs/deploy/cli.rst
@@ -111,6 +111,12 @@ Once ``mlc_chat_cli`` is installed, you are able to run any MLC-compiled model o
         - Model lib should be placed at ``./dist/prebuilt/lib/$(local_id)-$(arch).$(suffix)``.
         - Model weights and chat config are located under ``./dist/prebuilt/mlc-chat-$(local_id)/``.
 
+        .. note::
+          Please make sure that you have the same directory structure as above, because the CLI tool
+          relies on it to automatically search for model lib and weights. If you would like to directly
+          provide a full model lib path to override the auto-search, you can pass in a ``--model-lib-path`` argument
+          to the CLI
+
         .. collapse:: Example
 
           .. code:: shell
@@ -134,6 +140,12 @@ Once ``mlc_chat_cli`` is installed, you are able to run any MLC-compiled model o
         - Model libraries should be placed at ``./dist/$(local_id)/$(local_id)-$(arch).$(suffix)``.
         - Model weights and chat config are located under ``./dist/$(local_id)/params/``.
 
+        .. note::
+          Please make sure that you have the same directory structure as above, because the CLI tool
+          relies on it to automatically search for model lib and weights. If you would like to directly
+          provide a full model lib path to override the auto-search, you can pass in a ``--model-lib-path`` argument
+          to the CLI
+
         .. collapse:: Example
 
           .. code:: shell
diff --git a/docs/deploy/python.rst b/docs/deploy/python.rst
@@ -51,6 +51,11 @@ If you do not have the MLC-compiled ``model`` ready:
       - Model lib should be placed at ``./dist/prebuilt/lib/$(model)-$(arch).$(suffix)``.
       - Model weights and chat config are located under ``./dist/prebuilt/mlc-chat-$(model)/``.
 
+      .. note::
+         Please make sure that you have the same directory structure as above, because Python API
+         relies on it to automatically search for model lib and weights. If you would like to directly
+         provide a full model lib path to override the auto-search, you can specify ``ChatModule.model_lib_path``
+
       .. collapse:: Example
 
          .. code:: shell
@@ -74,6 +79,11 @@ If you do not have the MLC-compiled ``model`` ready:
       - Model libraries should be placed at ``./dist/$(model)/$(model)-$(arch).$(suffix)``.
       - Model weights and chat config are located under ``./dist/$(model)/params/``.
 
+      .. note::
+         Please make sure that you have the same directory structure as above, because Python API
+         relies on it to automatically search for model lib and weights. If you would like to directly
+         provide a full model lib path to override the auto-search, you can specify ``ChatModule.model_lib_path``
+
       .. collapse:: Example
 
          .. code:: shell
@@ -157,7 +167,7 @@ You can also checkout the :doc:`/prebuilt_models` page to run other models.
 |
 
 .. note:: 
-   You could also specify the address of ``model`` and ``lib_path`` explicitly. If
+   You could also specify the address of ``model`` and ``model_lib_path`` explicitly. If
    you only specify ``model`` as ``model_name`` and ``quantize_mode``, we will
    do a search for you. See more in the documentation of :meth:`mlc_chat.ChatModule.__init__`.
 
diff --git a/python/mlc_chat/chat_module.py b/python/mlc_chat/chat_module.py
@@ -5,6 +5,7 @@
 import logging
 import os
 import sys
+import warnings
 from dataclasses import asdict, dataclass, fields
 from enum import Enum
 from typing import List, Optional
@@ -351,6 +352,12 @@ def _get_chat_config(config_file_path: str, user_chat_config: Optional[ChatConfi
         # We override using user's chat config
         for field in fields(user_chat_config):
             field_name = field.name
+            if field_name == 'model_lib':
+                warn_msg = ('WARNING: Do not override "model_lib" in ChatConfig. '
+                            'This override will be ignored. '
+                            'Please use ChatModule.model_lib_path to override the full model library path instead.')
+                warnings.warn(warn_msg)
+                continue
             field_value = getattr(user_chat_config, field_name)
             if field_value is not None:
                 setattr(final_chat_config, field_name, field_value)
@@ -389,7 +396,7 @@ def _get_lib_module_path(
     model: str,
     model_path: str,
     chat_config: ChatConfig,
-    lib_path: Optional[str],
+    model_lib_path: Optional[str],
     device_name: str,
     config_file_path: str,
 ) -> str:
@@ -403,7 +410,7 @@ def _get_lib_module_path(
         Model path found by `_get_model_path`.
     chat_config : ChatConfig
         Chat config after potential overrides. Returned by ``_get_chat_config``.
-    lib_path : Optional[str]
+    model_lib_path : Optional[str]
         User's input. Supposedly a full path to model library. Prioritized to use.
     device_name : str
         User's input. Used to construct the library model file name.
@@ -412,21 +419,21 @@ def _get_lib_module_path(
 
     Returns
     ------
-    lib_path : str
+    model_lib_path : str
         The path pointing to the model library we find.
 
     Raises
     ------
     FileNotFoundError: if we cannot find a valid model library file.
     """
-    # 1. Use user's lib_path if provided
-    if lib_path is not None:
-        if os.path.isfile(lib_path):
-            logging.info(f"Using library model: {lib_path}")
-            return lib_path
+    # 1. Use user's model_lib_path if provided
+    if model_lib_path is not None:
+        if os.path.isfile(model_lib_path):
+            logging.info(f"Using library model: {model_lib_path}")
+            return model_lib_path
         else:
             err_msg = (
-                f"The `lib_path` you passed in is not a file: {lib_path}.\nPlease checkout "
+                f"The `model_lib_path` you passed in is not a file: {model_lib_path}.\nPlease checkout "
                 f"{_PYTHON_GET_STARTED_TUTORIAL_URL} for an example on how to load a model."
             )
             raise FileNotFoundError(err_msg)
@@ -482,7 +489,7 @@ def _get_lib_module_path(
         err_msg += f"- {candidate}\n"
     err_msg += (
         "If you would like to directly specify the model library path, you may "
-        "consider passing in the `lib_path` parameter.\n"
+        "consider passing in the `ChatModule.model_lib_path` parameter.\n"
         f"Please checkout {_PYTHON_GET_STARTED_TUTORIAL_URL} for an example "
         "on how to load a model."
     )
@@ -659,7 +666,7 @@ class ChatModule:
         A ``ChatConfig`` instance partially filled. Will be used to override the
         ``mlc-chat-config.json``.
 
-    lib_path : Optional[str]
+    model_lib_path : Optional[str]
         The full path to the model library file to use (e.g. a ``.so`` file).
         If unspecified, we will use the provided ``model`` to search over
         possible paths.
@@ -670,7 +677,7 @@ def __init__(
         model: str,
         device: str = "auto",
         chat_config: Optional[ChatConfig] = None,
-        lib_path: Optional[str] = None,
+        model_lib_path: Optional[str] = None,
     ):
         device_err_msg = (
             f"Invalid device name: {device}. Please enter the device in the form "
@@ -732,15 +739,15 @@ def __init__(
         self.chat_config = _get_chat_config(self.config_file_path, chat_config)
 
         # 5. Look up model library
-        self.lib_path = _get_lib_module_path(
-            model, self.model_path, self.chat_config, lib_path, device_name, self.config_file_path
+        self.model_lib_path = _get_lib_module_path(
+            model, self.model_path, self.chat_config, model_lib_path, device_name, self.config_file_path
         )
 
         # 6. Call reload
         user_chat_config_json_str = _convert_chat_config_to_json_str(
             self.chat_config, self.chat_config.conv_template
         )
-        self._reload(self.lib_path, self.model_path, user_chat_config_json_str)
+        self._reload(self.model_lib_path, self.model_path, user_chat_config_json_str)
 
     def generate(
         self,
