Add Xgboost Model

qlib/contrib/model/xgboost.py

@@ -0,0 +1,64 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+import numpy as np
+import pandas as pd
+import xgboost as xgb
+
+from ...model.base import Model
+from ...data.dataset import DatasetH
+from ...data.dataset.handler import DataHandlerLP
+
+
+class XGBModel(Model):
+    """XGBModel Model"""
+
+    def __init__(self, obj="mse", **kwargs):
+        if obj not in {"mse", "binary"}:
+            raise NotImplementedError
+        self._params = {"obj": obj}
+        self._params.update(kwargs)
+        self.model = None
+
+    def fit(
+        self,
+        dataset: DatasetH,
+        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
+        )
+        x_train, y_train = df_train["feature"], df_train["label"]
+        x_valid, y_valid = df_valid["feature"], df_valid["label"]
+
+        # Lightgbm need 1D array as its label
+        if y_train.values.ndim == 2 and y_train.values.shape[1] == 1:
+            y_train_1d, y_valid_1d = np.squeeze(y_train.values), np.squeeze(y_valid.values)
+        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)
+        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,
+            **kwargs
+        )
+        evals_result["train"] = list(evals_result["train"].values())[0]
+        evals_result["valid"] = list(evals_result["valid"].values())[0]
+
+    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(np.squeeze(x_test.values))), index=x_test.index)