Path type (tracel-ai#699)

hojjatabdollahi · Aug 27, 2023 · 0f7864f · 0f7864f
1 parent 7f558bd
commit 0f7864f
Show file tree

Hide file tree

Showing 7 changed files with 10 additions and 10 deletions.
diff --git a/burn-book/src/basic-workflow/inference.md b/burn-book/src/basic-workflow/inference.md
@@ -31,7 +31,7 @@ Now let's create a simple `infer` method in which we will load our trained model
 ```rust , ignore
 pub fn infer<B: Backend>(artifact_dir: &str, device: B::Device, item: MNISTItem) {
     let config =
-        TrainingConfig::load(&format!("{artifact_dir}/config.json")).expect("A config exists");
+        TrainingConfig::load(format!("{artifact_dir}/config.json")).expect("A config exists");
     let record = CompactRecorder::new()
         .load(format!("{artifact_dir}/model").into())
         .expect("Failed to load trained model");

diff --git a/burn-book/src/basic-workflow/training.md b/burn-book/src/basic-workflow/training.md
@@ -78,7 +78,7 @@ pub struct TrainingConfig {
 pub fn train<B: ADBackend>(artifact_dir: &str, config: TrainingConfig, device: B::Device) {
     std::fs::create_dir_all(artifact_dir).ok();
     config
-        .save(&format!("{artifact_dir}/config.json"))
+        .save(format!("{artifact_dir}/config.json"))
         .expect("Save without error");
 
     B::seed(config.seed);

diff --git a/burn-core/src/config.rs b/burn-core/src/config.rs
@@ -44,7 +44,7 @@ pub trait Config: serde::Serialize + serde::de::DeserializeOwned {
     ///
     /// The output of the save operation.
     #[cfg(feature = "std")]
-    fn save(&self, file: &str) -> std::io::Result<()> {
+    fn save<P: AsRef<std::path::Path>>(&self, file: P) -> std::io::Result<()> {
         std::fs::write(file, config_to_json(self))
     }
 
@@ -58,9 +58,9 @@ pub trait Config: serde::Serialize + serde::de::DeserializeOwned {
     ///
     /// The loaded configuration.
     #[cfg(feature = "std")]
-    fn load(file: &str) -> Result<Self, ConfigError> {
-        let content = std::fs::read_to_string(file)
-            .map_err(|_| ConfigError::FileNotFound(file.to_string()))?;
+    fn load<P: AsRef<std::path::Path>>(file: P) -> Result<Self, ConfigError> {
+        let content = std::fs::read_to_string(file.as_ref().clone())
+            .map_err(|_| ConfigError::FileNotFound(file.as_ref().to_string_lossy().to_string()))?;
         config_from_str(&content)
     }
 

diff --git a/examples/guide/src/inference.rs b/examples/guide/src/inference.rs
@@ -10,7 +10,7 @@ use burn_dataset::source::huggingface::MNISTItem;
 
 pub fn infer<B: Backend>(artifact_dir: &str, device: B::Device, item: MNISTItem) {
     let config =
-        TrainingConfig::load(&format!("{artifact_dir}/config.json")).expect("A config exists");
+        TrainingConfig::load(format!("{artifact_dir}/config.json")).expect("A config exists");
     let record = CompactRecorder::new()
         .load(format!("{artifact_dir}/model").into())
         .expect("Failed to load trained model");

diff --git a/examples/guide/src/training.rs b/examples/guide/src/training.rs
@@ -67,7 +67,7 @@ pub struct TrainingConfig {
 pub fn train<B: ADBackend>(artifact_dir: &str, config: TrainingConfig, device: B::Device) {
     std::fs::create_dir_all(artifact_dir).ok();
     config
-        .save(&format!("{artifact_dir}/config.json"))
+        .save(format!("{artifact_dir}/config.json"))
         .expect("Save without error");
 
     B::seed(config.seed);

diff --git a/examples/text-classification/src/training.rs b/examples/text-classification/src/training.rs
@@ -107,7 +107,7 @@ pub fn train<B: ADBackend, D: TextClassificationDataset + 'static>(
     let model_trained = learner.fit(dataloader_train, dataloader_test);
 
     // Save the configuration and the trained model
-    config.save(&format!("{artifact_dir}/config.json")).unwrap();
+    config.save(format!("{artifact_dir}/config.json")).unwrap();
     CompactRecorder::new()
         .record(
             model_trained.into_record(),

diff --git a/examples/text-generation/src/training.rs b/examples/text-generation/src/training.rs
@@ -83,7 +83,7 @@ pub fn train<B: ADBackend, D: Dataset<TextGenerationItem> + 'static>(
 
     let model_trained = learner.fit(dataloader_train, dataloader_test);
 
-    config.save(&format!("{artifact_dir}/config.json")).unwrap();
+    config.save(format!("{artifact_dir}/config.json")).unwrap();
 
     DefaultRecorder::new()
         .record(