From 48796ecd77ede240587a5ab9856badfb16d886f6 Mon Sep 17 00:00:00 2001 From: Meng Zhang Date: Sun, 28 May 2023 14:36:11 -0700 Subject: [PATCH] feat: add `tabby download` command (#157) * simplify fmt-display * cleanup * move tabby-admin to reduce nest * add model downloader * get rid of model-type * improve commands * fix fmt --- Cargo.lock | 148 +++++++++++++++++- .../ctranslate2-bindings/src/ctranslate2.cc | 10 +- .../admin => tabby-admin}/.gitattributes | 0 .../serve/admin => tabby-admin}/.gitignore | 0 .../src/serve/admin => tabby-admin}/README.md | 0 .../dist/assets/index.91561b80.js | 0 .../dist/assets/index.9874f585.css | 0 .../images/authentication/create-account.jpg | Bin .../dist/images/authentication/login.jpg | Bin .../images/authentication/reset-password.jpg | Bin .../dist/images/logo.svg | 0 .../dist/images/products/apple-imac-1.png | 0 .../dist/images/products/apple-imac-2.png | 0 .../dist/images/products/apple-imac-3.png | 0 .../dist/images/users/bonnie-green-2x.png | 0 .../dist/images/users/bonnie-green.png | 0 .../dist/images/users/helene-engels.png | 0 .../dist/images/users/jese-leos-2x.png | 0 .../dist/images/users/jese-leos.png | 0 .../dist/images/users/joseph-mcfall.png | 0 .../dist/images/users/lana-byrd.png | 0 .../dist/images/users/leslie-livingston.png | 0 .../dist/images/users/michael-gough.png | 0 .../dist/images/users/neil-sims.png | 0 .../dist/images/users/robert-brown.png | 0 .../dist/images/users/roberta-casas-2x.png | 0 .../dist/images/users/roberta-casas.png | 0 .../dist/images/users/thomas-lean.png | 0 .../admin => tabby-admin}/dist/index.html | 0 .../serve/admin => tabby-admin}/index.html | 0 .../serve/admin => tabby-admin}/package.json | 0 .../admin => tabby-admin}/postcss.config.mjs | 0 .../images/authentication/create-account.jpg | Bin .../public/images/authentication/login.jpg | Bin .../images/authentication/reset-password.jpg | Bin .../public/images/logo.svg | 0 .../public/images/products/apple-imac-1.png | 0 .../public/images/products/apple-imac-2.png | 0 .../public/images/products/apple-imac-3.png | 0 .../public/images/users/bonnie-green-2x.png | 0 .../public/images/users/bonnie-green.png | 0 .../public/images/users/helene-engels.png | 0 .../public/images/users/jese-leos-2x.png | 0 .../public/images/users/jese-leos.png | 0 .../public/images/users/joseph-mcfall.png | 0 .../public/images/users/lana-byrd.png | 0 .../public/images/users/leslie-livingston.png | 0 .../public/images/users/michael-gough.png | 0 .../public/images/users/neil-sims.png | 0 .../public/images/users/robert-brown.png | 0 .../public/images/users/roberta-casas-2x.png | 0 .../public/images/users/roberta-casas.png | 0 .../public/images/users/thomas-lean.png | 0 .../src/components/navbar.tsx | 0 .../src/components/sidebar.tsx | 0 .../src/flowbite-theme.ts | 0 .../serve/admin => tabby-admin}/src/index.css | 0 .../serve/admin => tabby-admin}/src/index.tsx | 0 .../src/layouts/navbar-sidebar.tsx | 0 .../src/pages/authentication/sign-in.tsx | 0 .../src/pages/authentication/sign-up.tsx | 0 .../src/pages/e-commerce/products.tsx | 0 .../admin => tabby-admin}/src/pages/index.tsx | 0 .../src/pages/users/list.tsx | 0 .../admin => tabby-admin}/src/svgmap.d.ts | 0 .../admin => tabby-admin}/tailwind.config.cjs | 0 .../serve/admin => tabby-admin}/tsconfig.json | 0 .../admin => tabby-admin}/vite.config.ts | 0 .../src/serve/admin => tabby-admin}/yarn.lock | 0 crates/tabby/Cargo.toml | 5 + crates/tabby/src/download.rs | 99 ++++++++++++ crates/tabby/src/main.rs | 15 +- crates/tabby/src/serve/admin.rs | 2 +- crates/tabby/src/serve/completions.rs | 37 ++++- crates/tabby/src/serve/mod.rs | 116 ++++++-------- 75 files changed, 351 insertions(+), 81 deletions(-) rename crates/{tabby/src/serve/admin => tabby-admin}/.gitattributes (100%) rename crates/{tabby/src/serve/admin => tabby-admin}/.gitignore (100%) rename crates/{tabby/src/serve/admin => tabby-admin}/README.md (100%) rename crates/{tabby/src/serve/admin => tabby-admin}/dist/assets/index.91561b80.js (100%) rename crates/{tabby/src/serve/admin => tabby-admin}/dist/assets/index.9874f585.css (100%) rename crates/{tabby/src/serve/admin => tabby-admin}/dist/images/authentication/create-account.jpg (100%) rename crates/{tabby/src/serve/admin => tabby-admin}/dist/images/authentication/login.jpg (100%) rename crates/{tabby/src/serve/admin => tabby-admin}/dist/images/authentication/reset-password.jpg (100%) rename crates/{tabby/src/serve/admin => tabby-admin}/dist/images/logo.svg (100%) rename crates/{tabby/src/serve/admin => tabby-admin}/dist/images/products/apple-imac-1.png (100%) rename crates/{tabby/src/serve/admin => tabby-admin}/dist/images/products/apple-imac-2.png (100%) rename crates/{tabby/src/serve/admin => tabby-admin}/dist/images/products/apple-imac-3.png (100%) rename crates/{tabby/src/serve/admin => tabby-admin}/dist/images/users/bonnie-green-2x.png (100%) rename crates/{tabby/src/serve/admin => tabby-admin}/dist/images/users/bonnie-green.png (100%) rename crates/{tabby/src/serve/admin => tabby-admin}/dist/images/users/helene-engels.png (100%) rename crates/{tabby/src/serve/admin => tabby-admin}/dist/images/users/jese-leos-2x.png (100%) rename crates/{tabby/src/serve/admin => tabby-admin}/dist/images/users/jese-leos.png (100%) rename crates/{tabby/src/serve/admin => tabby-admin}/dist/images/users/joseph-mcfall.png (100%) rename crates/{tabby/src/serve/admin => tabby-admin}/dist/images/users/lana-byrd.png (100%) rename crates/{tabby/src/serve/admin => tabby-admin}/dist/images/users/leslie-livingston.png (100%) rename crates/{tabby/src/serve/admin => tabby-admin}/dist/images/users/michael-gough.png (100%) rename crates/{tabby/src/serve/admin => tabby-admin}/dist/images/users/neil-sims.png (100%) rename crates/{tabby/src/serve/admin => tabby-admin}/dist/images/users/robert-brown.png (100%) rename crates/{tabby/src/serve/admin => tabby-admin}/dist/images/users/roberta-casas-2x.png (100%) rename crates/{tabby/src/serve/admin => tabby-admin}/dist/images/users/roberta-casas.png (100%) rename crates/{tabby/src/serve/admin => tabby-admin}/dist/images/users/thomas-lean.png (100%) rename crates/{tabby/src/serve/admin => tabby-admin}/dist/index.html (100%) rename crates/{tabby/src/serve/admin => tabby-admin}/index.html (100%) rename crates/{tabby/src/serve/admin => tabby-admin}/package.json (100%) rename crates/{tabby/src/serve/admin => tabby-admin}/postcss.config.mjs (100%) rename crates/{tabby/src/serve/admin => tabby-admin}/public/images/authentication/create-account.jpg (100%) rename crates/{tabby/src/serve/admin => tabby-admin}/public/images/authentication/login.jpg (100%) rename crates/{tabby/src/serve/admin => tabby-admin}/public/images/authentication/reset-password.jpg (100%) rename crates/{tabby/src/serve/admin => tabby-admin}/public/images/logo.svg (100%) rename crates/{tabby/src/serve/admin => tabby-admin}/public/images/products/apple-imac-1.png (100%) rename crates/{tabby/src/serve/admin => tabby-admin}/public/images/products/apple-imac-2.png (100%) rename crates/{tabby/src/serve/admin => tabby-admin}/public/images/products/apple-imac-3.png (100%) rename crates/{tabby/src/serve/admin => tabby-admin}/public/images/users/bonnie-green-2x.png (100%) rename crates/{tabby/src/serve/admin => tabby-admin}/public/images/users/bonnie-green.png (100%) rename crates/{tabby/src/serve/admin => tabby-admin}/public/images/users/helene-engels.png (100%) rename crates/{tabby/src/serve/admin => tabby-admin}/public/images/users/jese-leos-2x.png (100%) rename crates/{tabby/src/serve/admin => tabby-admin}/public/images/users/jese-leos.png (100%) rename crates/{tabby/src/serve/admin => tabby-admin}/public/images/users/joseph-mcfall.png (100%) rename crates/{tabby/src/serve/admin => tabby-admin}/public/images/users/lana-byrd.png (100%) rename crates/{tabby/src/serve/admin => tabby-admin}/public/images/users/leslie-livingston.png (100%) rename crates/{tabby/src/serve/admin => tabby-admin}/public/images/users/michael-gough.png (100%) rename crates/{tabby/src/serve/admin => tabby-admin}/public/images/users/neil-sims.png (100%) rename crates/{tabby/src/serve/admin => tabby-admin}/public/images/users/robert-brown.png (100%) rename crates/{tabby/src/serve/admin => tabby-admin}/public/images/users/roberta-casas-2x.png (100%) rename crates/{tabby/src/serve/admin => tabby-admin}/public/images/users/roberta-casas.png (100%) rename crates/{tabby/src/serve/admin => tabby-admin}/public/images/users/thomas-lean.png (100%) rename crates/{tabby/src/serve/admin => tabby-admin}/src/components/navbar.tsx (100%) rename crates/{tabby/src/serve/admin => tabby-admin}/src/components/sidebar.tsx (100%) rename crates/{tabby/src/serve/admin => tabby-admin}/src/flowbite-theme.ts (100%) rename crates/{tabby/src/serve/admin => tabby-admin}/src/index.css (100%) rename crates/{tabby/src/serve/admin => tabby-admin}/src/index.tsx (100%) rename crates/{tabby/src/serve/admin => tabby-admin}/src/layouts/navbar-sidebar.tsx (100%) rename crates/{tabby/src/serve/admin => tabby-admin}/src/pages/authentication/sign-in.tsx (100%) rename crates/{tabby/src/serve/admin => tabby-admin}/src/pages/authentication/sign-up.tsx (100%) rename crates/{tabby/src/serve/admin => tabby-admin}/src/pages/e-commerce/products.tsx (100%) rename crates/{tabby/src/serve/admin => tabby-admin}/src/pages/index.tsx (100%) rename crates/{tabby/src/serve/admin => tabby-admin}/src/pages/users/list.tsx (100%) rename crates/{tabby/src/serve/admin => tabby-admin}/src/svgmap.d.ts (100%) rename crates/{tabby/src/serve/admin => tabby-admin}/tailwind.config.cjs (100%) rename crates/{tabby/src/serve/admin => tabby-admin}/tsconfig.json (100%) rename crates/{tabby/src/serve/admin => tabby-admin}/vite.config.ts (100%) rename crates/{tabby/src/serve/admin => tabby-admin}/yarn.lock (100%) create mode 100644 crates/tabby/src/download.rs diff --git a/Cargo.lock b/Cargo.lock index 368e528..475c904 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2,6 +2,15 @@ # It is not intended for manual editing. version = 3 +[[package]] +name = "addr2line" +version = "0.19.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a76fd60b23679b7d19bd066031410fb7e458ccc5e958eb5c325888ce4baedc97" +dependencies = [ + "gimli", +] + [[package]] name = "adler" version = "1.0.2" @@ -152,6 +161,21 @@ dependencies = [ "tower-service", ] +[[package]] +name = "backtrace" +version = "0.3.67" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "233d376d6d185f2a3093e58f283f60f880315b6c60075b01f36b3b85154564ca" +dependencies = [ + "addr2line", + "cc", + "cfg-if", + "libc", + "miniz_oxide 0.6.2", + "object", + "rustc-demangle", +] + [[package]] name = "base64" version = "0.13.1" @@ -652,6 +676,16 @@ dependencies = [ "libc", ] +[[package]] +name = "error-chain" +version = "0.12.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2d2f06b9cac1506ece98fe3231e3cc9c4410ec3d5b1f24ae1c8946f0742cdefc" +dependencies = [ + "backtrace", + "version_check", +] + [[package]] name = "esaxx-rs" version = "0.1.8" @@ -689,7 +723,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3b9429470923de8e8cbd4d2dc513535400b4b3fef0319fb5c4e1f520a7bef743" dependencies = [ "crc32fast", - "miniz_oxide", + "miniz_oxide 0.7.1", ] [[package]] @@ -753,6 +787,17 @@ version = "0.3.28" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4fff74096e71ed47f8e023204cfd0aa1289cd54ae5430a9523be060cdb849964" +[[package]] +name = "futures-macro" +version = "0.3.28" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "89ca545a94061b6365f2c7355b4b32bd20df3ff95f02da9329b34ccc3bd6ee72" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.18", +] + [[package]] name = "futures-sink" version = "0.3.28" @@ -773,6 +818,8 @@ checksum = "26b01e40b772d54cf6c6d721c1d1abd0647a0106a12ecaa1c186273392a69533" dependencies = [ "futures-core", "futures-io", + "futures-macro", + "futures-sink", "futures-task", "memchr", "pin-project-lite", @@ -801,6 +848,12 @@ dependencies = [ "wasi", ] +[[package]] +name = "gimli" +version = "0.27.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ad0a93d233ebf96623465aad4046a8d3aa4da22d4f4beba5388838c8a434bbb4" + [[package]] name = "glob" version = "0.3.1" @@ -996,6 +1049,18 @@ dependencies = [ "regex", ] +[[package]] +name = "indicatif" +version = "0.17.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cef509aa9bc73864d6756f0d34d35504af3cf0844373afe9b8669a5b8005a729" +dependencies = [ + "console", + "number_prefix 0.4.0", + "portable-atomic 0.3.20", + "unicode-width", +] + [[package]] name = "inout" version = "0.1.3" @@ -1190,6 +1255,15 @@ version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a" +[[package]] +name = "miniz_oxide" +version = "0.6.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b275950c28b37e794e8c55d88aeb5e139d0ce23fdbbeda68f8d7174abdf9e8fa" +dependencies = [ + "adler", +] + [[package]] name = "miniz_oxide" version = "0.7.1" @@ -1282,6 +1356,15 @@ version = "0.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "830b246a0e5f20af87141b25c173cd1b609bd7779a4617d6ec582abaf90870f3" +[[package]] +name = "object" +version = "0.30.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ea86265d3d3dcb6a27fc51bd29a4bf387fae9d2986b823079d4986af253eb439" +dependencies = [ + "memchr", +] + [[package]] name = "once_cell" version = "1.17.1" @@ -1450,6 +1533,21 @@ version = "0.3.27" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "26072860ba924cbfa98ea39c8c19b4dd6a4a25423dbdf219c1eca91aa0cf6964" +[[package]] +name = "portable-atomic" +version = "0.3.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e30165d31df606f5726b090ec7592c308a0eaf61721ff64c9a3018e344a8753e" +dependencies = [ + "portable-atomic 1.3.2", +] + +[[package]] +name = "portable-atomic" +version = "1.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dc59d1bcc64fc5d021d67521f818db868368028108d37f0e98d74e33f68297b5" + [[package]] name = "ppv-lite86" version = "0.2.17" @@ -1642,10 +1740,12 @@ dependencies = [ "serde_urlencoded", "tokio", "tokio-native-tls", + "tokio-util", "tower-service", "url", "wasm-bindgen", "wasm-bindgen-futures", + "wasm-streams", "web-sys", "winreg", ] @@ -1689,6 +1789,12 @@ dependencies = [ "walkdir", ] +[[package]] +name = "rustc-demangle" +version = "0.1.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d626bb9dae77e28219937af045c257c28bfd3f69333c512553507f5f9798cb76" + [[package]] name = "rustix" version = "0.37.19" @@ -1903,6 +2009,28 @@ version = "0.10.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "73473c0e59e6d5812c5dfe2a064a6444949f089e20eec9a2e5506596494e4623" +[[package]] +name = "strum" +version = "0.24.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "063e6045c0e62079840579a7e47a355ae92f60eb74daaf156fb1e84ba164e63f" +dependencies = [ + "strum_macros", +] + +[[package]] +name = "strum_macros" +version = "0.24.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e385be0d24f186b4ce2f9982191e7101bb737312ad61c1f2f984f34bcf85d59" +dependencies = [ + "heck", + "proc-macro2", + "quote", + "rustversion", + "syn 1.0.109", +] + [[package]] name = "subtle" version = "2.5.0" @@ -1945,14 +2073,19 @@ dependencies = [ "clap", "ctranslate2-bindings", "env_logger", + "error-chain", + "futures-util", "hyper", + "indicatif 0.17.3", "lazy_static", "log", "mime_guess", "regex", + "reqwest", "rust-embed", "serde", "serde_json", + "strum", "tokio", "tower", "tower-http", @@ -2454,6 +2587,19 @@ version = "0.2.86" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ed9d5b4305409d1fc9482fee2d7f9bcbf24b3972bf59817ef757e23982242a93" +[[package]] +name = "wasm-streams" +version = "0.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6bbae3363c08332cadccd13b67db371814cd214c2524020932f0804b8cf7c078" +dependencies = [ + "futures-util", + "js-sys", + "wasm-bindgen", + "wasm-bindgen-futures", + "web-sys", +] + [[package]] name = "web-sys" version = "0.3.63" diff --git a/crates/ctranslate2-bindings/src/ctranslate2.cc b/crates/ctranslate2-bindings/src/ctranslate2.cc index 00c2b6f..915e367 100644 --- a/crates/ctranslate2-bindings/src/ctranslate2.cc +++ b/crates/ctranslate2-bindings/src/ctranslate2.cc @@ -91,9 +91,15 @@ std::unique_ptr create_engine( loader.device_indices = std::vector(device_indices.begin(), device_indices.end()); loader.num_replicas_per_device = num_replicas_per_device; - if (model_type_str == "decoder") { + if (loader.device == ctranslate2::Device::CPU) { + loader.compute_type = ctranslate2::ComputeType::INT8; + } else if (loader.device == ctranslate2::Device::CUDA) { + loader.compute_type = ctranslate2::ComputeType::FLOAT16; + } + + if (model_type_str == "AutoModelForCausalLM") { return DecoderImpl::create(loader); - } else if (model_type_str == "encoder-decoder") { + } else if (model_type_str == "AutoModelForSeq2SeqLM") { return EncoderDecoderImpl::create(loader); } else { return nullptr; diff --git a/crates/tabby/src/serve/admin/.gitattributes b/crates/tabby-admin/.gitattributes similarity index 100% rename from crates/tabby/src/serve/admin/.gitattributes rename to crates/tabby-admin/.gitattributes diff --git a/crates/tabby/src/serve/admin/.gitignore b/crates/tabby-admin/.gitignore similarity index 100% rename from crates/tabby/src/serve/admin/.gitignore rename to crates/tabby-admin/.gitignore diff --git a/crates/tabby/src/serve/admin/README.md b/crates/tabby-admin/README.md similarity index 100% rename from crates/tabby/src/serve/admin/README.md rename to crates/tabby-admin/README.md diff --git a/crates/tabby/src/serve/admin/dist/assets/index.91561b80.js b/crates/tabby-admin/dist/assets/index.91561b80.js similarity index 100% rename from crates/tabby/src/serve/admin/dist/assets/index.91561b80.js rename to crates/tabby-admin/dist/assets/index.91561b80.js diff --git a/crates/tabby/src/serve/admin/dist/assets/index.9874f585.css b/crates/tabby-admin/dist/assets/index.9874f585.css similarity index 100% rename from crates/tabby/src/serve/admin/dist/assets/index.9874f585.css rename to crates/tabby-admin/dist/assets/index.9874f585.css diff --git a/crates/tabby/src/serve/admin/dist/images/authentication/create-account.jpg b/crates/tabby-admin/dist/images/authentication/create-account.jpg similarity index 100% rename from crates/tabby/src/serve/admin/dist/images/authentication/create-account.jpg rename to crates/tabby-admin/dist/images/authentication/create-account.jpg diff --git a/crates/tabby/src/serve/admin/dist/images/authentication/login.jpg b/crates/tabby-admin/dist/images/authentication/login.jpg similarity index 100% rename from crates/tabby/src/serve/admin/dist/images/authentication/login.jpg rename to crates/tabby-admin/dist/images/authentication/login.jpg diff --git a/crates/tabby/src/serve/admin/dist/images/authentication/reset-password.jpg b/crates/tabby-admin/dist/images/authentication/reset-password.jpg similarity index 100% rename from crates/tabby/src/serve/admin/dist/images/authentication/reset-password.jpg rename to crates/tabby-admin/dist/images/authentication/reset-password.jpg diff --git a/crates/tabby/src/serve/admin/dist/images/logo.svg b/crates/tabby-admin/dist/images/logo.svg similarity index 100% rename from crates/tabby/src/serve/admin/dist/images/logo.svg rename to crates/tabby-admin/dist/images/logo.svg diff --git a/crates/tabby/src/serve/admin/dist/images/products/apple-imac-1.png b/crates/tabby-admin/dist/images/products/apple-imac-1.png similarity index 100% rename from crates/tabby/src/serve/admin/dist/images/products/apple-imac-1.png rename to crates/tabby-admin/dist/images/products/apple-imac-1.png diff --git a/crates/tabby/src/serve/admin/dist/images/products/apple-imac-2.png b/crates/tabby-admin/dist/images/products/apple-imac-2.png similarity index 100% rename from crates/tabby/src/serve/admin/dist/images/products/apple-imac-2.png rename to crates/tabby-admin/dist/images/products/apple-imac-2.png diff --git a/crates/tabby/src/serve/admin/dist/images/products/apple-imac-3.png b/crates/tabby-admin/dist/images/products/apple-imac-3.png similarity index 100% rename from crates/tabby/src/serve/admin/dist/images/products/apple-imac-3.png rename to crates/tabby-admin/dist/images/products/apple-imac-3.png diff --git a/crates/tabby/src/serve/admin/dist/images/users/bonnie-green-2x.png b/crates/tabby-admin/dist/images/users/bonnie-green-2x.png similarity index 100% rename from crates/tabby/src/serve/admin/dist/images/users/bonnie-green-2x.png rename to crates/tabby-admin/dist/images/users/bonnie-green-2x.png diff --git a/crates/tabby/src/serve/admin/dist/images/users/bonnie-green.png b/crates/tabby-admin/dist/images/users/bonnie-green.png similarity index 100% rename from crates/tabby/src/serve/admin/dist/images/users/bonnie-green.png rename to crates/tabby-admin/dist/images/users/bonnie-green.png diff --git a/crates/tabby/src/serve/admin/dist/images/users/helene-engels.png b/crates/tabby-admin/dist/images/users/helene-engels.png similarity index 100% rename from crates/tabby/src/serve/admin/dist/images/users/helene-engels.png rename to crates/tabby-admin/dist/images/users/helene-engels.png diff --git a/crates/tabby/src/serve/admin/dist/images/users/jese-leos-2x.png b/crates/tabby-admin/dist/images/users/jese-leos-2x.png similarity index 100% rename from crates/tabby/src/serve/admin/dist/images/users/jese-leos-2x.png rename to crates/tabby-admin/dist/images/users/jese-leos-2x.png diff --git a/crates/tabby/src/serve/admin/dist/images/users/jese-leos.png b/crates/tabby-admin/dist/images/users/jese-leos.png similarity index 100% rename from crates/tabby/src/serve/admin/dist/images/users/jese-leos.png rename to crates/tabby-admin/dist/images/users/jese-leos.png diff --git a/crates/tabby/src/serve/admin/dist/images/users/joseph-mcfall.png b/crates/tabby-admin/dist/images/users/joseph-mcfall.png similarity index 100% rename from crates/tabby/src/serve/admin/dist/images/users/joseph-mcfall.png rename to crates/tabby-admin/dist/images/users/joseph-mcfall.png diff --git a/crates/tabby/src/serve/admin/dist/images/users/lana-byrd.png b/crates/tabby-admin/dist/images/users/lana-byrd.png similarity index 100% rename from crates/tabby/src/serve/admin/dist/images/users/lana-byrd.png rename to crates/tabby-admin/dist/images/users/lana-byrd.png diff --git a/crates/tabby/src/serve/admin/dist/images/users/leslie-livingston.png b/crates/tabby-admin/dist/images/users/leslie-livingston.png similarity index 100% rename from crates/tabby/src/serve/admin/dist/images/users/leslie-livingston.png rename to crates/tabby-admin/dist/images/users/leslie-livingston.png diff --git a/crates/tabby/src/serve/admin/dist/images/users/michael-gough.png b/crates/tabby-admin/dist/images/users/michael-gough.png similarity index 100% rename from crates/tabby/src/serve/admin/dist/images/users/michael-gough.png rename to crates/tabby-admin/dist/images/users/michael-gough.png diff --git a/crates/tabby/src/serve/admin/dist/images/users/neil-sims.png b/crates/tabby-admin/dist/images/users/neil-sims.png similarity index 100% rename from crates/tabby/src/serve/admin/dist/images/users/neil-sims.png rename to crates/tabby-admin/dist/images/users/neil-sims.png diff --git a/crates/tabby/src/serve/admin/dist/images/users/robert-brown.png b/crates/tabby-admin/dist/images/users/robert-brown.png similarity index 100% rename from crates/tabby/src/serve/admin/dist/images/users/robert-brown.png rename to crates/tabby-admin/dist/images/users/robert-brown.png diff --git a/crates/tabby/src/serve/admin/dist/images/users/roberta-casas-2x.png b/crates/tabby-admin/dist/images/users/roberta-casas-2x.png similarity index 100% rename from crates/tabby/src/serve/admin/dist/images/users/roberta-casas-2x.png rename to crates/tabby-admin/dist/images/users/roberta-casas-2x.png diff --git a/crates/tabby/src/serve/admin/dist/images/users/roberta-casas.png b/crates/tabby-admin/dist/images/users/roberta-casas.png similarity index 100% rename from crates/tabby/src/serve/admin/dist/images/users/roberta-casas.png rename to crates/tabby-admin/dist/images/users/roberta-casas.png diff --git a/crates/tabby/src/serve/admin/dist/images/users/thomas-lean.png b/crates/tabby-admin/dist/images/users/thomas-lean.png similarity index 100% rename from crates/tabby/src/serve/admin/dist/images/users/thomas-lean.png rename to crates/tabby-admin/dist/images/users/thomas-lean.png diff --git a/crates/tabby/src/serve/admin/dist/index.html b/crates/tabby-admin/dist/index.html similarity index 100% rename from crates/tabby/src/serve/admin/dist/index.html rename to crates/tabby-admin/dist/index.html diff --git a/crates/tabby/src/serve/admin/index.html b/crates/tabby-admin/index.html similarity index 100% rename from crates/tabby/src/serve/admin/index.html rename to crates/tabby-admin/index.html diff --git a/crates/tabby/src/serve/admin/package.json b/crates/tabby-admin/package.json similarity index 100% rename from crates/tabby/src/serve/admin/package.json rename to crates/tabby-admin/package.json diff --git a/crates/tabby/src/serve/admin/postcss.config.mjs b/crates/tabby-admin/postcss.config.mjs similarity index 100% rename from crates/tabby/src/serve/admin/postcss.config.mjs rename to crates/tabby-admin/postcss.config.mjs diff --git a/crates/tabby/src/serve/admin/public/images/authentication/create-account.jpg b/crates/tabby-admin/public/images/authentication/create-account.jpg similarity index 100% rename from crates/tabby/src/serve/admin/public/images/authentication/create-account.jpg rename to crates/tabby-admin/public/images/authentication/create-account.jpg diff --git a/crates/tabby/src/serve/admin/public/images/authentication/login.jpg b/crates/tabby-admin/public/images/authentication/login.jpg similarity index 100% rename from crates/tabby/src/serve/admin/public/images/authentication/login.jpg rename to crates/tabby-admin/public/images/authentication/login.jpg diff --git a/crates/tabby/src/serve/admin/public/images/authentication/reset-password.jpg b/crates/tabby-admin/public/images/authentication/reset-password.jpg similarity index 100% rename from crates/tabby/src/serve/admin/public/images/authentication/reset-password.jpg rename to crates/tabby-admin/public/images/authentication/reset-password.jpg diff --git a/crates/tabby/src/serve/admin/public/images/logo.svg b/crates/tabby-admin/public/images/logo.svg similarity index 100% rename from crates/tabby/src/serve/admin/public/images/logo.svg rename to crates/tabby-admin/public/images/logo.svg diff --git a/crates/tabby/src/serve/admin/public/images/products/apple-imac-1.png b/crates/tabby-admin/public/images/products/apple-imac-1.png similarity index 100% rename from crates/tabby/src/serve/admin/public/images/products/apple-imac-1.png rename to crates/tabby-admin/public/images/products/apple-imac-1.png diff --git a/crates/tabby/src/serve/admin/public/images/products/apple-imac-2.png b/crates/tabby-admin/public/images/products/apple-imac-2.png similarity index 100% rename from crates/tabby/src/serve/admin/public/images/products/apple-imac-2.png rename to crates/tabby-admin/public/images/products/apple-imac-2.png diff --git a/crates/tabby/src/serve/admin/public/images/products/apple-imac-3.png b/crates/tabby-admin/public/images/products/apple-imac-3.png similarity index 100% rename from crates/tabby/src/serve/admin/public/images/products/apple-imac-3.png rename to crates/tabby-admin/public/images/products/apple-imac-3.png diff --git a/crates/tabby/src/serve/admin/public/images/users/bonnie-green-2x.png b/crates/tabby-admin/public/images/users/bonnie-green-2x.png similarity index 100% rename from crates/tabby/src/serve/admin/public/images/users/bonnie-green-2x.png rename to crates/tabby-admin/public/images/users/bonnie-green-2x.png diff --git a/crates/tabby/src/serve/admin/public/images/users/bonnie-green.png b/crates/tabby-admin/public/images/users/bonnie-green.png similarity index 100% rename from crates/tabby/src/serve/admin/public/images/users/bonnie-green.png rename to crates/tabby-admin/public/images/users/bonnie-green.png diff --git a/crates/tabby/src/serve/admin/public/images/users/helene-engels.png b/crates/tabby-admin/public/images/users/helene-engels.png similarity index 100% rename from crates/tabby/src/serve/admin/public/images/users/helene-engels.png rename to crates/tabby-admin/public/images/users/helene-engels.png diff --git a/crates/tabby/src/serve/admin/public/images/users/jese-leos-2x.png b/crates/tabby-admin/public/images/users/jese-leos-2x.png similarity index 100% rename from crates/tabby/src/serve/admin/public/images/users/jese-leos-2x.png rename to crates/tabby-admin/public/images/users/jese-leos-2x.png diff --git a/crates/tabby/src/serve/admin/public/images/users/jese-leos.png b/crates/tabby-admin/public/images/users/jese-leos.png similarity index 100% rename from crates/tabby/src/serve/admin/public/images/users/jese-leos.png rename to crates/tabby-admin/public/images/users/jese-leos.png diff --git a/crates/tabby/src/serve/admin/public/images/users/joseph-mcfall.png b/crates/tabby-admin/public/images/users/joseph-mcfall.png similarity index 100% rename from crates/tabby/src/serve/admin/public/images/users/joseph-mcfall.png rename to crates/tabby-admin/public/images/users/joseph-mcfall.png diff --git a/crates/tabby/src/serve/admin/public/images/users/lana-byrd.png b/crates/tabby-admin/public/images/users/lana-byrd.png similarity index 100% rename from crates/tabby/src/serve/admin/public/images/users/lana-byrd.png rename to crates/tabby-admin/public/images/users/lana-byrd.png diff --git a/crates/tabby/src/serve/admin/public/images/users/leslie-livingston.png b/crates/tabby-admin/public/images/users/leslie-livingston.png similarity index 100% rename from crates/tabby/src/serve/admin/public/images/users/leslie-livingston.png rename to crates/tabby-admin/public/images/users/leslie-livingston.png diff --git a/crates/tabby/src/serve/admin/public/images/users/michael-gough.png b/crates/tabby-admin/public/images/users/michael-gough.png similarity index 100% rename from crates/tabby/src/serve/admin/public/images/users/michael-gough.png rename to crates/tabby-admin/public/images/users/michael-gough.png diff --git a/crates/tabby/src/serve/admin/public/images/users/neil-sims.png b/crates/tabby-admin/public/images/users/neil-sims.png similarity index 100% rename from crates/tabby/src/serve/admin/public/images/users/neil-sims.png rename to crates/tabby-admin/public/images/users/neil-sims.png diff --git a/crates/tabby/src/serve/admin/public/images/users/robert-brown.png b/crates/tabby-admin/public/images/users/robert-brown.png similarity index 100% rename from crates/tabby/src/serve/admin/public/images/users/robert-brown.png rename to crates/tabby-admin/public/images/users/robert-brown.png diff --git a/crates/tabby/src/serve/admin/public/images/users/roberta-casas-2x.png b/crates/tabby-admin/public/images/users/roberta-casas-2x.png similarity index 100% rename from crates/tabby/src/serve/admin/public/images/users/roberta-casas-2x.png rename to crates/tabby-admin/public/images/users/roberta-casas-2x.png diff --git a/crates/tabby/src/serve/admin/public/images/users/roberta-casas.png b/crates/tabby-admin/public/images/users/roberta-casas.png similarity index 100% rename from crates/tabby/src/serve/admin/public/images/users/roberta-casas.png rename to crates/tabby-admin/public/images/users/roberta-casas.png diff --git a/crates/tabby/src/serve/admin/public/images/users/thomas-lean.png b/crates/tabby-admin/public/images/users/thomas-lean.png similarity index 100% rename from crates/tabby/src/serve/admin/public/images/users/thomas-lean.png rename to crates/tabby-admin/public/images/users/thomas-lean.png diff --git a/crates/tabby/src/serve/admin/src/components/navbar.tsx b/crates/tabby-admin/src/components/navbar.tsx similarity index 100% rename from crates/tabby/src/serve/admin/src/components/navbar.tsx rename to crates/tabby-admin/src/components/navbar.tsx diff --git a/crates/tabby/src/serve/admin/src/components/sidebar.tsx b/crates/tabby-admin/src/components/sidebar.tsx similarity index 100% rename from crates/tabby/src/serve/admin/src/components/sidebar.tsx rename to crates/tabby-admin/src/components/sidebar.tsx diff --git a/crates/tabby/src/serve/admin/src/flowbite-theme.ts b/crates/tabby-admin/src/flowbite-theme.ts similarity index 100% rename from crates/tabby/src/serve/admin/src/flowbite-theme.ts rename to crates/tabby-admin/src/flowbite-theme.ts diff --git a/crates/tabby/src/serve/admin/src/index.css b/crates/tabby-admin/src/index.css similarity index 100% rename from crates/tabby/src/serve/admin/src/index.css rename to crates/tabby-admin/src/index.css diff --git a/crates/tabby/src/serve/admin/src/index.tsx b/crates/tabby-admin/src/index.tsx similarity index 100% rename from crates/tabby/src/serve/admin/src/index.tsx rename to crates/tabby-admin/src/index.tsx diff --git a/crates/tabby/src/serve/admin/src/layouts/navbar-sidebar.tsx b/crates/tabby-admin/src/layouts/navbar-sidebar.tsx similarity index 100% rename from crates/tabby/src/serve/admin/src/layouts/navbar-sidebar.tsx rename to crates/tabby-admin/src/layouts/navbar-sidebar.tsx diff --git a/crates/tabby/src/serve/admin/src/pages/authentication/sign-in.tsx b/crates/tabby-admin/src/pages/authentication/sign-in.tsx similarity index 100% rename from crates/tabby/src/serve/admin/src/pages/authentication/sign-in.tsx rename to crates/tabby-admin/src/pages/authentication/sign-in.tsx diff --git a/crates/tabby/src/serve/admin/src/pages/authentication/sign-up.tsx b/crates/tabby-admin/src/pages/authentication/sign-up.tsx similarity index 100% rename from crates/tabby/src/serve/admin/src/pages/authentication/sign-up.tsx rename to crates/tabby-admin/src/pages/authentication/sign-up.tsx diff --git a/crates/tabby/src/serve/admin/src/pages/e-commerce/products.tsx b/crates/tabby-admin/src/pages/e-commerce/products.tsx similarity index 100% rename from crates/tabby/src/serve/admin/src/pages/e-commerce/products.tsx rename to crates/tabby-admin/src/pages/e-commerce/products.tsx diff --git a/crates/tabby/src/serve/admin/src/pages/index.tsx b/crates/tabby-admin/src/pages/index.tsx similarity index 100% rename from crates/tabby/src/serve/admin/src/pages/index.tsx rename to crates/tabby-admin/src/pages/index.tsx diff --git a/crates/tabby/src/serve/admin/src/pages/users/list.tsx b/crates/tabby-admin/src/pages/users/list.tsx similarity index 100% rename from crates/tabby/src/serve/admin/src/pages/users/list.tsx rename to crates/tabby-admin/src/pages/users/list.tsx diff --git a/crates/tabby/src/serve/admin/src/svgmap.d.ts b/crates/tabby-admin/src/svgmap.d.ts similarity index 100% rename from crates/tabby/src/serve/admin/src/svgmap.d.ts rename to crates/tabby-admin/src/svgmap.d.ts diff --git a/crates/tabby/src/serve/admin/tailwind.config.cjs b/crates/tabby-admin/tailwind.config.cjs similarity index 100% rename from crates/tabby/src/serve/admin/tailwind.config.cjs rename to crates/tabby-admin/tailwind.config.cjs diff --git a/crates/tabby/src/serve/admin/tsconfig.json b/crates/tabby-admin/tsconfig.json similarity index 100% rename from crates/tabby/src/serve/admin/tsconfig.json rename to crates/tabby-admin/tsconfig.json diff --git a/crates/tabby/src/serve/admin/vite.config.ts b/crates/tabby-admin/vite.config.ts similarity index 100% rename from crates/tabby/src/serve/admin/vite.config.ts rename to crates/tabby-admin/vite.config.ts diff --git a/crates/tabby/src/serve/admin/yarn.lock b/crates/tabby-admin/yarn.lock similarity index 100% rename from crates/tabby/src/serve/admin/yarn.lock rename to crates/tabby-admin/yarn.lock diff --git a/crates/tabby/Cargo.toml b/crates/tabby/Cargo.toml index 11723bb..7ff8399 100644 --- a/crates/tabby/Cargo.toml +++ b/crates/tabby/Cargo.toml @@ -21,6 +21,11 @@ regex = "1.8.3" lazy_static = "1.4.0" rust-embed = "6.6.1" mime_guess = "2.0.4" +strum = { version = "0.24", features = ["derive"] } +reqwest = { version = "0.11.18", features = ["stream"] } +error-chain = "0.12.4" +indicatif = "0.17.3" +futures-util = "0.3.28" [dependencies.uuid] version = "1.3.3" diff --git a/crates/tabby/src/download.rs b/crates/tabby/src/download.rs new file mode 100644 index 0000000..7b8ec6f --- /dev/null +++ b/crates/tabby/src/download.rs @@ -0,0 +1,99 @@ +use std::cmp; +use std::fs; +use std::io::Write; +use std::path::{Path, PathBuf}; + +use clap::Args; +use error_chain::error_chain; +use futures_util::StreamExt; +use indicatif::{ProgressBar, ProgressStyle}; + +#[derive(Args)] +pub struct DownloadArgs { + /// model id to fetch. + #[clap(long)] + model: String, +} + +error_chain! { + foreign_links { + Io(std::io::Error); + HttpRequest(reqwest::Error); + TemplateError(indicatif::style::TemplateError); + } +} + +pub async fn main(args: &DownloadArgs) -> Result<()> { + download_model(&args.model).await.unwrap(); + Ok(()) +} + +async fn download_model(model_id: &str) -> Result<()> { + download_metadata(model_id).await?; + download_model_file(model_id, "tokenizer.json").await?; + download_model_file(model_id, &format!("ctranslate2/config.json")).await?; + download_model_file(model_id, &format!("ctranslate2/vocabulary.txt")).await?; + download_model_file(model_id, &format!("ctranslate2/shared_vocabulary.txt")).await?; + download_model_file(model_id, &format!("ctranslate2/model.bin")).await?; + Ok(()) +} + +fn get_model_dir(model_id: &str) -> PathBuf { + let home = std::env::var("HOME").unwrap(); + let tabby_root = format!("{}/.tabby", home); + let model_dir = Path::new(&tabby_root).join("models").join(model_id); + model_dir +} + +async fn download_metadata(model_id: &str) -> Result<()> { + let url = format!("https://huggingface.co/api/models/{}", model_id); + let fname = "metadata.json"; + let filepath = get_model_dir(model_id).join(fname).display().to_string(); + download_file(&format!("{}/{}", model_id, fname), &url, &filepath).await +} + +async fn download_model_file(model_id: &str, fname: &str) -> Result<()> { + // Create url. + let url = format!("https://huggingface.co/{}/resolve/main/{}", model_id, fname); + + // Create destination path. + let filepath = get_model_dir(model_id).join(fname).display().to_string(); + download_file(&format!("{}/{}", model_id, fname), &url, &filepath).await +} + +async fn download_file(name: &str, url: &str, path: &str) -> Result<()> { + fs::create_dir_all(Path::new(path).parent().unwrap())?; + + // Reqwest setup + let res = reqwest::get(url) + .await + .or(Err(format!("Failed to GET from '{}'", url)))?; + + let total_size = res + .content_length() + .ok_or(format!("Failed to get content length from '{}'", url))?; + + // Indicatif setup + let pb = ProgressBar::new(total_size); + pb.set_style(ProgressStyle::default_bar() + .template("{msg}\n{spinner:.green} [{elapsed_precise}] [{wide_bar:.cyan/blue}] {bytes}/{total_bytes} ({bytes_per_sec}, {eta})")? + .progress_chars("#>-")); + pb.set_message(format!("Downloading {}", &name)); + + // download chunks + let mut file = fs::File::create(&path).or(Err(format!("Failed to create file '{}'", &path)))?; + let mut downloaded: u64 = 0; + let mut stream = res.bytes_stream(); + + while let Some(item) = stream.next().await { + let chunk = item.or(Err(format!("Error while downloading file")))?; + file.write_all(&chunk) + .or(Err(format!("Error while writing to file")))?; + let new = cmp::min(downloaded + (chunk.len() as u64), total_size); + downloaded = new; + pb.set_position(new); + } + + pb.finish_with_message(format!("Downloaded {}", &name)); + return Ok(()); +} diff --git a/crates/tabby/src/main.rs b/crates/tabby/src/main.rs index 7e3605e..d7d9d80 100644 --- a/crates/tabby/src/main.rs +++ b/crates/tabby/src/main.rs @@ -1,3 +1,6 @@ +mod download; +mod serve; + use clap::{Parser, Subcommand}; #[derive(Parser)] @@ -12,21 +15,25 @@ struct Cli { pub enum Commands { /// Serve the model Serve(serve::ServeArgs), -} -mod serve; + /// Download the model + Download(download::DownloadArgs), +} #[tokio::main] async fn main() { let cli = Cli::parse(); - // You can check for the existence of subcommands, and if found use their - // matches just as you would the top level cmd match &cli.command { Commands::Serve(args) => { serve::main(args) .await .expect("Error happens during the serve"); } + Commands::Download(args) => { + download::main(args) + .await + .expect("Error happens during the download"); + } } } diff --git a/crates/tabby/src/serve/admin.rs b/crates/tabby/src/serve/admin.rs index 791b6e8..c98faa5 100644 --- a/crates/tabby/src/serve/admin.rs +++ b/crates/tabby/src/serve/admin.rs @@ -5,7 +5,7 @@ use axum::{ }; #[derive(rust_embed::RustEmbed)] -#[folder = "src/serve/admin/dist/"] +#[folder = "../tabby-admin/dist/"] struct AdminAssets; struct AdminStaticFile(pub T); diff --git a/crates/tabby/src/serve/completions.rs b/crates/tabby/src/serve/completions.rs index 99d9a24..3f973d9 100644 --- a/crates/tabby/src/serve/completions.rs +++ b/crates/tabby/src/serve/completions.rs @@ -1,8 +1,9 @@ use axum::{extract::State, Json}; use ctranslate2_bindings::{ - TextInferenceEngine, TextInferenceEngineCreateOptions, TextInferenceOptionsBuilder, + TextInferenceEngine, TextInferenceEngineCreateOptionsBuilder, TextInferenceOptionsBuilder, }; use serde::{Deserialize, Serialize}; +use std::path::Path; use std::sync::Arc; use utoipa::ToSchema; @@ -64,7 +65,22 @@ pub struct CompletionState { } impl CompletionState { - pub fn new(options: TextInferenceEngineCreateOptions) -> Self { + pub fn new(args: &crate::serve::ServeArgs) -> Self { + let home = std::env::var("HOME").unwrap(); + let tabby_root = format!("{}/.tabby", home); + let model_dir = Path::new(&tabby_root).join("models").join(&args.model); + let metadata = read_metadata(&model_dir); + + let device = format!("{}", args.device); + let options = TextInferenceEngineCreateOptionsBuilder::default() + .model_path(model_dir.join("ctranslate2").display().to_string()) + .tokenizer_path(model_dir.join("tokenizer.json").display().to_string()) + .device(device) + .model_type(metadata.transformers_info.auto_model) + .device_indices(args.device_indices.clone()) + .num_replicas_per_device(args.num_replicas_per_device) + .build() + .unwrap(); let engine = TextInferenceEngine::create(options); Self { engine } } @@ -78,3 +94,20 @@ fn timestamp() -> u64 { .expect("Time went backwards") .as_secs() } + +#[derive(Deserialize)] +#[serde(rename_all = "camelCase")] +struct Metadata { + transformers_info: TransformersInfo, +} + +#[derive(Deserialize)] +struct TransformersInfo { + auto_model: String, +} + +fn read_metadata(model_dir: &std::path::PathBuf) -> Metadata { + let file = std::fs::File::open(model_dir.join("metadata.json")).unwrap(); + let reader = std::io::BufReader::new(file); + serde_json::from_reader(reader).unwrap() +} diff --git a/crates/tabby/src/serve/mod.rs b/crates/tabby/src/serve/mod.rs index e960c06..305bfe2 100644 --- a/crates/tabby/src/serve/mod.rs +++ b/crates/tabby/src/serve/mod.rs @@ -1,20 +1,19 @@ +mod admin; +mod completions; +mod events; + +use crate::Cli; +use axum::{routing, Router, Server}; +use clap::{error::ErrorKind, Args, CommandFactory}; +use hyper::Error; use std::{ net::{Ipv4Addr, SocketAddr}, sync::Arc, }; - -use axum::{routing, Router, Server}; -use clap::Args; -use ctranslate2_bindings::TextInferenceEngineCreateOptionsBuilder; -use hyper::Error; -use std::path::Path; use tower_http::cors::CorsLayer; use utoipa::OpenApi; use utoipa_swagger_ui::SwaggerUi; -mod completions; -mod events; - #[derive(OpenApi)] #[openapi( paths(events::log_event, completions::completion,), @@ -27,70 +26,47 @@ mod events; )] struct ApiDoc; -#[derive(clap::ValueEnum, Clone)] +#[derive(clap::ValueEnum, strum::Display, PartialEq, Clone)] pub enum Device { + #[strum(serialize = "cpu")] CPU, + + #[strum(serialize = "cuda")] CUDA, } -impl std::fmt::Display for Device { - fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { - let printable = match *self { - Device::CPU => "cpu", - Device::CUDA => "cuda", - }; - write!(f, "{}", printable) - } -} - -#[derive(clap::ValueEnum, Clone)] -pub enum ModelType { - EncoderDecoder, - Decoder, -} - -impl std::fmt::Display for ModelType { - fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { - let printable = match *self { - ModelType::EncoderDecoder => "encoder-decoder", - ModelType::Decoder => "decoder", - }; - write!(f, "{}", printable) - } -} - #[derive(Args)] pub struct ServeArgs { - /// path to model for serving + /// Model id for serving. #[clap(long)] model: String, - /// model type for serving - #[clap(long, default_value_t=ModelType::Decoder)] - model_type: ModelType, - #[clap(long, default_value_t = 8080)] port: u16, + /// Device to run model inference. #[clap(long, default_value_t=Device::CPU)] device: Device, + /// GPU indices to run models, only applicable for CUDA. #[clap(long, default_values_t=[0])] device_indices: Vec, - /// num_replicas_per_device + /// Number of replicas per device, only applicable for CPU. #[clap(long, default_value_t = 1)] num_replicas_per_device: usize, + /// *INTERNAL ONLY* #[clap(long, default_value_t = false)] experimental_admin_panel: bool, } pub async fn main(args: &ServeArgs) -> Result<(), Error> { + valid_args(args); let app = Router::new() .merge(SwaggerUi::new("/swagger-ui").url("/api-docs/openapi.json", ApiDoc::openapi())) .nest("/v1", api_router(args)) - .fallback(fallback(args)); + .fallback(fallback(args.experimental_admin_panel)); let address = SocketAddr::from((Ipv4Addr::UNSPECIFIED, args.port)); println!("Listening at {}", address); @@ -100,41 +76,39 @@ pub async fn main(args: &ServeArgs) -> Result<(), Error> { fn api_router(args: &ServeArgs) -> Router { Router::new() .route("/events", routing::post(events::log_event)) - .route("/completions", routing::post(completions::completion)) - .with_state(Arc::new(new_completion_state(args))) + .route( + "/completions", + routing::post(completions::completion) + .with_state(Arc::new(completions::CompletionState::new(args))), + ) .layer(CorsLayer::permissive()) } -mod admin; -fn fallback(args: &ServeArgs) -> routing::MethodRouter { - if args.experimental_admin_panel { +fn fallback(experimental_admin_panel: bool) -> routing::MethodRouter { + if experimental_admin_panel { routing::get(admin::handler) } else { routing::get(|| async { axum::response::Redirect::temporary("/swagger-ui") }) } } -fn new_completion_state(args: &ServeArgs) -> completions::CompletionState { - let device = format!("{}", args.device); - let options = TextInferenceEngineCreateOptionsBuilder::default() - .model_path( - Path::new(&args.model) - .join("ctranslate2") - .join(device.clone()) - .display() - .to_string(), - ) - .tokenizer_path( - Path::new(&args.model) - .join("tokenizer.json") - .display() - .to_string(), - ) - .device(device) - .model_type(format!("{}", args.model_type)) - .device_indices(args.device_indices.clone()) - .num_replicas_per_device(args.num_replicas_per_device) - .build() - .unwrap(); - completions::CompletionState::new(options) +fn valid_args(args: &ServeArgs) { + if args.device == Device::CUDA && args.num_replicas_per_device != 1 { + Cli::command() + .error( + ErrorKind::ValueValidation, + "CUDA device only supports 1 replicas per device", + ) + .exit(); + } + + if args.device == Device::CPU && (args.device_indices.len() != 1 || args.device_indices[0] != 0) + { + Cli::command() + .error( + ErrorKind::ValueValidation, + "CPU device only supports device indices = [0]", + ) + .exit(); + } }