识别任务类型

4paradigm · Aug 9, 2021 · 373c58e · 373c58e
1 parent 44755fa
commit 373c58e
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/autox/autox.py b/autox/autox.py
@@ -30,6 +30,12 @@ def __init__(self, target, train_name, test_name, path, feature_type = {}, id =
         self.dfs_['FE_all'] = None
         self.sub = None
 
+        # 识别任务类型
+        if self.dfs_[self.info_['train_name']][self.info_['target']].nunique() == 2:
+            self.info_['task_type'] = 'binary'
+        else:
+            self.info_['task_type'] = 'regression'
+
     def concat_train_test(self):
         self.info_['shape_of_train'] = len(self.dfs_[self.info_['train_name']])
         self.info_['shape_of_test'] = len(self.dfs_[self.info_['test_name']])