Skip to main content

entrenar/hf_pipeline/fetcher/
options.rs

1//! Fetch options for HuggingFace model downloads.
2//!
3//! Provides configuration for model fetching including revision, files, and security settings.
4
5use std::path::PathBuf;
6
7/// Options for model fetching
8#[derive(Debug, Clone)]
9pub struct FetchOptions {
10    /// Git revision (branch, tag, or commit)
11    pub revision: String,
12    /// Specific files to download
13    pub files: Vec<String>,
14    /// Allow PyTorch pickle files (SECURITY RISK)
15    pub allow_pytorch_pickle: bool,
16    /// Expected SHA256 hash for verification
17    pub verify_sha256: Option<String>,
18    /// Cache directory
19    pub cache_dir: Option<PathBuf>,
20}
21
22impl Default for FetchOptions {
23    fn default() -> Self {
24        Self {
25            revision: "main".into(),
26            files: vec![],
27            allow_pytorch_pickle: false,
28            verify_sha256: None,
29            cache_dir: None,
30        }
31    }
32}
33
34impl FetchOptions {
35    /// Create new options
36    #[must_use]
37    pub fn new() -> Self {
38        Self::default()
39    }
40
41    /// Set revision
42    #[must_use]
43    pub fn revision(mut self, rev: impl Into<String>) -> Self {
44        self.revision = rev.into();
45        self
46    }
47
48    /// Set files to download
49    #[must_use]
50    pub fn files(mut self, files: &[&str]) -> Self {
51        self.files = files.iter().map(|s| (*s).to_string()).collect();
52        self
53    }
54
55    /// Allow PyTorch pickle files (SECURITY RISK)
56    #[must_use]
57    pub fn allow_pytorch_pickle(mut self, allow: bool) -> Self {
58        self.allow_pytorch_pickle = allow;
59        self
60    }
61
62    /// Set SHA256 hash for verification
63    #[must_use]
64    pub fn verify_sha256(mut self, hash: impl Into<String>) -> Self {
65        self.verify_sha256 = Some(hash.into());
66        self
67    }
68
69    /// Set cache directory
70    #[must_use]
71    pub fn cache_dir(mut self, dir: impl Into<PathBuf>) -> Self {
72        self.cache_dir = Some(dir.into());
73        self
74    }
75}