From 6656662d795ad1d4fbf2b32d6c094ef6eb2dd88e Mon Sep 17 00:00:00 2001 From: kerthcet Date: Tue, 25 Feb 2025 16:26:52 +0800 Subject: [PATCH] Use .puma as the root folder Signed-off-by: kerthcet --- Cargo.lock | 65 ++++++++++++++++++++++++++++++++++++++-- Cargo.toml | 1 + src/cli/cmds.rs | 7 +++-- src/downloader/ollama.rs | 12 ++++---- src/main.rs | 4 +++ src/util/file.rs | 18 +++++++++++ src/util/mod.rs | 1 + src/util/request.rs | 9 ++++-- 8 files changed, 103 insertions(+), 14 deletions(-) create mode 100644 src/util/file.rs diff --git a/Cargo.lock b/Cargo.lock index c8a0962..522f786 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -238,6 +238,15 @@ dependencies = [ "memchr", ] +[[package]] +name = "dirs" +version = "6.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c3e8aa94d75141228480295a7d0e7feb620b1a5ad9f12bc40be62411e38cce4e" +dependencies = [ + "dirs-sys", +] + [[package]] name = "dirs-next" version = "2.0.0" @@ -248,6 +257,18 @@ dependencies = [ "dirs-sys-next", ] +[[package]] +name = "dirs-sys" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e01a3366d27ee9890022452ee61b2b63a67e6f13f58900b651ff5665f0bb1fab" +dependencies = [ + "libc", + "option-ext", + "redox_users 0.5.0", + "windows-sys 0.59.0", +] + [[package]] name = "dirs-sys-next" version = "0.1.2" @@ -255,7 +276,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4ebda144c4fe02d1f7ea1a7d9641b6fc6b580adcfa024ae48797ecdeb6825b4d" dependencies = [ "libc", - "redox_users", + "redox_users 0.4.6", "winapi", ] @@ -948,6 +969,12 @@ dependencies = [ "vcpkg", ] +[[package]] +name = "option-ext" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "04744f49eae99ab78e0d5c0b603ab218f515ea8cfe5a456d7629ad883a3b6e7d" + [[package]] name = "parking_lot" version = "0.12.3" @@ -1029,6 +1056,7 @@ name = "puma" version = "0.0.1" dependencies = [ "clap", + "dirs", "env_logger", "indicatif", "log", @@ -1065,7 +1093,18 @@ checksum = "ba009ff324d1fc1b900bd1fdb31564febe58a8ccc8a6fdbb93b543d33b13ca43" dependencies = [ "getrandom 0.2.15", "libredox", - "thiserror", + "thiserror 1.0.69", +] + +[[package]] +name = "redox_users" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dd6f9d3d47bdd2ad6945c5015a226ec6155d0bcdfd8f7cd29f86b71f8de99d2b" +dependencies = [ + "getrandom 0.2.15", + "libredox", + "thiserror 2.0.11", ] [[package]] @@ -1448,7 +1487,16 @@ version = "1.0.69" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b6aaf5339b578ea85b50e080feb250a3e8ae8cfcdff9a461c9ec2904bc923f52" dependencies = [ - "thiserror-impl", + "thiserror-impl 1.0.69", +] + +[[package]] +name = "thiserror" +version = "2.0.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d452f284b73e6d76dd36758a0c8684b1d5be31f92b89d07fd5822175732206fc" +dependencies = [ + "thiserror-impl 2.0.11", ] [[package]] @@ -1462,6 +1510,17 @@ dependencies = [ "syn", ] +[[package]] +name = "thiserror-impl" +version = "2.0.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "26afc1baea8a989337eeb52b6e72a039780ce45c3edfcc9c5b9d112feeb173c2" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "tinystr" version = "0.7.6" diff --git a/Cargo.toml b/Cargo.toml index 6afb608..d46c79e 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -16,3 +16,4 @@ serde_derive = "1.0" env_logger = "0.11.6" log = "0.4.26" indicatif = "0.17.11" +dirs = "6.0.0" diff --git a/src/cli/cmds.rs b/src/cli/cmds.rs index e4be907..ac0dafc 100644 --- a/src/cli/cmds.rs +++ b/src/cli/cmds.rs @@ -2,6 +2,7 @@ use clap::{Parser, Subcommand}; use prettytable::{format, row, Table}; use crate::downloader::ollama::OllamaDownloader; +use crate::util::file; #[derive(Parser)] #[command(name = "PUMA")] @@ -105,9 +106,9 @@ pub async fn run(cli: Cli) { } Provider::Ollama => { let d = OllamaDownloader::new(&args.model); - d.download_model("/Users/kerthcet/Workspaces/InftyAI/puma/tmp") - .await - .unwrap(); + let model_path = file::root_home().join(file::model_folder_name(&args.model)); + file::create_folder_if_not_exists(&model_path).unwrap(); + d.download_model(&model_path).await.unwrap(); } Provider::Modelscope => { println!("Downloading model from Modelscope..."); diff --git a/src/downloader/ollama.rs b/src/downloader/ollama.rs index 803ed49..a929bcc 100644 --- a/src/downloader/ollama.rs +++ b/src/downloader/ollama.rs @@ -1,3 +1,4 @@ +use std::path::PathBuf; use std::sync::Arc; use indicatif::{MultiProgress, ProgressStyle}; @@ -50,7 +51,7 @@ impl OllamaDownloader { } } - pub async fn download_model(&self, path: &str) -> Result<(), DownloadError> { + pub async fn download_model(&self, path: &PathBuf) -> Result<(), DownloadError> { info!( "Downloading model {} from ollama provider...", self.model_name @@ -89,7 +90,7 @@ impl OllamaDownloader { let semaphore: Arc = Arc::clone(&semaphore); let arc_m = Arc::clone(&arc_m); - let full_path = format!("{}/{}", path, layer.path()); + let full_path = path.join(layer.path()); let size = layer.size; let sty = sty.clone(); @@ -98,19 +99,18 @@ impl OllamaDownloader { // TODO: return the error. let _ = request::download_file( client, - layer_url.clone(), + layer_url, size, layer.path().to_string(), - full_path, + &full_path, arc_m, sty.clone(), ) .await .map_err(|e| { error!( - "Failed to download file {} from {}: {}", + "Failed to download file {}: {}", layer.path(), - layer_url, e.to_string() ); }); diff --git a/src/main.rs b/src/main.rs index 8670eb3..6547036 100644 --- a/src/main.rs +++ b/src/main.rs @@ -7,10 +7,14 @@ use env_logger; use tokio::runtime::Builder; use crate::cli::cmds::{run, Cli}; +use crate::util::file; fn main() { env_logger::init(); + // create the root folder. + file::create_folder_if_not_exists(&file::root_home()).unwrap(); + let runtime = Builder::new_multi_thread() .worker_threads(4) .enable_all() diff --git a/src/util/file.rs b/src/util/file.rs new file mode 100644 index 0000000..fcf616c --- /dev/null +++ b/src/util/file.rs @@ -0,0 +1,18 @@ +use std::fs; +use std::path::PathBuf; + +use dirs::home_dir; + +pub fn create_folder_if_not_exists(folder_path: &PathBuf) -> std::io::Result<()> { + fs::create_dir_all(folder_path)?; + Ok(()) +} + +pub fn model_folder_name(model_name: &str) -> String { + model_name.replace(":", "--").replace("/", "--") +} + +pub fn root_home() -> PathBuf { + let home = home_dir().expect("Failed to get home directory"); + home.join(".puma") +} diff --git a/src/util/mod.rs b/src/util/mod.rs index be9378d..c756da3 100644 --- a/src/util/mod.rs +++ b/src/util/mod.rs @@ -1 +1,2 @@ +pub mod file; pub mod request; diff --git a/src/util/request.rs b/src/util/request.rs index 442917f..a82369d 100644 --- a/src/util/request.rs +++ b/src/util/request.rs @@ -3,6 +3,7 @@ use std::error::Error; use std::fs::File; use std::io; use std::os::unix::fs::FileExt; +use std::path::PathBuf; use std::sync::Arc; use indicatif::{MultiProgress, ProgressBar, ProgressStyle}; @@ -20,11 +21,15 @@ pub async fn download_file( url: String, content_length: u64, filename: String, - output_path: String, + output_path: &PathBuf, m: Arc, sty: ProgressStyle, ) -> Result<(), Box> { - debug!("Start to download file {} to {}", filename, output_path); + debug!( + "Start to download file {} to {}", + filename, + output_path.display() + ); let mut tasks = Vec::new(); let mut start = 0;