Merge pull request #7585 from aemr3/add-eval-set-catboost

Add eval set to CatboostClassifier
This commit is contained in:
Robert Caulk 2022-10-15 16:08:13 +02:00 committed by GitHub
commit 46ba3bb357
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -30,6 +30,14 @@ class CatboostClassifier(BaseClassifierModel):
label=data_dictionary["train_labels"],
weight=data_dictionary["train_weights"],
)
if self.freqai_info.get("data_split_parameters", {}).get("test_size", 0.1) == 0:
test_data = None
else:
test_data = Pool(
data=data_dictionary["test_features"],
label=data_dictionary["test_labels"],
weight=data_dictionary["test_weights"],
)
cbr = CatBoostClassifier(
allow_writing_files=True,
@ -40,6 +48,6 @@ class CatboostClassifier(BaseClassifierModel):
init_model = self.get_init_model(dk.pair)
cbr.fit(train_data, init_model=init_model)
cbr.fit(X=train_data, eval_set=test_data, init_model=init_model)
return cbr