candding 0.1.0

//! Locating model files on disk, and with the `download` feature fetching them into the hugging face cache.

pub mod cache;
#[cfg(feature = "download")]
pub mod download;
pub mod pattern;

use std::path::Path;
use std::path::PathBuf;

pub use cache::HubConfig;

use crate::error::CanddingError;
use crate::error::Result;
use crate::registry::WeightsFormat;

/// Files a safetensors model needs; onnx, bin, gguf and openvino exports are skipped.
pub const INCLUDE_PATTERNS: &[&str] = &["*.json", "*.safetensors", "1_Pooling/*", "*_Dense/*"];

/// The same list plus three root-only pickles: `pytorch_model.bin` and bge-m3's `sparse_linear.pt` and
/// `colbert_linear.pt`. Patterns match the whole path and `*` crosses `/`, so they are not directory-scoped.
const PICKLE_INCLUDE_PATTERNS: &[&str] = &[
  "*.json",
  "*.safetensors",
  "1_Pooling/*",
  "*_Dense/*",
  "pytorch_model.bin",
  "sparse_linear.pt",
  "colbert_linear.pt",
];

/// The same base list plus `*.npy`, `*.vocab` and `stopwords.txt`, for `Qdrant/minicoil-v1`'s NumPy head;
/// none matches its ONNX tree, so a fetch still skips the 458 MB this port never reads.
const NPY_INCLUDE_PATTERNS: &[&str] = &[
  "*.json",
  "*.safetensors",
  "1_Pooling/*",
  "*_Dense/*",
  "*.npy",
  "*.vocab",
  "stopwords.txt",
];

/// `WeightsFormat::NoWeights` has nothing to fetch, so its list is empty; `candding fetch` explains that itself
/// before it reaches a download.
pub fn include_patterns(format: WeightsFormat) -> &'static [&'static str] {
  match format {
    WeightsFormat::Safetensors => INCLUDE_PATTERNS,
    WeightsFormat::Pickle => PICKLE_INCLUDE_PATTERNS,
    WeightsFormat::Npy => NPY_INCLUDE_PATTERNS,
    WeightsFormat::NoWeights => &[],
  }
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub enum ModelSource {
  Hub { repo_id: String, revision: String },
  Local(PathBuf),
}

impl ModelSource {
  /// An existing directory is a local model; anything else is a hub id.
  pub fn parse(spec: &str, revision: &str) -> Self {
    let path = Path::new(spec);
    if path.is_dir() {
      ModelSource::Local(path.to_path_buf())
    } else {
      ModelSource::Hub {
        repo_id: spec.to_string(),
        revision: revision.to_string(),
      }
    }
  }
}

/// The CLI command that fills the cache with the files `resolve` would have downloaded.
pub fn download_hint(repo_id: &str, revision: &str) -> String {
  let mut hint = format!("candding fetch {repo_id}");
  if revision != "main" {
    hint.push_str(&format!(" --revision {revision}"));
  }
  hint
}

pub fn resolve(source: &ModelSource, format: WeightsFormat) -> Result<PathBuf> {
  match source {
    ModelSource::Local(dir) => dir
      .is_dir()
      .then(|| dir.clone())
      .ok_or_else(|| CanddingError::MissingDir(dir.clone())),
    ModelSource::Hub { repo_id, revision } => {
      resolve_hub(&HubConfig::from_env(), repo_id, revision, format)
    }
  }
}

/// Whether `resolve` finds `source` on disk, with no download.
pub(crate) fn is_present(source: &ModelSource) -> bool {
  match source {
    ModelSource::Local(_) => true,
    ModelSource::Hub { repo_id, revision } => {
      cache::snapshot_dir(&HubConfig::from_env().cache, repo_id, revision).is_some()
    }
  }
}

/// Fetches what the cache lacks at `revision`, contacting the hub even when a snapshot is already there.
#[cfg(feature = "download")]
pub fn fetch(repo_id: &str, revision: &str, format: WeightsFormat) -> Result<PathBuf> {
  download::fetch(&HubConfig::from_env(), repo_id, revision, format)
}

fn resolve_hub(
  config: &HubConfig,
  repo_id: &str,
  revision: &str,
  format: WeightsFormat,
) -> Result<PathBuf> {
  if let Some(dir) = cache::snapshot_dir(&config.cache, repo_id, revision) {
    return Ok(dir);
  }
  fetch_missing(config, repo_id, revision, format)
}

#[cfg(feature = "download")]
fn fetch_missing(
  config: &HubConfig,
  repo_id: &str,
  revision: &str,
  format: WeightsFormat,
) -> Result<PathBuf> {
  download::fetch(config, repo_id, revision, format)
}

#[cfg(not(feature = "download"))]
fn fetch_missing(
  _config: &HubConfig,
  repo_id: &str,
  revision: &str,
  _format: WeightsFormat,
) -> Result<PathBuf> {
  Err(CanddingError::ModelNotFound {
    id: repo_id.to_string(),
    hint: download_hint(repo_id, revision),
  })
}

#[cfg(test)]
mod tests {
  use tempfile::tempdir;

  use super::*;

  #[test]
  fn parse_prefers_existing_directory() {
    let dir = tempdir().unwrap();
    let spec = dir.path().to_str().unwrap();
    assert_eq!(
      ModelSource::parse(spec, "main"),
      ModelSource::Local(dir.path().to_path_buf())
    );
    assert_eq!(
      ModelSource::parse("BAAI/bge-small-en-v1.5", "v2"),
      ModelSource::Hub {
        repo_id: "BAAI/bge-small-en-v1.5".into(),
        revision: "v2".into()
      }
    );
  }

  #[test]
  fn hint_names_the_fetch_command_and_any_revision_but_main() {
    assert_eq!(
      download_hint("org/model", "abc"),
      "candding fetch org/model --revision abc"
    );
    assert_eq!(
      download_hint("org/model", "main"),
      "candding fetch org/model"
    );
  }

  /// huggingface_hub matches a pattern against the whole relative path, so the three root patterns leave subdirectory files behind.
  #[test]
  fn the_pickle_patterns_add_the_root_bin_and_bge_m3s_two_side_heads() {
    let safetensors = include_patterns(WeightsFormat::Safetensors);
    let pickle = include_patterns(WeightsFormat::Pickle);
    assert_eq!(safetensors, INCLUDE_PATTERNS);
    assert_eq!(&pickle[..safetensors.len()], safetensors);
    assert_eq!(
      pickle[safetensors.len()..],
      ["pytorch_model.bin", "sparse_linear.pt", "colbert_linear.pt"]
    );
    assert!(pickle.contains(&"sparse_linear.pt"));
    assert!(pickle.contains(&"colbert_linear.pt"));
    assert!(!safetensors.contains(&"sparse_linear.pt"));
    assert!(!safetensors.contains(&"colbert_linear.pt"));
    assert!(!safetensors.contains(&"pytorch_model.bin"));
  }

  #[test]
  fn the_npy_patterns_add_the_npy_vocab_and_stopwords_patterns() {
    let safetensors = include_patterns(WeightsFormat::Safetensors);
    let npy = include_patterns(WeightsFormat::Npy);
    assert_eq!(&npy[..safetensors.len()], safetensors);
    assert_eq!(
      npy[safetensors.len()..],
      ["*.npy", "*.vocab", "stopwords.txt"]
    );
    assert!(!safetensors.contains(&"*.npy"));
    assert!(!npy.contains(&"pytorch_model.bin"));
  }

  #[test]
  fn missing_local_dir_errors() {
    let err = resolve(
      &ModelSource::Local(PathBuf::from("/definitely/missing")),
      WeightsFormat::Safetensors,
    )
    .unwrap_err();
    assert!(matches!(err, CanddingError::MissingDir(_)));
  }

  #[cfg(not(feature = "download"))]
  #[test]
  fn missing_hub_model_reports_hint() {
    let cache = tempdir().unwrap();
    let config = HubConfig {
      cache: cache.path().to_path_buf(),
      offline: false,
      endpoint: None,
      token: None,
    };
    let err = resolve_hub(&config, "org/missing", "main", WeightsFormat::Safetensors).unwrap_err();
    let message = err.to_string();
    assert!(message.contains("`download` feature"), "{message}");
    match err {
      CanddingError::ModelNotFound { id, hint } => {
        assert_eq!(id, "org/missing");
        assert_eq!(hint, "candding fetch org/missing");
      }
      other => panic!("unexpected {other:?}"),
    }
  }
}