Add eval set to CatboostClassifier

This commit is contained in:
Emre 2022-10-13 23:01:09 +03:00
parent c6d2eed4fc
commit 7f05b44376
No known key found for this signature in database
GPG Key ID: 0EAD2EE11B666ABA

View File

@ -30,6 +30,14 @@ class CatboostClassifier(BaseClassifierModel):
label=data_dictionary["train_labels"], label=data_dictionary["train_labels"],
weight=data_dictionary["train_weights"], 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( cbr = CatBoostClassifier(
allow_writing_files=True, allow_writing_files=True,
@ -40,6 +48,6 @@ class CatboostClassifier(BaseClassifierModel):
init_model = self.get_init_model(dk.pair) 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 return cbr