1
0
mirror of https://github.com/microsoft/qlib.git synced 2026-07-15 16:56:54 +08:00

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

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

View File

@@ -4,7 +4,7 @@
import numpy as np
import pandas as pd
import lightgbm as lgb
from typing import Text, Union
from ...model.base import ModelFT
from ...data.dataset import DatasetH
from ...data.dataset.handler import DataHandlerLP
@@ -61,7 +61,7 @@ class LGBModel(ModelFT):
evals_result["train"] = list(evals_result["train"].values())[0]
evals_result["valid"] = list(evals_result["valid"].values())[0]
def predict(self, dataset, segment="test"):
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(segment, col_set="feature", data_key=DataHandlerLP.DK_I)