diff --git a/cleanlab/datalab/datalab.py b/cleanlab/datalab/datalab.py index 407d0e5a9..a7f813dcc 100644 --- a/cleanlab/datalab/datalab.py +++ b/cleanlab/datalab/datalab.py @@ -592,7 +592,7 @@ def list_default_issue_types(self) -> List[str]: default_issue_types.extend(DEFAULT_CLEANVISION_ISSUES.keys()) return default_issue_types - def save(self, path: str, force: bool = False) -> None: + def save(self, path: str, force: bool = False, save_dataset: bool = True) -> None: """Saves this DatalabĂ‚ object to file (all files are in folder at `path/`). We do not guarantee saved Datalab can be loaded from future versions of cleanlab. @@ -608,7 +608,7 @@ def save(self, path: str, force: bool = False) -> None: ---- You have to save the Dataset yourself separately if you want it saved to file. """ - _Serializer.serialize(path=path, datalab=self, force=force) + _Serializer.serialize(path=path, datalab=self, force=force, save_dataset=save_dataset) save_message = f"Saved Datalab to folder: {path}" print(save_message) diff --git a/cleanlab/datalab/internal/serialize.py b/cleanlab/datalab/internal/serialize.py index 487bed354..0b58f8e4d 100644 --- a/cleanlab/datalab/internal/serialize.py +++ b/cleanlab/datalab/internal/serialize.py @@ -76,7 +76,7 @@ def _validate_version(datalab: Datalab) -> None: ) @classmethod - def serialize(cls, path: str, datalab: Datalab, force: bool, save_dataset: bool = True) -> None: + def serialize(cls, path: str, datalab: Datalab, force: bool, save_dataset: bool) -> None: """Serializes the datalab object to disk. Parameters diff --git a/tests/datalab/datalab/test_datalab.py b/tests/datalab/datalab/test_datalab.py index ffc82ec71..4661d97e7 100644 --- a/tests/datalab/datalab/test_datalab.py +++ b/tests/datalab/datalab/test_datalab.py @@ -451,6 +451,14 @@ def test_save(self, lab, tmp_path, mock_issues, mock_issue_summary, monkeypatch) def test_load(self, lab, tmp_path, dataset, mock_issues, mock_issue_summary, monkeypatch): """Test that the save and load methods work.""" + # Load without saved datalab raises warning + lab.save(tmp_path, force=True, save_dataset=False) + with pytest.warns( + UserWarning, + match="The Datalab object was loaded without a dataset and will have limited functionality.", + ): + loaded_lab = Datalab.load(tmp_path) + monkeypatch.setattr(lab, "issues", mock_issues) monkeypatch.setattr(lab, "issue_summary", mock_issue_summary)