add alg name in learner

mljar · Nov 12, 2018 · 4b83d72 · 4b83d72
1 parent e0e6e10
commit 4b83d72
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 2 deletions.
diff --git a/supervised/models/learner.py b/supervised/models/learner.py
@@ -8,7 +8,9 @@ class Learner():
     def __init__(self, params):
         self.params = params
         self.stop_training = False
-        self.version = None
+        self.library_version = None
+        self.algorithm_name = 'Unknown'
+        self.algorithm_short_name = 'Unknown'
         self.model = None
         self.uid = str(uuid.uuid4())
 

diff --git a/supervised/models/learner_xgboost.py b/supervised/models/learner_xgboost.py
@@ -33,7 +33,9 @@ class XgbLearner(Learner):
     '''
     def __init__(self, params):
         Learner.__init__(self, params)
-        self.version = xgb.__version__
+        self.library_version = xgb.__version__
+        self.algorithm_name = 'Extreme Gradient Boosting'
+        self.algorithm_short_name = 'Xgboost'
         self.model_file = self.uid + '.xgb.model'
         self.model_file_path = '/tmp/' + self.model_file
 

diff --git a/supervised/models/tests/test_learner_factory.py b/supervised/models/tests/test_learner_factory.py
@@ -15,3 +15,4 @@ def test_fit(self):
             'eval_metric': 'logloss'
         }
         learner = LearnerFactory.get_learner(params)
+        self.assertEqual(learner.algorithm_short_name, 'Xgboost')