load("@fbcode_macros//build_defs:python_library.bzl", "python_library")
load("@fbsource//xplat/executorch/backends/qualcomm/qnn_version.bzl", "get_qnn_library_version")
load("@fbcode_macros//build_defs:python_binary.bzl", "python_binary")
load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime")

oncall("executorch")

python_library(
    name = "static_llama",
    srcs = [
        "model/static_llama.py",
    ],
    deps = [
        "//caffe2:torch",
    ],
)

python_library(
    name = "llama_lib",
    srcs = ["llama.py"],
    deps = [
        "//executorch/examples/models/llama:source_transformation",
        "//caffe2:torch",
        "//executorch/backends/qualcomm/partition:partition",
        "//executorch/backends/qualcomm/quantizer:quantizer",
        "//executorch/devtools/backend_debug:delegation_info",
        "//executorch/devtools:lib",
        "//executorch/examples/models:models",
        "//executorch/examples/qualcomm/oss_scripts/llama:static_llama",
        "//executorch/examples/qualcomm:utils",
        "//executorch/extension/export_util:export_util",
        "//executorch/extension/llm/export:export_lib",
        "//executorch/extension/pybindings:aten_lib",
    ],
)

python_binary(
    name = "llama",
    main_function = "executorch.examples.qualcomm.oss_scripts.llama.llama.main",
    preload_deps = [
        "//executorch/extension/llm/custom_ops:model_sharding_py",
    ],
    deps = [
        ":llama_lib",
    ],
)

runtime.command_alias(
    name = "llama_qnn",
    env = {
        "LD_LIBRARY_PATH": "$(location fbsource//third-party/qualcomm/qnn/qnn-{0}:qnn_offline_compile_libs)".format(get_qnn_library_version()),
        # Place holder to pass the QNN_SDK_ROOT check in executorch/examples/qualcomm/utils.py
        "QNN_SDK_ROOT": "",
    },
    exe = ":llama",
)