1
0
mirror of https://github.com/microsoft/qlib.git synced 2026-07-01 18:11:18 +08:00

change code style in Catboost and XGboost

This commit is contained in:
Hong Zhang
2020-11-26 14:15:54 +08:00
parent e8bb5061b0
commit 305b9372c7
2 changed files with 25 additions and 25 deletions

View File

@@ -34,14 +34,14 @@ class CatBoostModel(Model):
def fit(
self,
dataset: DatasetH,
num_boost_round=1000,
early_stopping_rounds=50,
verbose_eval=20,
evals_result=dict(),
num_boost_round = 1000,
early_stopping_rounds = 50,
verbose_eval = 20,
evals_result = dict(),
**kwargs
):
df_train, df_valid = dataset.prepare(
["train", "valid"], col_set=["feature", "label"], data_key=DataHandlerLP.DK_L
["train", "valid"], col_set = ["feature", "label"], data_key = DataHandlerLP.DK_L
)
x_train, y_train = df_train["feature"], df_train["label"]
x_valid, y_valid = df_valid["feature"], df_valid["label"]
@@ -52,8 +52,8 @@ class CatBoostModel(Model):
else:
raise ValueError("CatBoost doesn't support multi-label training")
train_pool = Pool(data=x_train, label=y_train_1d)
valid_pool = Pool(data=x_valid, label=y_valid_1d)
train_pool = Pool(data = x_train, label = y_train_1d)
valid_pool = Pool(data = x_valid, label = y_valid_1d)
# Initialize the catboost model
self._params["iterations"] = num_boost_round
@@ -63,7 +63,7 @@ class CatBoostModel(Model):
self.model = CatBoost(self._params, **kwargs)
# train the model
self.model.fit(train_pool, eval_set=valid_pool, use_best_model=True, **kwargs)
self.model.fit(train_pool, eval_set = valid_pool, use_best_model = True, **kwargs)
evals_result = self.model.get_evals_result()
evals_result["train"] = list(evals_result["learn"].values())[0]
@@ -72,8 +72,8 @@ class CatBoostModel(Model):
def predict(self, dataset):
if self.model is None:
raise ValueError("model is not fitted yet!")
x_test = dataset.prepare("test", col_set="feature")
return pd.Series(self.model.predict(x_test.values), index=x_test.index)
x_test = dataset.prepare("test", col_set = "feature")
return pd.Series(self.model.predict(x_test.values), index = x_test.index)
if __name__ == "__main__":

View File

@@ -30,15 +30,15 @@ class XGBModel(Model):
def fit(
self,
dataset: DatasetH,
num_boost_round=1000,
early_stopping_rounds=50,
verbose_eval=20,
evals_result=dict(),
num_boost_round = 1000,
early_stopping_rounds = 50,
verbose_eval = 20,
evals_result = dict(),
**kwargs
):
df_train, df_valid = dataset.prepare(
["train", "valid"], col_set=["feature", "label"], data_key=DataHandlerLP.DK_L
["train", "valid"], col_set = ["feature", "label"], data_key = DataHandlerLP.DK_L
)
x_train, y_train = df_train["feature"], df_train["label"]
x_valid, y_valid = df_valid["feature"], df_valid["label"]
@@ -49,16 +49,16 @@ class XGBModel(Model):
else:
raise ValueError("XGBoost doesn't support multi-label training")
dtrain = xgb.DMatrix(x_train.values, label=y_train_1d)
dvalid = xgb.DMatrix(x_valid.values, label=y_valid_1d)
dtrain = xgb.DMatrix(x_train.values, label = y_train_1d)
dvalid = xgb.DMatrix(x_valid.values, label = y_valid_1d)
self.model = xgb.train(
self._params,
dtrain=dtrain,
num_boost_round=num_boost_round,
evals=[(dtrain, "train"), (dvalid, "valid")],
early_stopping_rounds=early_stopping_rounds,
verbose_eval=verbose_eval,
evals_result=evals_result,
dtrain = dtrain,
num_boost_round = num_boost_round,
evals = [(dtrain, "train"), (dvalid, "valid")],
early_stopping_rounds = early_stopping_rounds,
verbose_eval = verbose_eval,
evals_result = evals_result,
**kwargs
)
evals_result["train"] = list(evals_result["train"].values())[0]
@@ -67,5 +67,5 @@ class XGBModel(Model):
def predict(self, dataset):
if self.model is None:
raise ValueError("model is not fitted yet!")
x_test = dataset.prepare("test", col_set="feature")
return pd.Series(self.model.predict(xgb.DMatrix(x_test.values)), index=x_test.index)
x_test = dataset.prepare("test", col_set = "feature")
return pd.Series(self.model.predict(xgb.DMatrix(x_test.values)), index = x_test.index)