fix: upload tokenizer/config files with hub checkpoints

shared/client/src/state/init.rs

@@ -39,6 +39,29 @@ use super::{
 };
 use iroh_blobs::api::Tag;
 
+/// Files from the model's repo that are re-uploaded alongside every
+/// checkpoint safetensors file, so a checkpoint repo remains a complete,
+/// loadable model repository on its own.
+const CHECKPOINT_EXTRA_FILE_NAMES: [&str; 10] = [
+    "added_tokens.json",
+    "config.json",
+    "generation_config.json",
+    "merges.txt",
+    "special_tokens_map.json",
+    "tokenizer.json",
+    "tokenizer_config.json",
+    "tokenizer.model",
+    "vocab.json",
+    "chat_template.jinja",
+];
+
+fn is_checkpoint_extra_file(file: &PathBuf) -> bool {
+    let Some(name) = file.file_name().and_then(|name| name.to_str()) else {
+        return false;
+    };
+    CHECKPOINT_EXTRA_FILE_NAMES.contains(&name) || name.ends_with(".py")
+}
+
 pub struct RunInitConfig {
     // identity for connecting to the data server
     pub identity: NodeIdentity,
@@ -369,14 +392,7 @@ impl RunInitConfigAndIO {
                                 let repo_files = model_is_local;
                                 let checkpoint_extra_files = repo_files
                                     .iter()
-                                    .filter(|file| {
-                                        file.ends_with("config.json")
-                                            || file.ends_with("tokenizer.json")
-                                            || file.ends_with("tokenizer_config.json")
-                                            || file.ends_with("special_tokens_map.json")
-                                            || file.ends_with("generation_config.json")
-                                            || file.ends_with(".py")
-                                    })
+                                    .filter(|file| is_checkpoint_extra_file(file))
                                     .cloned()
                                     .collect();
                                 let tokenizer = Arc::new(auto_tokenizer(&repo_files)?);
@@ -481,14 +497,7 @@ impl RunInitConfigAndIO {
 
                                 let checkpoint_extra_files = repo_files
                                     .iter()
-                                    .filter(|file| {
-                                        file.ends_with("config.json")
-                                            || file.ends_with("tokenizer.json")
-                                            || file.ends_with("tokenizer_config.json")
-                                            || file.ends_with("special_tokens_map.json")
-                                            || file.ends_with("generation_config.json")
-                                            || file.ends_with(".py")
-                                    })
+                                    .filter(|file| is_checkpoint_extra_file(file))
                                     .cloned()
                                     .collect();
                                 let tokenizer = Arc::new(auto_tokenizer(&repo_files)?);
@@ -919,3 +928,39 @@ impl RunInitConfigAndIO {
         ))
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::{CHECKPOINT_EXTRA_FILE_NAMES, is_checkpoint_extra_file};
+    use std::path::Path;
+
+    #[test]
+    fn checkpoint_extra_file_filter_includes_model_metadata() {
+        for file_name in CHECKPOINT_EXTRA_FILE_NAMES {
+            assert!(
+                is_checkpoint_extra_file(&PathBuf::from(file_name)),
+                "expected {file_name} to be included"
+            );
+        }
+        assert!(is_checkpoint_extra_file(&PathBuf::from("modeling_custom.py")));
+        assert!(is_checkpoint_extra_file(&PathBuf::from(
+            "some/nested/dir/tokenizer_config.json"
+        )));
+    }
+
+    #[test]
+    fn checkpoint_extra_file_filter_excludes_weights_and_readmes() {
+        for file_name in [
+            "model.safetensors",
+            "model-00001-of-00002.safetensors",
+            "model.safetensors.index.json",
+            "README.md",
+            "training_args.bin",
+        ] {
+            assert!(
+                !is_checkpoint_extra_file(&PathBuf::from(file_name)),
+                "expected {file_name} to be excluded"
+            );
+        }
+    }
+}

shared/data-provider/src/file_extensions.rs

@@ -1,2 +1,10 @@
 pub const DATA_FILE_EXTENSIONS: [&str; 3] = ["npy", "bin", "ds"];
 pub const PARQUET_EXTENSION: &str = "parquet";
+
+/// File extensions downloaded when fetching a model repository. Besides the
+/// weights, these cover the config/tokenizer files that must be re-uploaded
+/// with every checkpoint so the checkpoint repo stays a complete, loadable
+/// model repository on its own.
+pub const MODEL_FILE_EXTENSIONS: [&str; 6] = [
+    ".safetensors", ".json", ".py", ".txt", ".model", ".jinja",
+];

shared/data-provider/src/gcs.rs

@@ -1,4 +1,5 @@
 use crate::errors::{DownloadError, UploadError};
+use crate::file_extensions::MODEL_FILE_EXTENSIONS;
 use chrono::{DateTime, Utc};
 use google_cloud_storage::client::{Client, ClientConfig};
 use google_cloud_storage::http::objects::upload::Media;
@@ -51,8 +52,6 @@ pub struct GcsManifestMetadata {
     pub run_id: String,
 }
 
-const MODEL_EXTENSIONS: [&str; 3] = [".safetensors", ".json", ".py"];
-
 fn get_cache_base(bucket: &str) -> PathBuf {
     // Use HF_HOME if set, otherwise fall back to ~/.cache
     std::env::var("HF_HOME")
@@ -185,7 +184,7 @@ pub async fn download_model_from_gcs_async(
             info!("No manifest found, downloading model without manifest");
             let cache_dir = get_cache_dir_no_manifest(bucket, prefix);
             std::fs::create_dir_all(&cache_dir)?;
-            download_files_no_manifest(&client, bucket, prefix, &cache_dir, &MODEL_EXTENSIONS).await
+            download_files_no_manifest(&client, bucket, prefix, &cache_dir, &MODEL_FILE_EXTENSIONS).await
         }
     }
 }

shared/data-provider/src/hub.rs

@@ -1,4 +1,5 @@
 use crate::errors::UploadError;
+use crate::file_extensions::MODEL_FILE_EXTENSIONS;
 use crate::hub::model::HubRepo;
 use hf_hub::{
     Cache, Repo, RepoType,
@@ -13,7 +14,6 @@ use std::{path::PathBuf, time::Instant};
 use tokio::sync::mpsc;
 use tracing::{error, info};
 
-const MODEL_EXTENSIONS: [&str; 3] = [".safetensors", ".json", ".py"];
 const DATASET_EXTENSIONS: [&str; 1] = [".parquet"];
 
 /// Strip leading/trailing whitespace and control characters from a repo identifier.
@@ -107,7 +107,7 @@ pub async fn download_model_repo_async(
         token,
         max_concurrent_downloads,
         progress_bar,
-        &MODEL_EXTENSIONS,
+        &MODEL_FILE_EXTENSIONS,
     )
     .await
 }
@@ -180,7 +180,7 @@ pub fn download_model_repo_sync(
         cache,
         token,
         progress_bar,
-        &MODEL_EXTENSIONS,
+        &MODEL_FILE_EXTENSIONS,
     )
 }
 
@@ -274,3 +274,42 @@ pub async fn upload_to_hub(
 
     Ok(())
 }
+
+#[cfg(test)]
+mod tests {
+    use super::{MODEL_FILE_EXTENSIONS, check_extensions};
+    use hf_hub::api::Siblings;
+
+    #[test]
+    fn model_extensions_include_tokenizer_artifacts() {
+        for file_name in [
+            "merges.txt",
+            "tokenizer.model",
+            "chat_template.jinja",
+            "config.json",
+            "added_tokens.json",
+            "model.safetensors",
+        ] {
+            let sibling = Siblings {
+                rfilename: file_name.to_string(),
+            };
+            assert!(
+                check_extensions(&sibling, &MODEL_FILE_EXTENSIONS),
+                "expected {file_name} to be downloaded"
+            );
+        }
+    }
+
+    #[test]
+    fn model_extensions_exclude_non_model_files() {
+        for file_name in ["README.md", "training_args.bin"] {
+            let sibling = Siblings {
+                rfilename: file_name.to_string(),
+            };
+            assert!(
+                !check_extensions(&sibling, &MODEL_FILE_EXTENSIONS),
+                "expected {file_name} to be skipped"
+            );
+        }
+    }
+}

shared/data-provider/src/lib.rs

@@ -16,7 +16,7 @@ pub use data_provider::DataProvider;
 pub use dataset::{Dataset, Field, Row, Split};
 pub use dummy::DummyDataProvider;
 pub use errors::{DownloadError, UploadError};
-pub use file_extensions::{DATA_FILE_EXTENSIONS, PARQUET_EXTENSION};
+pub use file_extensions::{DATA_FILE_EXTENSIONS, MODEL_FILE_EXTENSIONS, PARQUET_EXTENSION};
 pub use gcs::{
     GcsCheckpointManifest, GcsManifestMetadata, GcsUploadInfo, ManifestFileEntry, ManifestMetadata,
     download_model_from_gcs_async, download_model_from_gcs_sync, upload_to_gcs,