//! model_prefetch: pre-download model files from Hugging Face Hub into the local cache.
//!
//! Why: Avoid slow first runs and enable offline use by warming up the HF cache for one or more
//! models. This helper fetches tokenizer/config and weights (single-file or sharded) so that
//! other tools in this workspace (e.g., model_generation) can start instantly.
//!
//! Quick start
//! - List presets: `cargo run -p model_prefetch -- --list`
//! - Prefetch the default model (SmolLM2-1.7B-Instruct): `cargo run -p model_prefetch --`
//! - Prefetch presets: `cargo run -p model_prefetch -- --which SmoLM2-1.7B-Instruct --which tiny-llama-1.1b-chat`
//! - Prefetch by repo id: `cargo run -p model_prefetch -- --model-id HuggingFaceTB/SmolLM2-1.7B-Instruct`
//! - Specific revision: `... --revision main`
//! - Choose files: `... --files tokenizer.json --files config.json --files model.safetensors.index.json`
//! - Dry-run (print planned ops only): `... --dry-run`
//!
//! Notes
//! - If `--files` is omitted, the tool will try: tokenizer.json, config.json, and model weights.
//! It prefers sharded weights via `model.safetensors.index.json`, falling back to a single
//! `model.safetensors` if needed.
//! - Files are stored in the Hugging Face cache (e.g., `~/.cache/huggingface/hub` on Linux/macOS).
//! If a repo requires auth, set `HF_TOKEN` or run `huggingface-cli login`.
//! - Success prints `Prefetch complete` and exits 0; on errors, a summary is shown and the exit code is non‑zero.
use anyhow::{Context, Result, anyhow, bail};
use clap::{Parser, ValueEnum};
use hf_hub::{Repo, RepoType, api::sync::Api};
use std::io::Write;
#[derive(Clone, Debug, Copy, PartialEq, Eq, ValueEnum)]
enum Which {
V1,
V2,
V3,
V31,
V3Instruct,
V31Instruct,
V32_1b,
V32_1bInstruct,
V32_3b,
V32_3bInstruct,
#[value(name = "solar-10.7b")]
Solar10_7B,
#[value(name = "tiny-llama-1.1b-chat")]
TinyLlama1_1BChat,
#[value(name = "SmoLM2-1.7B")]
SmolLM2_1B,
#[value(name = "SmoLM2-1.7B-Instruct")]
SmolLM2_1BInstruct,
#[value(name = "SmoLM2-360M")]
SmolLM2_360M,
#[value(name = "SmoLM2-360M-Instruct")]
SmolLM2_360MInstruct,
#[value(name = "SmoLM2-135M")]
SmolLM2_135M,
#[value(name = "SmoLM2-135M-Instruct")]
SmolLM2_135MInstruct,
}
fn which_to_model_id(w: Which) -> &'static str {
match w {
Which::V1 => "Narsil/amall-7b",
Which::V2 => "meta-llama/Llama-2-7b-hf",
Which::V3 => "meta-llama/Meta-Llama-3-8B",
Which::V3Instruct => "meta-llama/Meta-Llama-3-8B-Instruct",
Which::V31 => "meta-llama/Llama-3.1-8B",
Which::V31Instruct => "meta-llama/Llama-3.1-8B-Instruct",
Which::V32_1b => "meta-llama/Llama-3.2-1B",
Which::V32_1bInstruct => "meta-llama/Llama-3.2-1B-Instruct",
Which::V32_3b => "meta-llama/Llama-3.2-3B",
Which::V32_3bInstruct => "meta-llama/Llama-3.2-3B-Instruct",
Which::Solar10_7B => "upstage/SOLAR-10.7B-v1.0",
Which::TinyLlama1_1BChat => "TinyLlama/TinyLlama-1.1B-Chat-v1.0",
Which::SmolLM2_135M => "HuggingFaceTB/SmolLM2-135M",
Which::SmolLM2_135MInstruct => "HuggingFaceTB/SmolLM2-135M-Instruct",
Which::SmolLM2_360M => "HuggingFaceTB/SmolLM2-360M",
Which::SmolLM2_360MInstruct => "HuggingFaceTB/SmolLM2-360M-Instruct",
Which::SmolLM2_1B => "HuggingFaceTB/SmolLM2-1.7B",
Which::SmolLM2_1BInstruct => "HuggingFaceTB/SmolLM2-1.7B-Instruct",
}
}
#[derive(Parser, Debug)]
#[command(
author,
version,
about = "Prefetch model files into the local Hugging Face cache",
long_about = "Pre-download tokenizer/config and weights for one or more models from Hugging Face Hub so other tools can start instantly.\n\nExamples:\n - List presets: cargo run -p model_prefetch -- --list\n - Default model: cargo run -p model_prefetch --\n - Specific repo: cargo run -p model_prefetch -- --model-id HuggingFaceTB/SmolLM2-1.7B-Instruct\n - Presets: cargo run -p model_prefetch -- --which SmoLM2-1.7B-Instruct --which tiny-llama-1.1b-chat\n - Choose files: cargo run -p model_prefetch -- --files tokenizer.json --files config.json --files model.safetensors.index.json\n - Dry run: cargo run -p model_prefetch -- --dry-run",
disable_version_flag = true
)]
struct Args {
/// One or more explicit model IDs to prefetch (e.g., HuggingFaceTB/SmolLM2-1.7B-Instruct).
#[arg(long)]
model_id: Vec<String>,
/// One or more preset model choices to prefetch.
#[arg(long, value_enum)]
which: Vec<Which>,
/// Model revision/branch to use (default: main).
#[arg(long, default_value = "main")]
revision: String,
/// Specific files to fetch (if omitted, will fetch tokenizer.json, config.json, and weights).
#[arg(long)]
files: Vec<String>,
/// Only print planned operations, do not download.
#[arg(long)]
dry_run: bool,
/// List built-in presets and exit.
#[arg(long)]
list: bool,
}
fn main() -> Result<()> {
let args = Args::parse();
if args.list {
println!(
"Built-in presets (which -> model id):\n- V1 -> {}\n- V2 -> {}\n- V3 -> {}\n- V3Instruct -> {}\n- V31 -> {}\n- V31Instruct -> {}\n- V32_1b -> {}\n- V32_1bInstruct -> {}\n- V32_3b -> {}\n- V32_3bInstruct -> {}\n- solar-10.7b -> {}\n- tiny-llama-1.1b-chat -> {}\n- SmoLM2-135M -> {}\n- SmoLM2-135M-Instruct -> {}\n- SmoLM2-360M -> {}\n- SmoLM2-360M-Instruct -> {}\n- SmoLM2-1.7B -> {}\n- SmoLM2-1.7B-Instruct -> {}",
which_to_model_id(Which::V1),
which_to_model_id(Which::V2),
which_to_model_id(Which::V3),
which_to_model_id(Which::V3Instruct),
which_to_model_id(Which::V31),
which_to_model_id(Which::V31Instruct),
which_to_model_id(Which::V32_1b),
which_to_model_id(Which::V32_1bInstruct),
which_to_model_id(Which::V32_3b),
which_to_model_id(Which::V32_3bInstruct),
which_to_model_id(Which::Solar10_7B),
which_to_model_id(Which::TinyLlama1_1BChat),
which_to_model_id(Which::SmolLM2_135M),
which_to_model_id(Which::SmolLM2_135MInstruct),
which_to_model_id(Which::SmolLM2_360M),
which_to_model_id(Which::SmolLM2_360MInstruct),
which_to_model_id(Which::SmolLM2_1B),
which_to_model_id(Which::SmolLM2_1BInstruct),
);
return Ok(());
}
let mut ids: Vec<String> = Vec::new();
ids.extend(args.model_id.iter().cloned());
for w in &args.which {
ids.push(which_to_model_id(*w).to_string());
}
if ids.is_empty() {
// Default to the same as the main CLI default.
ids.push(which_to_model_id(Which::SmolLM2_1BInstruct).to_string());
}
let api = Api::new()?;
let mut failures: Vec<(String, String)> = Vec::new();
for (idx, id) in ids.iter().enumerate() {
eprintln!(
"({}/{}) Preparing repo {}@{}...",
idx + 1,
ids.len(),
id,
&args.revision
);
std::io::stderr().flush().ok();
let repo = api.repo(Repo::with_revision(
id.clone(),
RepoType::Model,
args.revision.clone(),
));
// Determine planned files
let planned = if !args.files.is_empty() {
args.files.clone()
} else {
vec![
"tokenizer.json".to_string(),
"config.json".to_string(),
"model.safetensors".to_string(),
]
};
if args.dry_run {
println!("[dry-run] {} -> files: {}", id, planned.join(", "));
continue;
}
// Always try tokenizer and config first (if present).
for f in ["tokenizer.json", "config.json"] {
if planned.iter().any(|p| p == f) {
match repo.get(f) {
Ok(local) => {
eprintln!("fetched {} -> {}", f, local.to_string_lossy());
}
Err(err) => {
eprintln!("warning: could not fetch {} for {}: {}", f, id, err);
}
}
}
}
// Weights.
let mut got_weights = false;
if planned.iter().any(|p| p == "model.safetensors.index.json") {
match candle_examples::hub_load_safetensors(&repo, "model.safetensors.index.json") {
Ok(files) => {
got_weights = true;
for f in files {
eprintln!("prefetched shard {}", f.to_string_lossy());
}
}
Err(err) => {
eprintln!("warning: index fetch failed for {}: {}", id, err);
}
}
} else if planned.iter().any(|p| p == "model.safetensors") || args.files.is_empty() {
// Try index first, then single-file fallback.
match candle_examples::hub_load_safetensors(&repo, "model.safetensors.index.json") {
Ok(files) => {
got_weights = true;
for f in files {
eprintln!("prefetched shard {}", f.to_string_lossy());
}
}
Err(_) => match repo.get("model.safetensors") {
Ok(local) => {
got_weights = true;
eprintln!("fetched model.safetensors -> {}", local.to_string_lossy());
}
Err(err) => {
eprintln!(
"warning: could not fetch model weights for {}: {} (try specifying --files)",
id, err
);
}
},
}
}
// Also try any other explicitly listed files beyond the known ones.
for f in &args.files {
if f == "tokenizer.json"
|| f == "config.json"
|| f == "model.safetensors"
|| f == "model.safetensors.index.json"
{
continue;
}
match repo.get(f) {
Ok(local) => {
eprintln!("fetched {} -> {}", f, local.to_string_lossy());
}
Err(err) => {
failures.push((id.clone(), format!("{}: {}", f, err)));
}
}
}
if !got_weights && args.files.is_empty() {
failures.push((
id.clone(),
"no weights fetched (missing model.safetensors or index)".to_string(),
));
}
}
if failures.is_empty() {
println!("Prefetch complete for {} model(s).", ids.len());
Ok(())
} else {
eprintln!("Some models/files failed to prefetch:");
for (id, err) in &failures {
eprintln!(" - {} -> {}", id, err);
}
bail!("prefetch finished with {} error(s)", failures.len())
}
}