liuyu/qlib
1
0
Fork 0
mirror of https://github.com/microsoft/qlib.git synced 2026-07-13 07:46:53 +08:00
Code Activity

Collect all contrib models in __init__ and add unit tests for init

Browse Source
This commit is contained in:
D-X-Y
2021-03-28 10:39:28 +00:00
parent 8a2e7b62af
commit 0386df7b16
17 changed files with 115 additions and 46 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
qlib/contrib/model/catboost_model.py
View File

@@ -3,6 +3,7 @@
import numpy as np
import pandas as pd
from typing import Text, Union
from catboost import Pool, CatBoost
from catboost.utils import get_gpu_device_count
@@ -62,10 +63,10 @@ class CatBoostModel(Model):
evals_result["train"] = list(evals_result["learn"].values())[0]
evals_result["valid"] = list(evals_result["validation"].values())[0]
def predict(self, dataset):
def predict(self, dataset: DatasetH, segment: Union[Text, slice] = "test"):
if self.model is None:
raise ValueError("model is not fitted yet!")
x_test = dataset.prepare("test", col_set="feature")
x_test = dataset.prepare(segment, col_set="feature", data_key=DataHandlerLP.DK_I)
return pd.Series(self.model.predict(x_test.values), index=x_test.index)