Adding capability of taking auxiliary data

amlb/benchmark.py

@@ -19,7 +19,7 @@
 from typing import List, Union
 
 from .job import Job, JobError, SimpleJobRunner, MultiThreadingJobRunner
-from .datasets import DataLoader, DataSourceType
+from .datasets import DataLoader, DataSourceType, DatasetWithAuxiliaryData
 from .data import DatasetType
 from .datautils import read_csv
 from .resources import get as rget, config as rconfig, output_dirs as routput_dirs
@@ -489,6 +489,10 @@ def load_data(self):
         else:
             raise ValueError("Tasks should have one property among [openml_task_id, openml_dataset_id, dataset].")
 
+        if hasattr(self._task_def, 'auxiliary_data'):
+            auxiliary_data = Benchmark.data_loader.load_auxiliary_data(DataSourceType.file, auxiliary_data=self._task_def.auxiliary_data, fold=self.fold)
+            self._dataset = DatasetWithAuxiliaryData(self._dataset, auxiliary_data)
+
     def as_job(self):
         job = Job(name=rconfig().token_separator.join([
                 'local',

amlb/datasets/__init__.py

@@ -1,6 +1,6 @@
 from enum import Enum, auto
 
-from .file import FileLoader
+from .file import FileLoader, DatasetWithAuxiliaryData
 from .openml import OpenmlLoader
 
 
@@ -24,5 +24,11 @@ def load(self, source: DataSourceType, *args, **kwargs):
         else:
             raise NotImplementedError(f"data source {source} is not supported yet")
 
+    def load_auxiliary_data(self, source: DataSourceType, *args, **kwargs):
+        if source == DataSourceType.file:
+            return self.file_loader.load_auxiliary_data(*args, **kwargs)
+        else:
+            raise NotImplementedError(f"data source {source} is not supported yet")
+
 
 __all__ = ["DataLoader", "DataSourceType"]

amlb/datasets/file.py

@@ -55,6 +55,84 @@ def load(self, dataset, fold=0):
         else:
             raise ValueError(f"Unsupported file type: {ext}")
 
+    @profile(logger=log)
+    def load_auxiliary_data(self, auxiliary_data, fold=0):
+        auxiliary_data = auxiliary_data if isinstance(auxiliary_data, ns) else ns(path=auxiliary_data)
+        log.debug("Loading auxiliary data %s", auxiliary_data)
+        paths = self._extract_auxiliary_paths(auxiliary_data.path if 'path' in auxiliary_data else auxiliary_data, fold=fold)
+        train_path = paths['train'][fold]
+        test_path = paths['test'][fold]
+        paths = dict(train=train_path, test=test_path)
+        return paths
+
+    def _extract_auxiliary_paths(self, auxiliary_data, fold=None):
+        train_search_pat = re.compile(r"(?:(.*)[_-])train_auxiliary(?:[_-](\d+))?\.\w+")
+        test_search_pat = re.compile(r"(?:(.*)[_-])test_auxiliary(?:[_-](\d+))?\.\w+")
+        if isinstance(auxiliary_data, (tuple, list)):
+            assert len(auxiliary_data) % 2 == 0, "auxiliary data list must contain an even number of paths: [train_auxiliary_0, test_auxiliary_0, train_auxiliary_1, test_auxiliary_1, ...]."
+            return self._extract_auxiliary_paths(ns(train=[p for i, p in enumerate(auxiliary_data) if i % 2 == 0],
+                                                   test=[p for i, p in enumerate(auxiliary_data) if i % 2 == 1]),
+                                                fold=fold)
+        elif isinstance(auxiliary_data, ns):
+            return dict(
+                train=[self._extract_auxiliary_paths(p)['train'][0]
+                       if i == fold else None
+                       for i, p in enumerate(as_list(auxiliary_data.train))],
+                test=[self._extract_auxiliary_paths(p)['train'][0]
+                      if i == fold else None
+                      for i, p in enumerate(as_list(auxiliary_data.test))] if 'test' in auxiliary_data else []
+            )
+        else:
+            assert isinstance(auxiliary_data, str)
+            auxiliary_data = os.path.expanduser(auxiliary_data)
+            auxiliary_data = auxiliary_data.format(**rconfig().common_dirs)
+
+        if os.path.exists(auxiliary_data):
+            if os.path.isfile(auxiliary_data):
+                # we leave the auxiliary data handling to the user
+                return dict(train=[auxiliary_data], test=[])
+            elif os.path.isdir(auxiliary_data):
+                files = list_all_files(auxiliary_data)
+                log.debug("Files found in auxiliary data folder %s: %s", auxiliary_data, files)
+                assert len(files) > 0, f"Empty folder: {auxiliary_data}"
+                if len(files) == 1:
+                    return dict(train=files, test=[])
+
+                train_matches = [m for m in [train_search_pat.search(f) for f in files] if m]
+                test_matches = [m for m in [test_search_pat.search(f) for f in files] if m]
+                # verify they're for the same dataset (just based on name)
+                assert train_matches, f"Folder {auxiliary_data} must contain at least one training auxiliary data."
+                root_names = {m[1] for m in (train_matches+test_matches)}
+                assert len(root_names) == 1, f"All dataset files in {auxiliary_data} should follow the same naming: xxxxx_train_auxiliary_N.ext or xxxxx_test_auxiliary_N.ext with N starting from 0."
+
+                train_no_fold = next((m[0] for m in train_matches if m[2] is None), None)
+                test_no_fold = next((m[0] for m in test_matches if m[2] is None), None)
+                if train_no_fold and test_no_fold:
+                    return dict(train=[train_no_fold], test=[test_no_fold])
+
+                paths = dict(train=[], test=[])
+                fold = 0
+                while fold >= 0:
+                    train = next((m[0] for m in train_matches if m[2] == str(fold)), None)
+                    test = next((m[0] for m in test_matches if m[2] == str(fold)), None)
+                    if train and test:
+                        paths['train'].append(train)
+                        paths['test'].append(test)
+                        fold += 1
+                    else:
+                        fold = -1
+                assert len(paths) > 0, f"No dataset file found in {auxiliary_data}: they should follow the naming xxxx_train_auxiliary.ext, xxxx_test_auxiliary.ext or xxxx_train_auxiliary_0.ext, xxxx_test_auxiliary_0.ext, xxxx_train_auxiliary_1.ext, ..."
+                return paths
+        elif is_valid_url(auxiliary_data):
+            cached_file = os.path.join(self._cache_dir, os.path.basename(auxiliary_data))
+            if not os.path.exists(cached_file):  # don't download if previously done
+                handler = get_file_handler(auxiliary_data)
+                assert handler.exists(auxiliary_data), f"Invalid path/url: {auxiliary_data}"
+                handler.download(auxiliary_data, dest_path=cached_file)
+            return self._extract_auxiliary_paths(cached_file)
+        else:
+            raise ValueError(f"Invalid dataset description: {auxiliary_data}")
+
     def _extract_train_test_paths(self, dataset, fold=None):
         if isinstance(dataset, (tuple, list)):
             assert len(dataset) % 2 == 0, "dataset list must contain an even number of paths: [train_0, test_0, train_1, test_1, ...]."
@@ -167,6 +245,59 @@ def _get_metadata(self, prop):
         return meta[prop]
 
 
+class DatasetWithAuxiliaryData:
+
+    def __init__(self, dataset: FileDataset, auxiliary_data_path):
+        self._dataset = dataset
+        self._train_auxiliary_data = auxiliary_data_path.get('train', None)
+        self._test_auxiliary_data = auxiliary_data_path.get('test', None)
+
+    @property
+    def train_auxiliary_data(self) -> str:
+        return self._train_auxiliary_data
+
+    @property
+    def test_auxiliary_data(self) -> str:
+        return self._test_auxiliary_data
+
+    @property
+    def type(self) -> DatasetType:
+        assert self._dataset.target is not None
+        return (DatasetType[self._dataset._type] if self._dataset._type is not None
+                else DatasetType.regression if self._dataset.target.values is None
+                else DatasetType.binary if len(self._dataset.target.values) == 2
+                else DatasetType.multiclass)
+
+    @property
+    def train(self) -> Datasplit:
+        return self._dataset._train
+
+    @property
+    def test(self) -> Datasplit:
+        return self._dataset._test
+
+    @property
+    def features(self) -> List[Feature]:
+        return self._get_metadata('features')
+
+    @property
+    def target(self) -> Feature:
+        return self._get_metadata('target')
+
+    @memoize
+    def _get_metadata(self, prop):
+        meta = self._dataset._train.load_metadata()
+        return meta[prop]
+
+    @profile(logger=log)
+    def release(self, properties=None):
+        """
+        Call this to release cached properties and optimize memory once in-memory data are not needed anymore.
+        :param properties:
+        """
+        self._dataset.release(properties)
+
+
 class FileDatasplit(Datasplit):
 
     def __init__(self, dataset: FileDataset, format: str, path: str):

frameworks/AutoGluon/__init__.py

@@ -19,6 +19,10 @@ def run(dataset: Dataset, config: TaskConfig):
         ),
         problem_type=dataset.type.name  # AutoGluon problem_type is using same names as amlb.data.DatasetType
     )
+    if hasattr(dataset, 'train_auxiliary_data'):
+        data['train_auxiliary_data'] = dict(path=dataset.train_auxiliary_data)
+    if hasattr(dataset, 'test_auxiliary_data'):
+        data['test_auxiliary_data'] = dict(path=dataset.test_auxiliary_data) 
 
     return run_in_venv(__file__, "exec.py",
                        input_data=data, dataset=dataset, config=config)