fix: dataset and tests

Added dataset saving fix and test on this
stochasticai · May 11, 2023 · e5a5b6d · e5a5b6d
1 parent ab787fa
commit e5a5b6d
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/src/xturing/datasets/instruction_dataset.py b/src/xturing/datasets/instruction_dataset.py
@@ -59,7 +59,6 @@ def __init__(
         else:
             path = Path(path)
             assert Path(path).exists(), "path does not exist"
-
             if path.is_dir():
                 self.data = load_from_disk(str(path))
             elif path.suffix == ".jsonl":
@@ -123,7 +122,7 @@ def __getitem__(self, idx):
         return self.data["train"][idx]
 
     def save(self, path):
-        return self.data.save_to_disk(path)
+        return self.data["train"].save_to_disk(path)
 
     @classmethod
     def generate_dataset(

diff --git a/tests/xturing/datasets/test_instruction_dataset.py b/tests/xturing/datasets/test_instruction_dataset.py
@@ -42,3 +42,4 @@ def test_features_dataset():
         "target": "second text",
         "instruction": "second instruction",
     }
+    dataset.save(".")