add test

Author: calderjo

tests/test_hf_datasets.py

@@ -1,7 +1,8 @@
 import unittest
 
-from datasets import Dataset
-
+import datasets
+import pandas as pd
+import warnings
 
 class TestHuggingFaceDatasets(unittest.TestCase):
 
@@ -10,7 +11,13 @@ def some_func(batch):
             batch['label'] = 'foo'
             return batch
 
-        df = Dataset.from_dict({'text': ['Kaggle rocks!']})
+        df = datasets.Dataset.from_dict({'text': ['Kaggle rocks!']})
         mapped_df = df.map(some_func)
 
-        self.assertEqual('foo', mapped_df[0]['label'])
+        self.assertEqual('foo', mapped_df[0]['label'])
+        
+    def test_load_dataset(self):
+        warnings.simplefilter(action='ignore', category=FutureWarning)
+        dataset = datasets.load_dataset("csv", data_files="/input/tests/data/train.csv")
+        full_data = pd.DataFrame(dataset['train'])
+        self.assertFalse(full_data.empty)