Skip to main content

ferrum_cli/commands/
pull.rs

1//! Pull command - Download a model from HuggingFace Hub
2
3use crate::config::CliConfig;
4use clap::Args;
5use colored::*;
6use ferrum_models::HfDownloader;
7use ferrum_types::Result;
8use std::path::PathBuf;
9
10#[derive(Args)]
11pub struct PullCommand {
12    /// Model name to download (e.g., tinyllama, qwen2.5:7b, TinyLlama/TinyLlama-1.1B-Chat-v1.0)
13    pub model: String,
14}
15
16pub async fn execute(cmd: PullCommand, config: CliConfig) -> Result<()> {
17    let model_id = resolve_model_alias(&cmd.model);
18
19    println!("{} {}", "Pulling".cyan().bold(), model_id);
20
21    // Get cache directory
22    let cache_dir = get_hf_cache_dir(&config);
23    println!("{}", format!("Cache: {}", cache_dir.display()).dimmed());
24
25    // Get HF token
26    let token = std::env::var("HF_TOKEN")
27        .or_else(|_| std::env::var("HUGGING_FACE_HUB_TOKEN"))
28        .ok();
29
30    // Create downloader with proxy support
31    let downloader = HfDownloader::new(cache_dir, token)?;
32
33    // Download
34    match downloader.download(&model_id, None).await {
35        Ok(path) => {
36            println!();
37            println!("{} Model ready at:", "✓".green().bold());
38            println!("  {}", path.display());
39            Ok(())
40        }
41        Err(e) => {
42            eprintln!();
43            eprintln!("{} Failed to pull model: {}", "✗".red().bold(), e);
44            eprintln!();
45            eprintln!("Tips:");
46            eprintln!("  • Check your internet connection");
47            eprintln!("  • Set proxy: HTTPS_PROXY=socks5h://host:port");
48            eprintln!("  • For private models, set HF_TOKEN environment variable");
49            Err(e)
50        }
51    }
52}
53
54fn resolve_model_alias(name: &str) -> String {
55    match name.to_lowercase().as_str() {
56        "tinyllama" | "tiny" => "TinyLlama/TinyLlama-1.1B-Chat-v1.0".to_string(),
57        "qwen2.5:0.5b" | "qwen:0.5b" => "Qwen/Qwen2.5-0.5B-Instruct".to_string(),
58        "qwen2.5:1.5b" | "qwen:1.5b" => "Qwen/Qwen2.5-1.5B-Instruct".to_string(),
59        "qwen2.5:3b" | "qwen:3b" => "Qwen/Qwen2.5-3B-Instruct".to_string(),
60        "qwen2.5:7b" | "qwen:7b" => "Qwen/Qwen2.5-7B-Instruct".to_string(),
61        "llama3.2:1b" => "meta-llama/Llama-3.2-1B-Instruct".to_string(),
62        "llama3.2:3b" => "meta-llama/Llama-3.2-3B-Instruct".to_string(),
63        _ => name.to_string(),
64    }
65}
66
67fn get_hf_cache_dir(config: &CliConfig) -> PathBuf {
68    if let Ok(hf_home) = std::env::var("HF_HOME") {
69        return PathBuf::from(hf_home);
70    }
71    let configured = shellexpand::tilde(&config.models.download.hf_cache_dir).to_string();
72    PathBuf::from(configured)
73}