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)