Returning None instead of "python" if dataset is unformatted (hugging…

…face#526) * returning None instead of "python" if dataset is unformatted * updated tests
Kaiseryang1224 · Aug 24, 2020 · af17331 · af17331
1 parent c93a19e
commit af17331
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 3 deletions.
diff --git a/src/nlp/arrow_dataset.py b/src/nlp/arrow_dataset.py
@@ -477,7 +477,7 @@ def __repr__(self):
     @property
     def format(self):
         return {
-            "type": "python" if self._format_type is None else self._format_type,
+            "type": self._format_type,
             "format_kwargs": self._format_kwargs,
             "columns": self.column_names if self._format_columns is None else self._format_columns,
             "output_all_columns": self._output_all_columns,

diff --git a/tests/test_arrow_dataset.py b/tests/test_arrow_dataset.py
@@ -106,7 +106,7 @@ def test_set_format_numpy_multiple_columns(self):
             self.assertListEqual(list(dset["col_1"].shape), [4])
             np.testing.assert_array_equal(dset["col_1"], np.array([3, 2, 1, 0]))
 
-        self.assertEqual(dset.format["type"], "python")
+        self.assertEqual(dset.format["type"], None)
         self.assertEqual(dset.format["format_kwargs"], {})
         self.assertEqual(dset.format["columns"], dset.column_names)
         self.assertEqual(dset.format["output_all_columns"], False)

diff --git a/tests/test_dataset_dict.py b/tests/test_dataset_dict.py
@@ -48,7 +48,7 @@ def test_set_format_numpy(self):
                 self.assertEqual(dset_split[0]["col_1"].item(), 3)
 
         for dset_split in dset.values():
-            self.assertEqual(dset_split.format["type"], "python")
+            self.assertEqual(dset_split.format["type"], None)
             self.assertEqual(dset_split.format["format_kwargs"], {})
             self.assertEqual(dset_split.format["columns"], dset_split.column_names)
             self.assertEqual(dset_split.format["output_all_columns"], False)