[wip] inference service

Cargo.toml

@@ -12,9 +12,17 @@ candle = { git = "https://github.com/huggingface/candle", package = "candle-core
 candle-flash-attn = { git = "https://github.com/huggingface/candle", package = "candle-flash-attn", version = "0.4.2" }
 candle-nn = { git = "https://github.com/huggingface/candle", package = "candle-nn", version = "0.4.2" }
 candle-transformers = { git = "https://github.com/huggingface/candle", package = "candle-transformers", version = "0.4.2" }
+config = "0.14.0"
 ed25519-consensus = "2.1.0"
+futures = "0.3.30"
+hf-hub = "0.3.2"
 serde = "1.0.197"
+serde_json = "1.0.114"
+rand = "0.8.5"
+reqwest = "0.12.1"
 thiserror = "1.0.58"
 tokenizers = "0.15.2"
 tokio = "1.36.0"
+toml = "0.8.12"
 tracing = "0.1.40"
+tracing-subscriber = "0.3.18"

atoma-inference/Cargo.toml

@@ -11,12 +11,23 @@ candle.workspace = true
 candle-flash-attn = { workspace = true, optional = true }
 candle-nn.workspace = true
 candle-transformers.workspace = true
+config.true = true
 ed25519-consensus.workspace = true
+futures.workspace = true
+hf-hub.workspace = true
+reqwest = { workspace = true, features = ["json"] }
 serde = { workspace = true, features = ["derive"] }
+serde_json.workspace = true
 thiserror.workspace = true
 tokenizers.workspace = true
 tokio = { workspace = true, features = ["full", "tracing"] }
 tracing.workspace = true
+tracing-subscriber.workspace = true
+
+[dev-dependencies]
+rand.workspace = true
+toml.workspace = true
+
 
 [features]
 accelerate = ["candle/accelerate", "candle-nn/accelerate", "candle-transformers/accelerate"]

atoma-inference/src/apis/hugging_face.rs

@@ -0,0 +1,117 @@
+use std::path::PathBuf;
+
+use async_trait::async_trait;
+use hf_hub::api::sync::{Api, ApiBuilder};
+
+use crate::models::ModelId;
+
+use super::{ApiError, ApiTrait};
+
+struct FilePaths {
+    file_paths: Vec<String>,
+}
+
+fn get_model_safe_tensors_from_hf(model_id: &ModelId) -> (String, FilePaths) {
+    match model_id.as_str() {
+        "Llama2_7b" => (
+            String::from("meta-llama/Llama-2-7b-hf"),
+            FilePaths {
+                file_paths: vec![
+                    "model-00001-of-00002.safetensors".to_string(),
+                    "model-00002-of-00002.safetensors".to_string(),
+                ],
+            },
+        ),
+        "Mamba3b" => (
+            String::from("state-spaces/mamba-2.8b-hf"),
+            FilePaths {
+                file_paths: vec![
+                    "model-00001-of-00003.safetensors".to_string(),
+                    "model-00002-of-00003.safetensors".to_string(),
+                    "model-00003-of-00003.safetensors".to_string(),
+                ],
+            },
+        ),
+        "Mistral7b" => (
+            String::from("mistralai/Mistral-7B-Instruct-v0.2"),
+            FilePaths {
+                file_paths: vec![
+                    "model-00001-of-00003.safetensors".to_string(),
+                    "model-00002-of-00003.safetensors".to_string(),
+                    "model-00003-of-00003.safetensors".to_string(),
+                ],
+            },
+        ),
+        "Mixtral8x7b" => (
+            String::from("mistralai/Mixtral-8x7B-Instruct-v0.1"),
+            FilePaths {
+                file_paths: vec![
+                    "model-00001-of-00019.safetensors".to_string(),
+                    "model-00002-of-00019.safetensors".to_string(),
+                    "model-00003-of-00019.safetensors".to_string(),
+                    "model-00004-of-00019.safetensors".to_string(),
+                    "model-00005-of-00019.safetensors".to_string(),
+                    "model-00006-of-00019.safetensors".to_string(),
+                    "model-00007-of-00019.safetensors".to_string(),
+                    "model-00008-of-00019.safetensors".to_string(),
+                    "model-00009-of-00019.safetensors".to_string(),
+                    "model-000010-of-00019.safetensors".to_string(),
+                    "model-000011-of-00019.safetensors".to_string(),
+                    "model-000012-of-00019.safetensors".to_string(),
+                    "model-000013-of-00019.safetensors".to_string(),
+                    "model-000014-of-00019.safetensors".to_string(),
+                    "model-000015-of-00019.safetensors".to_string(),
+                    "model-000016-of-00019.safetensors".to_string(),
+                    "model-000017-of-00019.safetensors".to_string(),
+                    "model-000018-of-00019.safetensors".to_string(),
+                    "model-000019-of-00019.safetensors".to_string(),
+                ],
+            },
+        ),
+        "StableDiffusion2" => (
+            String::from("stabilityai/stable-diffusion-2"),
+            FilePaths {
+                file_paths: vec!["768-v-ema.safetensors".to_string()],
+            },
+        ),
+        "StableDiffusionXl" => (
+            String::from("stabilityai/stable-diffusion-xl-base-1.0"),
+            FilePaths {
+                file_paths: vec![
+                    "sd_xl_base_1.0.safetensors".to_string(),
+                    "sd_xl_base_1.0_0.9vae.safetensors".to_string(),
+                    "sd_xl_offset_example-lora_1.0.safetensors".to_string(),
+                ],
+            },
+        ),
+        _ => {
+            panic!("Invalid model id")
+        }
+    }
+}
+
+#[async_trait]
+impl ApiTrait for Api {
+    fn create(api_key: String, cache_dir: PathBuf) -> Result<Self, ApiError>
+    where
+        Self: Sized,
+    {
+        Ok(ApiBuilder::new()
+            .with_progress(true)
+            .with_token(Some(api_key))
+            .with_cache_dir(cache_dir)
+            .build()?)
+    }
+
+    fn fetch(&self, model_id: &ModelId) -> Result<Vec<PathBuf>, ApiError> {
+        let (model_path, files) = get_model_safe_tensors_from_hf(model_id);
+        let api_repo = self.model(model_path);
+        let mut path_bufs = Vec::with_capacity(files.file_paths.len());
+
+        for file in files.file_paths {
+            path_bufs.push(api_repo.get(&file)?);
+        }
+
+        Ok(path_bufs)
+    }
+}

atoma-inference/src/apis/mod.rs

@@ -0,0 +1,29 @@
+pub mod hugging_face;
+use hf_hub::api::sync::ApiError as HuggingFaceError;
+
+use std::path::PathBuf;
+
+use thiserror::Error;
+
+use crate::models::ModelId;
+
+#[derive(Debug, Error)]
+pub enum ApiError {
+    #[error("Api Error: `{0}`")]
+    ApiError(String),
+    #[error("HuggingFace API error: `{0}`")]
+    HuggingFaceError(HuggingFaceError),
+}
+
+impl From<HuggingFaceError> for ApiError {
+    fn from(error: HuggingFaceError) -> Self {
+        Self::HuggingFaceError(error)
+    }
+}
+
+pub trait ApiTrait: Send {
+    fn fetch(&self, model_id: &ModelId) -> Result<Vec<PathBuf>, ApiError>;
+    fn create(api_key: String, cache_dir: PathBuf) -> Result<Self, ApiError>
+    where
+        Self: Sized;
+}

atoma-inference/src/lib.rs

@@ -1,6 +1,7 @@
-pub mod config;
-pub mod core_thread;
-pub mod models;
+pub mod model_thread;
 pub mod service;
 pub mod specs;
 pub mod types;
+
+pub mod apis;
+pub mod models;

atoma-inference/src/main.rs

@@ -1,3 +1,16 @@
-fn main() {
-    println!("Hello, world!");
+// use hf_hub::api::sync::Api;
+// use inference::service::ModelService;
+
+#[tokio::main]
+async fn main() {
+    tracing_subscriber::fmt::init();
+
+    // let (_, receiver) = tokio::sync::mpsc::channel(32);
+
+    // let _ = ModelService::start::<Model, Api>(
+    //     "../inference.toml".parse().unwrap(),
+    //     "../private_key".parse().unwrap(),
+    //     receiver,
+    // )
+    // .expect("Failed to start inference service");
 }