mod api;
use crate::error::{Result, TrustformersError};
use api::{CommitFile, LFS_INLINE_THRESHOLD_BYTES};
use std::path::{Path, PathBuf};
use tracing::{debug, info, warn};
const HF_HUB_URL: &str = "https://huggingface.co";
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum RepoType {
#[default]
Model,
Dataset,
Space,
}
impl RepoType {
pub fn as_str(&self) -> &'static str {
match self {
RepoType::Model => "model",
RepoType::Dataset => "dataset",
RepoType::Space => "space",
}
}
}
#[derive(Debug, Clone)]
pub struct UploadConfig {
pub token: String,
pub repo_id: String,
pub repo_type: RepoType,
pub revision: String,
pub commit_message: String,
pub create_if_missing: bool,
pub private: bool,
pub base_url: String,
pub dry_run: bool,
}
impl Default for UploadConfig {
fn default() -> Self {
Self {
token: String::new(),
repo_id: String::new(),
repo_type: RepoType::Model,
revision: "main".to_string(),
commit_message: "Upload via TrustformeRS".to_string(),
create_if_missing: true,
private: false,
base_url: HF_HUB_URL.to_string(),
dry_run: false,
}
}
}
#[derive(Debug, Clone)]
pub struct UploadFile {
pub local_path: PathBuf,
pub repo_path: String,
}
impl UploadFile {
pub fn new(local_path: impl Into<PathBuf>, repo_path: impl Into<String>) -> Self {
Self {
local_path: local_path.into(),
repo_path: repo_path.into(),
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct UploadResult {
pub repo_id: String,
pub revision: String,
pub commit_url: Option<String>,
pub commit_oid: Option<String>,
pub files_uploaded: Vec<String>,
pub dry_run: bool,
}
pub struct HubUploader {
config: UploadConfig,
}
impl HubUploader {
pub fn new(config: UploadConfig) -> Self {
Self { config }
}
pub fn validate(&self) -> Result<()> {
if self.config.token.is_empty() {
return Err(missing_credentials_error(&self.config.repo_id));
}
if self.config.repo_id.is_empty() {
return Err(TrustformersError::InvalidInput {
message: "Repository ID cannot be empty".to_string(),
parameter: Some("repo_id".to_string()),
expected: None,
received: None,
suggestion: None,
});
}
if !self.config.repo_id.contains('/') {
return Err(TrustformersError::InvalidInput {
message: "Repository ID must be in format 'username/repo-name'".to_string(),
parameter: Some("repo_id".to_string()),
expected: Some("username/repo-name".to_string()),
received: Some(self.config.repo_id.clone()),
suggestion: None,
});
}
if self.config.revision.is_empty() {
return Err(TrustformersError::InvalidInput {
message: "Revision/branch name cannot be empty".to_string(),
parameter: Some("revision".to_string()),
expected: None,
received: None,
suggestion: None,
});
}
Ok(())
}
pub fn repo_exists(&self) -> Result<bool> {
self.validate()?;
if self.config.dry_run {
debug!(repo_id = %self.config.repo_id, "repo_exists: dry run, skipping network call");
return Ok(false);
}
api::run_blocking(api::repo_exists(
&self.config.base_url,
self.config.repo_type,
&self.config.repo_id,
&self.config.token,
))
.map_err(|e| hub_error_to_trustformers(e, &self.config.repo_id))
}
pub fn create_repo(&self) -> Result<String> {
self.validate()?;
if self.config.dry_run {
let url = format!("{}/{}", self.config.base_url, self.config.repo_id);
info!(repo_id = %self.config.repo_id, "create_repo: dry run, not contacting the Hub");
return Ok(url);
}
api::run_blocking(api::create_repo(
&self.config.base_url,
self.config.repo_type,
&self.config.repo_id,
self.config.private,
&self.config.token,
))
.map_err(|e| hub_error_to_trustformers(e, &self.config.repo_id))
}
pub fn upload_file(&self, file: &UploadFile) -> Result<UploadResult> {
self.upload_files(std::slice::from_ref(file))
}
pub fn upload_files(&self, files: &[UploadFile]) -> Result<UploadResult> {
self.validate()?;
if files.is_empty() {
return Err(TrustformersError::InvalidInput {
message: "File list cannot be empty".to_string(),
parameter: Some("files".to_string()),
expected: None,
received: None,
suggestion: None,
});
}
let mut repo_paths = Vec::with_capacity(files.len());
let mut commit_files = Vec::with_capacity(files.len());
for file in files {
if !file.local_path.exists() {
return Err(TrustformersError::Io {
message: format!("File not found: {}", file.local_path.display()),
path: Some(file.local_path.display().to_string()),
suggestion: Some("Ensure all files exist before uploading".to_string()),
});
}
if file.repo_path.is_empty() {
return Err(TrustformersError::InvalidInput {
message: "Repository path cannot be empty for one of the files".to_string(),
parameter: Some("repo_path".to_string()),
expected: None,
received: None,
suggestion: None,
});
}
let content = std::fs::read(&file.local_path).map_err(|e| TrustformersError::Io {
message: format!("Cannot read file: {e}"),
path: Some(file.local_path.display().to_string()),
suggestion: None,
})?;
if content.len() as u64 >= LFS_INLINE_THRESHOLD_BYTES {
return Err(hub_error_to_trustformers(
HubError::LfsRequired {
path: file.repo_path.clone(),
size: content.len() as u64,
},
&self.config.repo_id,
));
}
repo_paths.push(file.repo_path.clone());
commit_files.push(CommitFile {
repo_path: file.repo_path.clone(),
content,
});
}
if self.config.dry_run {
info!(
file_count = files.len(),
repo_id = %self.config.repo_id,
"upload_files: dry run, not contacting the Hub"
);
return Ok(UploadResult {
repo_id: self.config.repo_id.clone(),
revision: self.config.revision.clone(),
commit_url: None,
commit_oid: None,
files_uploaded: repo_paths,
dry_run: true,
});
}
if self.config.create_if_missing && !self.repo_exists()? {
info!(repo_id = %self.config.repo_id, "Repository does not exist yet; creating it");
self.create_repo()?;
}
let outcome = api::run_blocking(api::commit(
&self.config.base_url,
self.config.repo_type,
&self.config.repo_id,
&self.config.revision,
&self.config.commit_message,
&commit_files,
&[],
&self.config.token,
))
.map_err(|e| hub_error_to_trustformers(e, &self.config.repo_id))?;
info!(
file_count = files.len(),
repo_id = %self.config.repo_id,
commit_url = ?outcome.commit_url,
"Uploaded files to Hub"
);
Ok(UploadResult {
repo_id: self.config.repo_id.clone(),
revision: self.config.revision.clone(),
commit_url: outcome.commit_url,
commit_oid: outcome.commit_oid,
files_uploaded: repo_paths,
dry_run: false,
})
}
pub fn upload_directory(&self, local_dir: &Path, repo_prefix: &str) -> Result<UploadResult> {
self.validate()?;
if !local_dir.is_dir() {
return Err(TrustformersError::Io {
message: format!("Not a directory: {}", local_dir.display()),
path: Some(local_dir.display().to_string()),
suggestion: Some("Provide a path to an existing directory".to_string()),
});
}
let files = collect_files_recursive(local_dir, local_dir, repo_prefix)?;
if files.is_empty() {
warn!(
dir = %local_dir.display(),
"Directory is empty; nothing to upload"
);
return Ok(UploadResult {
repo_id: self.config.repo_id.clone(),
revision: self.config.revision.clone(),
commit_url: None,
commit_oid: None,
files_uploaded: vec![],
dry_run: self.config.dry_run,
});
}
self.upload_files(&files)
}
pub fn delete_file(&self, repo_path: &str) -> Result<()> {
self.validate()?;
if repo_path.is_empty() {
return Err(TrustformersError::InvalidInput {
message: "Repository path cannot be empty".to_string(),
parameter: Some("repo_path".to_string()),
expected: None,
received: None,
suggestion: None,
});
}
if self.config.dry_run {
info!(repo_path = %repo_path, repo_id = %self.config.repo_id, "delete_file: dry run, not contacting the Hub");
return Ok(());
}
api::run_blocking(api::commit(
&self.config.base_url,
self.config.repo_type,
&self.config.repo_id,
&self.config.revision,
&format!("Delete {repo_path}"),
&[],
std::slice::from_ref(&repo_path.to_string()),
&self.config.token,
))
.map_err(|e| hub_error_to_trustformers(e, &self.config.repo_id))?;
info!(repo_path = %repo_path, repo_id = %self.config.repo_id, "Deleted file from Hub");
Ok(())
}
}
fn collect_files_recursive(
base_dir: &Path,
current_dir: &Path,
repo_prefix: &str,
) -> Result<Vec<UploadFile>> {
let mut files = Vec::new();
let entries = std::fs::read_dir(current_dir).map_err(|e| TrustformersError::Io {
message: format!("Cannot read directory: {e}"),
path: Some(current_dir.display().to_string()),
suggestion: None,
})?;
for entry_result in entries {
let entry = entry_result.map_err(|e| TrustformersError::Io {
message: format!("Cannot read directory entry: {e}"),
path: Some(current_dir.display().to_string()),
suggestion: None,
})?;
let path = entry.path();
if path.is_dir() {
let mut sub_files = collect_files_recursive(base_dir, &path, repo_prefix)?;
files.append(&mut sub_files);
} else {
let relative = path.strip_prefix(base_dir).map_err(|e| TrustformersError::Io {
message: format!("Path prefix stripping failed: {e}"),
path: Some(path.display().to_string()),
suggestion: None,
})?;
let repo_path = if repo_prefix.is_empty() {
relative.display().to_string()
} else {
format!("{}/{}", repo_prefix, relative.display())
};
let repo_path = repo_path.replace('\\', "/");
files.push(UploadFile {
local_path: path.clone(),
repo_path,
});
}
}
Ok(files)
}
pub struct HubUploaderBuilder {
config: UploadConfig,
}
impl HubUploaderBuilder {
pub fn new(token: impl Into<String>, repo_id: impl Into<String>) -> Self {
let config = UploadConfig {
token: token.into(),
repo_id: repo_id.into(),
..Default::default()
};
Self { config }
}
pub fn repo_type(mut self, repo_type: RepoType) -> Self {
self.config.repo_type = repo_type;
self
}
pub fn revision(mut self, revision: impl Into<String>) -> Self {
self.config.revision = revision.into();
self
}
pub fn commit_message(mut self, msg: impl Into<String>) -> Self {
self.config.commit_message = msg.into();
self
}
pub fn private(mut self, private: bool) -> Self {
self.config.private = private;
self
}
pub fn create_if_missing(mut self, create: bool) -> Self {
self.config.create_if_missing = create;
self
}
pub fn base_url(mut self, base_url: impl Into<String>) -> Self {
self.config.base_url = base_url.into();
self
}
pub fn dry_run(mut self, dry_run: bool) -> Self {
self.config.dry_run = dry_run;
self
}
pub fn build(self) -> Result<HubUploader> {
let uploader = HubUploader::new(self.config);
uploader.validate()?;
Ok(uploader)
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum HubError {
Unauthorized { message: String },
MissingCredentials { message: String },
NotFound {
repo_id: String,
path: Option<String>,
},
RequestFailed { status_code: u16, message: String },
Io {
message: String,
path: Option<String>,
},
InvalidInput { message: String },
Network { message: String },
FeatureUnavailable { message: String },
LfsRequired { path: String, size: u64 },
}
impl std::fmt::Display for HubError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
HubError::Unauthorized { message } => write!(f, "Unauthorized: {message}"),
HubError::MissingCredentials { message } => write!(f, "Missing credentials: {message}"),
HubError::NotFound { repo_id, path } => {
if let Some(p) = path {
write!(f, "Not found: {repo_id}/{p}")
} else {
write!(f, "Not found: {repo_id}")
}
},
HubError::RequestFailed {
status_code,
message,
} => {
write!(f, "Request failed (HTTP {status_code}): {message}")
},
HubError::Io { message, path } => {
if let Some(p) = path {
write!(f, "IO error at {p}: {message}")
} else {
write!(f, "IO error: {message}")
}
},
HubError::InvalidInput { message } => write!(f, "Invalid input: {message}"),
HubError::Network { message } => write!(f, "Network error: {message}"),
HubError::FeatureUnavailable { message } => write!(f, "Feature unavailable: {message}"),
HubError::LfsRequired { path, size } => {
write!(
f,
"'{path}' is {size} bytes, at or above the {}-byte inline-upload threshold, \
and would require real Git-LFS object storage, which is not implemented",
LFS_INLINE_THRESHOLD_BYTES
)
},
}
}
}
impl std::error::Error for HubError {}
impl From<TrustformersError> for HubError {
fn from(e: TrustformersError) -> Self {
HubError::RequestFailed {
status_code: 0,
message: e.to_string(),
}
}
}
fn hub_error_to_trustformers(error: HubError, repo_id: &str) -> TrustformersError {
match error {
HubError::MissingCredentials { .. } => missing_credentials_error(repo_id),
HubError::Unauthorized { message } => TrustformersError::Hub {
message,
model_id: repo_id.to_string(),
endpoint: None,
suggestion: Some("Check that the API token is valid and has write access".to_string()),
recovery_actions: vec![],
},
HubError::FeatureUnavailable { message } => TrustformersError::Hub {
message,
model_id: repo_id.to_string(),
endpoint: None,
suggestion: Some("Rebuild with `--features hub`".to_string()),
recovery_actions: vec![],
},
HubError::LfsRequired { path, size } => TrustformersError::Hub {
message: format!(
"'{path}' is {size} bytes and would require real Git-LFS object storage, which \
is not implemented"
),
model_id: repo_id.to_string(),
endpoint: None,
suggestion: Some("Upload large files through the Hub web UI or the CLI's `huggingface-cli upload-large-folder` until LFS object upload is implemented here".to_string()),
recovery_actions: vec![],
},
HubError::NotFound { repo_id, path } => TrustformersError::Hub {
message: format!("Not found: {repo_id}{}", path.map(|p| format!("/{p}")).unwrap_or_default()),
model_id: repo_id,
endpoint: None,
suggestion: None,
recovery_actions: vec![],
},
HubError::RequestFailed { status_code, message } => TrustformersError::Hub {
message: format!("HTTP {status_code}: {message}"),
model_id: repo_id.to_string(),
endpoint: None,
suggestion: None,
recovery_actions: vec![],
},
HubError::Network { message } => TrustformersError::Hub {
message,
model_id: repo_id.to_string(),
endpoint: None,
suggestion: Some("Check network connectivity".to_string()),
recovery_actions: vec![],
},
HubError::Io { message, path } => TrustformersError::Io {
message,
path,
suggestion: None,
},
HubError::InvalidInput { message } => TrustformersError::InvalidInput {
message,
parameter: None,
expected: None,
received: None,
suggestion: None,
},
}
}
fn missing_credentials_error(repo_id: &str) -> TrustformersError {
TrustformersError::Hub {
message: "Missing credentials: a Hub API token is required to upload".to_string(),
model_id: repo_id.to_string(),
endpoint: None,
suggestion: Some(
"Set `UploadConfig::token` (e.g. from the `HF_TOKEN` environment variable)".to_string(),
),
recovery_actions: vec![],
}
}
#[derive(Debug, Clone)]
pub struct HubUploadConfig {
pub repo_id: String,
pub token: String,
pub commit_message: String,
pub private: bool,
pub revision: Option<String>,
pub base_url: Option<String>,
pub dry_run: bool,
}
impl HubUploadConfig {
pub fn new(
repo_id: impl Into<String>,
token: impl Into<String>,
commit_message: impl Into<String>,
) -> Self {
Self {
repo_id: repo_id.into(),
token: token.into(),
commit_message: commit_message.into(),
private: false,
revision: None,
base_url: None,
dry_run: false,
}
}
pub fn with_private(mut self, private: bool) -> Self {
self.private = private;
self
}
pub fn with_revision(mut self, revision: impl Into<String>) -> Self {
self.revision = Some(revision.into());
self
}
pub fn with_base_url(mut self, base_url: impl Into<String>) -> Self {
self.base_url = Some(base_url.into());
self
}
pub fn with_dry_run(mut self, dry_run: bool) -> Self {
self.dry_run = dry_run;
self
}
pub fn effective_revision(&self) -> &str {
self.revision.as_deref().unwrap_or("main")
}
fn validate(&self) -> std::result::Result<(), HubError> {
if self.token.is_empty() {
return Err(HubError::MissingCredentials {
message: "API token cannot be empty".to_string(),
});
}
if self.repo_id.is_empty() {
return Err(HubError::InvalidInput {
message: "repo_id cannot be empty".to_string(),
});
}
if !self.repo_id.contains('/') {
return Err(HubError::InvalidInput {
message: format!(
"repo_id must be in 'username/repo-name' format, got '{}'",
self.repo_id
),
});
}
Ok(())
}
}
#[derive(Debug, Clone, Default)]
pub struct HubUploadProgress {
pub total_files: usize,
pub uploaded_files: usize,
pub total_bytes: u64,
pub uploaded_bytes: u64,
}
impl HubUploadProgress {
pub fn new(total_files: usize, total_bytes: u64) -> Self {
Self {
total_files,
uploaded_files: 0,
total_bytes,
uploaded_bytes: 0,
}
}
pub fn record_file(&mut self, bytes: u64) {
self.uploaded_files += 1;
self.uploaded_bytes += bytes;
}
pub fn fraction(&self) -> f64 {
if self.total_bytes == 0 {
if self.total_files == 0 {
1.0
} else {
self.uploaded_files as f64 / self.total_files as f64
}
} else {
self.uploaded_bytes as f64 / self.total_bytes as f64
}
}
pub fn is_complete(&self) -> bool {
self.uploaded_files >= self.total_files
}
}
pub fn sha256(data: &[u8]) -> String {
use sha2::{Digest, Sha256};
let mut hasher = Sha256::new();
hasher.update(data);
hex::encode(hasher.finalize())
}
pub fn sha256_file(path: &Path) -> std::result::Result<String, HubError> {
let data = std::fs::read(path).map_err(|e| HubError::Io {
message: format!("Cannot read file for hashing: {e}"),
path: Some(path.display().to_string()),
})?;
Ok(sha256(&data))
}
#[derive(Debug, Clone)]
pub struct SingleFileUploadResult {
pub remote_url: String,
pub commit_url: Option<String>,
pub commit_oid: Option<String>,
pub file_size: u64,
pub sha256: String,
}
impl HubUploader {
pub fn from_hub_config(cfg: HubUploadConfig) -> std::result::Result<Self, HubError> {
cfg.validate()?;
let revision = cfg.effective_revision().to_string();
let upload_config = UploadConfig {
token: cfg.token,
repo_id: cfg.repo_id,
repo_type: RepoType::Model,
revision,
commit_message: cfg.commit_message,
create_if_missing: true,
private: cfg.private,
base_url: cfg.base_url.unwrap_or_else(|| HF_HUB_URL.to_string()),
dry_run: cfg.dry_run,
};
Ok(Self::new(upload_config))
}
pub fn upload_file_path(
&self,
local_path: &str,
remote_path: &str,
) -> std::result::Result<SingleFileUploadResult, HubError> {
let path = Path::new(local_path);
if !path.exists() {
return Err(HubError::Io {
message: format!("File not found: {local_path}"),
path: Some(local_path.to_string()),
});
}
if remote_path.is_empty() {
return Err(HubError::InvalidInput {
message: "remote_path cannot be empty".to_string(),
});
}
let metadata = path.metadata().map_err(|e| HubError::Io {
message: format!("Cannot read file metadata: {e}"),
path: Some(local_path.to_string()),
})?;
let file_size = metadata.len();
let sha256 = sha256_file(path)?;
let result =
self.upload_file(&UploadFile::new(path, remote_path)).map_err(HubError::from)?;
let remote_url = format!(
"{}/{}/blob/{}/{}",
self.config.base_url, self.config.repo_id, self.config.revision, remote_path
);
Ok(SingleFileUploadResult {
remote_url,
commit_url: result.commit_url,
commit_oid: result.commit_oid,
file_size,
sha256,
})
}
pub fn upload_model(
&self,
model_dir: &str,
) -> std::result::Result<Vec<SingleFileUploadResult>, HubError> {
let base = Path::new(model_dir);
if !base.is_dir() {
return Err(HubError::Io {
message: format!("Not a directory: {model_dir}"),
path: Some(model_dir.to_string()),
});
}
self.upload_dir_filtered(base, |name| {
name.ends_with(".json")
|| name.ends_with(".safetensors")
|| name.ends_with(".bin")
|| name.ends_with(".pt")
|| name.ends_with(".ckpt")
|| name.ends_with(".msgpack")
|| name.ends_with(".model")
|| name == "README.md"
})
}
pub fn upload_tokenizer(
&self,
tokenizer_dir: &str,
) -> std::result::Result<Vec<SingleFileUploadResult>, HubError> {
let base = Path::new(tokenizer_dir);
if !base.is_dir() {
return Err(HubError::Io {
message: format!("Not a directory: {tokenizer_dir}"),
path: Some(tokenizer_dir.to_string()),
});
}
self.upload_dir_filtered(base, |name| {
name.ends_with("tokenizer.json")
|| name.ends_with("tokenizer_config.json")
|| name.ends_with("vocab.json")
|| name.ends_with("vocab.txt")
|| name.ends_with("merges.txt")
|| name.ends_with("special_tokens_map.json")
|| name.ends_with("added_tokens.json")
|| name.ends_with(".model")
|| name.ends_with("spiece.model")
})
}
pub fn create_repo_typed(&self, repo_type: RepoType) -> std::result::Result<String, HubError> {
let mut cfg = self.config.clone();
cfg.repo_type = repo_type;
let tmp = HubUploader::new(cfg);
tmp.create_repo().map_err(HubError::from)
}
pub fn delete_remote_file(&self, remote_path: &str) -> std::result::Result<(), HubError> {
self.delete_file(remote_path).map_err(HubError::from)
}
fn upload_dir_filtered<F>(
&self,
base: &Path,
filter: F,
) -> std::result::Result<Vec<SingleFileUploadResult>, HubError>
where
F: Fn(&str) -> bool,
{
let entries = collect_files_recursive_hub(base, base)?;
let mut results = Vec::new();
for (local_path, repo_path) in entries {
let file_name = local_path.file_name().and_then(|n| n.to_str()).unwrap_or("");
if !filter(file_name) {
continue;
}
let local_str = local_path.display().to_string();
let result = self.upload_file_path(&local_str, &repo_path)?;
results.push(result);
}
Ok(results)
}
}
fn collect_files_recursive_hub(
base: &Path,
current: &Path,
) -> std::result::Result<Vec<(PathBuf, String)>, HubError> {
let mut files = Vec::new();
let entries = std::fs::read_dir(current).map_err(|e| HubError::Io {
message: format!("Cannot read directory: {e}"),
path: Some(current.display().to_string()),
})?;
for entry_result in entries {
let entry = entry_result.map_err(|e| HubError::Io {
message: format!("Cannot read directory entry: {e}"),
path: Some(current.display().to_string()),
})?;
let path = entry.path();
if path.is_dir() {
let mut sub = collect_files_recursive_hub(base, &path)?;
files.append(&mut sub);
} else {
let relative = path.strip_prefix(base).map_err(|e| HubError::Io {
message: format!("Path strip prefix failed: {e}"),
path: Some(path.display().to_string()),
})?;
let repo_path = relative.display().to_string().replace('\\', "/");
files.push((path, repo_path));
}
}
Ok(files)
}
#[cfg(test)]
mod tests;