[Feature] set xgboost num boost round (#39)

RektPunk · Oct 5, 2024 · ddde704 · ddde704 · github-actions · Oct 5, 2024
1 parent 3055d02
commit ddde704
Show file tree

Hide file tree

Showing 4 changed files with 4 additions and 2 deletions.
diff --git a/pyproject.toml b/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "rektgbm"
-version = "0.1.5"
+version = "0.1.6"
 description = "No-brainer machine learning solution to achieve satisfactory performance"
 authors = ["RektPunk <[email protected]>"]
 license = "MIT"

diff --git a/rektgbm/__init__.py b/rektgbm/__init__.py
@@ -3,4 +3,4 @@
 from rektgbm.gbm import RektGBM
 from rektgbm.optimizer import RektOptimizer
 
-__version__ = "0.1.5"
+__version__ = "0.1.6"
diff --git a/rektgbm/engine.py b/rektgbm/engine.py
@@ -45,6 +45,7 @@ def fit(
             evals_result = {}
             self.model = xgb.train(
                 dtrain=dtrain,
+                num_boost_round=100,
                 verbose_eval=False,
                 params=self.params,
                 evals_result=evals_result,

diff --git a/tests/test_engine.py b/tests/test_engine.py
@@ -66,6 +66,7 @@ def test_rektengine_xgb_fit(mock_dataset, mock_xgb_model):
     xgb.train.assert_called_once_with(
         dtrain="dtrain_mock",
         verbose_eval=False,
+        num_boost_round=100,
         params={"eval_metric": "rmse"},
         evals_result={},
         evals=[("dvalid_mock", _VALID_STR)],