File size: 673 Bytes
bc45d7d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 | from dataclasses import asdict
import pickle
import xgboost
class Xgboost:
def __init__(self, config):
config = asdict(config)
self.model = xgboost.XGBClassifier(use_label_encoder=False, **config)
def __call__(self, x):
return self.model.predict(x)
def fit(self, x_train, y_train, x_val, y_val):
self.model.fit(
x_train,
y_train,
eval_set=[(x_train, y_train), (x_val, y_val)],
verbose=True,
)
def save(self, save_path):
pickle.dump(self.model, open(save_path, "wb"))
def load(self, load_path):
self.model = pickle.load(open(load_path, "rb"))
|