Merge pull request #157 from openml/fix/labels

FIX allow to retrieve class labels by the tasks class label flag

Author: mfeurer

openml/datasets/dataset.py

@@ -249,9 +249,24 @@ def get_data(self, target=None, target_dtype=int, include_row_id=False,
         else:
             return rval
 
-    def _retrieve_class_labels(self):
-        """Reads the datasets arff to determine the class-labels, and returns those.
-        If the task has no class labels (for example a regression problem) it returns None."""
+    def retrieve_class_labels(self, target_name='class'):
+        """Reads the datasets arff to determine the class-labels.
+
+        If the task has no class labels (for example a regression problem)
+        it returns None. Necessary because the data returned by get_data
+        only contains the indices of the classes, while OpenML needs the real
+        classname when uploading the results of a run.
+
+        Parameters
+        ----------
+        target_name : str
+            Name of the target attribute
+
+        Returns
+        -------
+        list
+        """
+
         # TODO improve performance, currently reads the whole file
         # Should make a method that only reads the attributes
         arffFileName = self.data_file
@@ -267,10 +282,8 @@ def _retrieve_class_labels(self):
             arffData = arff.ArffDecoder().decode(fh, return_type=return_type)
 
         dataAttributes = dict(arffData['attributes'])
-        if('class' in dataAttributes):
-            return dataAttributes['class']
-        elif('Class' in dataAttributes):
-            return dataAttributes['Class']
+        if target_name in dataAttributes:
+            return dataAttributes[target_name]
         else:
             return None
 

openml/runs/run.py

@@ -151,7 +151,7 @@ def run_task(task, model):
     arff_datacontent = []
 
     dataset = task.get_dataset()
-    X, Y = dataset.get_data(target=task.target_feature)
+    X, Y = dataset.get_data(target=task.target_name)
 
     class_labels = task.class_labels
     if class_labels is None:

openml/tasks/functions.py

@@ -236,7 +236,7 @@ def get_task(task_id):
 
     # TODO look into either adding the class labels to task xml, or other
     # way of reading it.
-    class_labels = dataset._retrieve_class_labels()
+    class_labels = dataset.retrieve_class_labels(task.target_name)
     task.class_labels = class_labels
     return task
 

openml/tasks/task.py

@@ -9,14 +9,14 @@
 
 
 class OpenMLTask(object):
-    def __init__(self, task_id, task_type, data_set_id, target_feature,
+    def __init__(self, task_id, task_type, data_set_id, target_name,
                  estimation_procedure_type, data_splits_url,
                  estimation_parameters, evaluation_measure, cost_matrix,
                  class_labels=None):
         self.task_id = int(task_id)
         self.task_type = task_type
         self.dataset_id = int(data_set_id)
-        self.target_feature = target_feature
+        self.target_name = target_name
         self.estimation_procedure = dict()
         self.estimation_procedure["type"] = estimation_procedure_type
         self.estimation_procedure["data_splits_url"] = data_splits_url
@@ -43,7 +43,7 @@ def get_X_and_y(self):
             target_dtype = float
         else:
             raise NotImplementedError(self.task_type)
-        X_and_y = dataset.get_data(target=self.target_feature,
+        X_and_y = dataset.get_data(target=self.target_name,
                                    target_dtype=target_dtype)
         return X_and_y
 

tests/datasets/test_datasets.py

@@ -207,6 +207,14 @@ def test_publish_dataset(self):
         dataset.publish()
         self.assertIsInstance(dataset.dataset_id, int)
 
+    def test__retrieve_class_labels(self):
+        openml.config.set_cache_directory(self.static_cache_dir)
+        labels = openml.datasets.get_dataset(2).retrieve_class_labels()
+        self.assertEqual(labels, ['1', '2', '3', '4', '5', 'U'])
+        labels = openml.datasets.get_dataset(2).retrieve_class_labels(
+            target_name='product-type')
+        self.assertEqual(labels, ['C', 'H', 'G'])
+
     def test_upload_dataset_with_url(self):
         dataset = OpenMLDataset(
             name="UploadTestWithURL", version=1, description="test",

tests/tasks/test_task.py

@@ -18,12 +18,12 @@ def test_get_clustering_task(self):
         self.assertRaisesRegexp(KeyError, 'oml:target_feature',
                                 openml.tasks.get_task, 10128)
 
-    @mock.patch('openml.datasets.get_dataset')
+    @mock.patch('openml.datasets.get_dataset', autospec=True)
     def test_get_dataset(self, patch):
         patch.return_value = mock.MagicMock()
         mm = mock.MagicMock()
-        patch.return_value._retrieve_class_labels = mm
-        patch.return_value._retrieve_class_labels.return_value = 'LA'
+        patch.return_value.retrieve_class_labels = mm
+        patch.return_value.retrieve_class_labels.return_value = 'LA'
         retval = openml.tasks.get_task(1)
         self.assertEqual(patch.call_count, 1)
         self.assertIsInstance(retval, openml.OpenMLTask)