Preserve Hugging Face metadata when checkpointing

shared/client/src/state/init.rs

@@ -21,7 +21,11 @@ use psyche_modeling::{
 };
 use psyche_network::{BlobTicket, SecretKey};
 use psyche_watcher::OpportunisticData;
-use std::{collections::HashMap, path::PathBuf, sync::Arc};
+use std::{
+    collections::HashMap,
+    path::{Path, PathBuf},
+    sync::Arc,
+};
 use tch::{Kind, Tensor};
 use thiserror::Error;
 use tokenizers::{ModelWrapper, Tokenizer, models::wordlevel::WordLevel};
@@ -153,6 +157,24 @@ struct RawLoadedModel {
 type OneshotModelParameterSender = oneshot::Sender<HashMap<String, Tensor>>;
 type OneShotModelConfigSender = oneshot::Sender<(String, Tokenizer, Vec<String>)>;
 
+fn is_checkpoint_extra_file(path: &Path) -> bool {
+    let Some(file_name) = path.file_name().and_then(|name| name.to_str()) else {
+        return false;
+    };
+
+    matches!(
+        file_name,
+        "added_tokens.json"
+            | "config.json"
+            | "generation_config.json"
+            | "merges.txt"
+            | "special_tokens_map.json"
+            | "tokenizer.json"
+            | "tokenizer_config.json"
+            | "vocab.json"
+    ) || file_name.ends_with(".py")
+}
+
 pub struct RunInitConfigAndIO {
     pub init_config: RunInitConfig,
 
@@ -369,14 +391,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 +496,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 +927,40 @@ impl RunInitConfigAndIO {
         ))
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::is_checkpoint_extra_file;
+    use std::path::Path;
+
+    #[test]
+    fn checkpoint_extra_file_filter_includes_model_metadata() {
+        for file_name in [
+            "added_tokens.json",
+            "config.json",
+            "generation_config.json",
+            "merges.txt",
+            "special_tokens_map.json",
+            "tokenizer.json",
+            "tokenizer_config.json",
+            "vocab.json",
+            "modeling_custom.py",
+        ] {
+            assert!(
+                is_checkpoint_extra_file(Path::new(file_name)),
+                "expected {file_name} to be included"
+            );
+        }
+
+        for file_name in [
+            "model.safetensors",
+            "model.safetensors.index.json",
+            "README.md",
+        ] {
+            assert!(
+                !is_checkpoint_extra_file(Path::new(file_name)),
+                "expected {file_name} to be excluded"
+            );
+        }
+    }
+}

shared/data-provider/src/hub.rs

@@ -13,7 +13,7 @@ use std::{path::PathBuf, time::Instant};
 use tokio::sync::mpsc;
 use tracing::{error, info};
 
-const MODEL_EXTENSIONS: [&str; 3] = [".safetensors", ".json", ".py"];
+const MODEL_EXTENSIONS: [&str; 4] = [".safetensors", ".json", ".py", ".txt"];
 const DATASET_EXTENSIONS: [&str; 1] = [".parquet"];
 
 /// Strip leading/trailing whitespace and control characters from a repo identifier.
@@ -274,3 +274,22 @@ pub async fn upload_to_hub(
 
     Ok(())
 }
+
+#[cfg(test)]
+mod tests {
+    use super::{check_extensions, MODEL_EXTENSIONS};
+    use hf_hub::api::Siblings;
+
+    #[test]
+    fn model_extensions_include_tokenizer_merges() {
+        let merges = Siblings {
+            rfilename: "merges.txt".to_string(),
+        };
+        let readme = Siblings {
+            rfilename: "README.md".to_string(),
+        };
+
+        assert!(check_extensions(&merges, &MODEL_EXTENSIONS));
+        assert!(!check_extensions(&readme, &MODEL_EXTENSIONS));
+    }
+}