Initial commit

llama_toolchain/cli/__init__.py

@@ -0,0 +1,5 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.

llama_toolchain/cli/download.py

@@ -0,0 +1,97 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+import argparse
+import os
+import textwrap
+from pathlib import Path
+
+from huggingface_hub import snapshot_download
+from huggingface_hub.utils import GatedRepoError, RepositoryNotFoundError
+
+from llama_toolchain.cli.subcommand import Subcommand
+from llama_toolchain.utils import DEFAULT_DUMP_DIR
+
+
+DEFAULT_CHECKPOINT_DIR = os.path.join(DEFAULT_DUMP_DIR, "checkpoints")
+
+
+class Download(Subcommand):
+    """Llama cli for downloading llama toolchain assets"""
+
+    def __init__(self, subparsers: argparse._SubParsersAction):
+        super().__init__()
+        self.parser = subparsers.add_parser(
+            "download",
+            prog="llama download",
+            description="Download a model from the Hugging Face Hub",
+            epilog=textwrap.dedent(
+                """\
+            # Here are some examples on how to use this command:
+
+            llama download --repo-id meta-llama/Llama-2-7b-hf --hf-token <HF_TOKEN>
+            llama download --repo-id meta-llama/Llama-2-7b-hf --output-dir /data/my_custom_dir --hf-token <HF_TOKEN>
+            HF_TOKEN=<HF_TOKEN> llama download --repo-id meta-llama/Llama-2-7b-hf
+
+            The output directory will be used to load models and tokenizers for inference.
+            """
+            ),
+            formatter_class=argparse.RawTextHelpFormatter,
+        )
+        self._add_arguments()
+        self.parser.set_defaults(func=self._run_download_cmd)
+
+    def _add_arguments(self):
+        self.parser.add_argument(
+            "repo_id",
+            type=str,
+            help="Name of the repository on Hugging Face Hub eg. llhf/Meta-Llama-3.1-70B-Instruct",
+        )
+        self.parser.add_argument(
+            "--hf-token",
+            type=str,
+            required=False,
+            default=os.getenv("HF_TOKEN", None),
+            help="Hugging Face API token. Needed for gated models like Llama2. Will also try to read environment variable `HF_TOKEN` as default.",
+        )
+        self.parser.add_argument(
+            "--ignore-patterns",
+            type=str,
+            required=False,
+            default="*.safetensors",
+            help="If provided, files matching any of the patterns are not downloaded. Defaults to ignoring "
+            "safetensors files to avoid downloading duplicate weights.",
+        )
+
+    def _run_download_cmd(self, args: argparse.Namespace):
+        model_name = args.repo_id.split("/")[-1]
+        output_dir = Path(DEFAULT_CHECKPOINT_DIR) / model_name
+        os.makedirs(output_dir, exist_ok=True)
+        try:
+            true_output_dir = snapshot_download(
+                args.repo_id,
+                local_dir=output_dir,
+                # "auto" will download to cache_dir and symlink files to local_dir
+                # avoiding unnecessary duplicate copies
+                local_dir_use_symlinks="auto",
+                ignore_patterns=args.ignore_patterns,
+                token=args.hf_token,
+            )
+        except GatedRepoError:
+            self.parser.error(
+                "It looks like you are trying to access a gated repository. Please ensure you "
+                "have access to the repository and have provided the proper Hugging Face API token "
+                "using the option `--hf-token` or by running `huggingface-cli login`."
+                "You can find your token by visiting https://huggingface.co/settings/tokens"
+            )
+        except RepositoryNotFoundError:
+            self.parser.error(
+                f"Repository '{args.repo_id}' not found on the Hugging Face Hub."
+            )
+        except Exception as e:
+            self.parser.error(e)
+
+        print(f"Successfully downloaded model to {true_output_dir}")

llama_toolchain/cli/inference/__init__.py

@@ -0,0 +1,5 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.

llama_toolchain/cli/inference/configure.py

@@ -0,0 +1,91 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+import argparse
+import os
+import textwrap
+
+from pathlib import Path
+
+import pkg_resources
+
+from llama_toolchain.cli.subcommand import Subcommand
+from llama_toolchain.utils import DEFAULT_DUMP_DIR
+
+
+CONFIGS_BASE_DIR = os.path.join(DEFAULT_DUMP_DIR, "configs")
+
+
+class InferenceConfigure(Subcommand):
+    """Llama cli for configuring llama toolchain configs"""
+
+    def __init__(self, subparsers: argparse._SubParsersAction):
+        super().__init__()
+        self.parser = subparsers.add_parser(
+            "configure",
+            prog="llama inference configure",
+            description="Configure llama toolchain inference configs",
+            epilog=textwrap.dedent(
+                """
+                Example:
+                    llama inference configure
+                """
+            ),
+            formatter_class=argparse.RawTextHelpFormatter,
+        )
+        self._add_arguments()
+        self.parser.set_defaults(func=self._run_inference_configure_cmd)
+
+    def _add_arguments(self):
+        pass
+
+    def read_user_inputs(self):
+        checkpoint_dir = input(
+            "Enter the checkpoint directory for the model (e.g., ~/.llama/checkpoints/Meta-Llama-3-8B/): "
+        )
+        model_parallel_size = input(
+            "Enter model parallel size (e.g., 1 for 8B / 8 for 70B and 405B): "
+        )
+        assert model_parallel_size.isdigit() and int(model_parallel_size) in {
+            1,
+            8,
+        }, "model parallel size must be 1 or 8"
+
+        return checkpoint_dir, model_parallel_size
+
+    def write_output_yaml(self, checkpoint_dir, model_parallel_size, yaml_output_path):
+        default_conf_path = pkg_resources.resource_filename(
+            "llama_toolchain", "data/default_inference_config.yaml"
+        )
+        with open(default_conf_path, "r") as f:
+            yaml_content = f.read()
+
+        yaml_content = yaml_content.format(
+            checkpoint_dir=checkpoint_dir,
+            model_parallel_size=model_parallel_size,
+        )
+
+        with open(yaml_output_path, "w") as yaml_file:
+            yaml_file.write(yaml_content.strip())
+
+        print(f"YAML configuration has been written to {yaml_output_path}")
+
+    def _run_inference_configure_cmd(self, args: argparse.Namespace) -> None:
+        checkpoint_dir, model_parallel_size = self.read_user_inputs()
+        checkpoint_dir = os.path.expanduser(checkpoint_dir)
+
+        assert (
+            Path(checkpoint_dir).exists() and Path(checkpoint_dir).is_dir()
+        ), f"{checkpoint_dir} does not exist or it not a directory"
+
+        os.makedirs(CONFIGS_BASE_DIR, exist_ok=True)
+        yaml_output_path = Path(CONFIGS_BASE_DIR) / "inference.yaml"
+
+        self.write_output_yaml(
+            checkpoint_dir,
+            model_parallel_size,
+            yaml_output_path,
+        )

llama_toolchain/cli/inference/inference.py

@@ -0,0 +1,36 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+import argparse
+import textwrap
+
+from llama_toolchain.cli.inference.configure import InferenceConfigure
+from llama_toolchain.cli.inference.start import InferenceStart
+from llama_toolchain.cli.subcommand import Subcommand
+
+
+class InferenceParser(Subcommand):
+    """Llama cli for inference apis"""
+
+    def __init__(self, subparsers: argparse._SubParsersAction):
+        super().__init__()
+        self.parser = subparsers.add_parser(
+            "inference",
+            prog="llama inference",
+            description="Run inference on a llama model",
+            epilog=textwrap.dedent(
+                """
+                Example:
+                    llama inference start <options>
+                """
+            ),
+        )
+
+        subparsers = self.parser.add_subparsers(title="inference_subcommands")
+
+        # Add sub-commandsa
+        InferenceStart.create(subparsers)
+        InferenceConfigure.create(subparsers)

llama_toolchain/cli/inference/start.py

@@ -0,0 +1,57 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+import argparse
+import textwrap
+
+from llama_toolchain.cli.subcommand import Subcommand
+
+from llama_toolchain.inference.server import main as inference_server_init
+
+
+class InferenceStart(Subcommand):
+    """Llama Inference cli for starting inference server"""
+
+    def __init__(self, subparsers: argparse._SubParsersAction):
+        super().__init__()
+        self.parser = subparsers.add_parser(
+            "start",
+            prog="llama inference start",
+            description="Start an inference server",
+            epilog=textwrap.dedent(
+                """
+                Example:
+                    llama inference start <options>
+                """
+            ),
+            formatter_class=argparse.RawTextHelpFormatter,
+        )
+        self._add_arguments()
+        self.parser.set_defaults(func=self._run_inference_start_cmd)
+
+    def _add_arguments(self):
+        self.parser.add_argument(
+            "--port",
+            type=int,
+            help="Port to run the server on. Defaults to 5000",
+            default=5000,
+        )
+        self.parser.add_argument(
+            "--disable-ipv6",
+            action="store_true",
+            help="Disable IPv6 support",
+            default=False,
+        )
+        self.parser.add_argument(
+            "--config", type=str, help="Path to config file", default="inference"
+        )
+
+    def _run_inference_start_cmd(self, args: argparse.Namespace) -> None:
+        inference_server_init(
+            config_path=args.config,
+            port=args.port,
+            disable_ipv6=args.disable_ipv6,
+        )

llama_toolchain/cli/llama.py

@@ -0,0 +1,58 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+import argparse
+
+from llama_toolchain.cli.download import Download
+from llama_toolchain.cli.inference.inference import InferenceParser
+from llama_toolchain.cli.model.model import ModelParser
+
+
+class LlamaCLIParser:
+    """Defines CLI parser for Llama CLI"""
+
+    def __init__(self):
+        self.parser = argparse.ArgumentParser(
+            prog="llama",
+            description="Welcome to the LLama cli",
+            add_help=True,
+        )
+
+        # Default command is to print help
+        self.parser.set_defaults(func=lambda args: self.parser.print_help())
+
+        subparsers = self.parser.add_subparsers(title="subcommands")
+
+        # Add sub-commands
+        Download.create(subparsers)
+        InferenceParser.create(subparsers)
+        ModelParser.create(subparsers)
+
+        # Import sub-commands from agentic_system if they exist
+        try:
+            from llama_agentic_system.cli.subcommand_modules import SUBCOMMAND_MODULES
+
+            for module in SUBCOMMAND_MODULES:
+                module.create(subparsers)
+
+        except ImportError:
+            pass
+
+    def parse_args(self) -> argparse.Namespace:
+        return self.parser.parse_args()
+
+    def run(self, args: argparse.Namespace) -> None:
+        args.func(args)
+
+
+def main():
+    parser = LlamaCLIParser()
+    args = parser.parse_args()
+    parser.run(args)
+
+
+if __name__ == "__main__":
+    main()

llama_toolchain/cli/model/__init__.py

@@ -0,0 +1,5 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.

llama_toolchain/cli/model/model.py

@@ -0,0 +1,35 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+import argparse
+import textwrap
+
+from llama_toolchain.cli.model.template import ModelTemplate
+from llama_toolchain.cli.subcommand import Subcommand
+
+
+class ModelParser(Subcommand):
+    """Llama cli for model interface apis"""
+
+    def __init__(self, subparsers: argparse._SubParsersAction):
+        super().__init__()
+        self.parser = subparsers.add_parser(
+            "model",
+            prog="llama model",
+            description="Describe llama model interfaces",
+            epilog=textwrap.dedent(
+                """
+                Example:
+                    llama model <subcommand> <options>
+                """
+            ),
+        )
+
+        subparsers = self.parser.add_subparsers(title="model_subcommands")
+
+        # Add sub-commandsa
+        # ModelDescribe.create(subparsers)
+        ModelTemplate.create(subparsers)

llama_toolchain/cli/model/template.py

@@ -0,0 +1,57 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+import argparse
+import textwrap
+
+from llama_models.llama3_1.api.interface import (
+    list_jinja_templates,
+    render_jinja_template,
+)
+
+from llama_toolchain.cli.subcommand import Subcommand
+
+
+class ModelTemplate(Subcommand):
+    """Llama model cli for describe a model template (message formats)"""
+
+    def __init__(self, subparsers: argparse._SubParsersAction):
+        super().__init__()
+        self.parser = subparsers.add_parser(
+            "template",
+            prog="llama model template",
+            description="Show llama model message formats",
+            epilog=textwrap.dedent(
+                """
+                Example:
+                    llama model template <options>
+                """
+            ),
+            formatter_class=argparse.RawTextHelpFormatter,
+        )
+        self._add_arguments()
+        self.parser.set_defaults(func=self._run_model_template_cmd)
+
+    def _add_arguments(self):
+        self.parser.add_argument(
+            "-m",
+            "--model-family",
+            type=str,
+            default="llama3_1",
+            help="Model Family (llama3_1, llama3_X, etc.)",
+        )
+        self.parser.add_argument(
+            "--template",
+            type=str,
+            help="Usecase template name (system_message, user_message, assistant_message, tool_message)...",
+            required=False,
+        )
+
+    def _run_model_template_cmd(self, args: argparse.Namespace) -> None:
+        if args.template:
+            render_jinja_template(args.template)
+        else:
+            list_jinja_templates()

llama_toolchain/cli/subcommand.py

@@ -0,0 +1,19 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+
+class Subcommand:
+    """All llama cli subcommands must inherit from this class"""
+
+    def __init__(self, *args, **kwargs):
+        pass
+
+    @classmethod
+    def create(cls, *args, **kwargs):
+        return cls(*args, **kwargs)
+
+    def _add_arguments(self):
+        pass