Skip to content

Commit

Permalink
Restore loading model from buffer. (#5360)
Browse files Browse the repository at this point in the history
  • Loading branch information
trivialfis authored Feb 26, 2020
1 parent f2b8cd2 commit 0fd455e
Show file tree
Hide file tree
Showing 2 changed files with 16 additions and 2 deletions.
11 changes: 9 additions & 2 deletions python-package/xgboost/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -1052,7 +1052,7 @@ def __init__(self, params=None, cache=(), model_file=None):
_check_call(
_LIB.XGBoosterUnserializeFromBuffer(self.handle, ptr, length))
self.__dict__.update(state)
elif isinstance(model_file, (STRING_TYPES, os_PathLike)):
elif isinstance(model_file, (STRING_TYPES, os_PathLike, bytearray)):
self.load_model(model_file)
elif model_file is None:
pass
Expand Down Expand Up @@ -1512,7 +1512,8 @@ def save_raw(self):
return ctypes2buffer(cptr, length.value)

def load_model(self, fname):
"""Load the model from a file, local or as URI.
"""Load the model from a file or bytearray. Path to file can be local
or as an URI.
The model is loaded from an XGBoost format which is universal among the
various XGBoost interfaces. Auxiliary attributes of the Python Booster
Expand All @@ -1530,6 +1531,12 @@ def load_model(self, fname):
# from URL.
_check_call(_LIB.XGBoosterLoadModel(
self.handle, c_str(os_fspath(fname))))
elif isinstance(fname, bytearray):
buf = fname
length = c_bst_ulong(len(buf))
ptr = (ctypes.c_char * len(buf)).from_buffer(buf)
_check_call(_LIB.XGBoosterLoadModelFromBuffer(self.handle, ptr,
length))
else:
raise TypeError('Unknown file type: ', fname)

Expand Down
7 changes: 7 additions & 0 deletions tests/python/test_basic_models.py
Original file line number Diff line number Diff line change
Expand Up @@ -300,6 +300,13 @@ def test_model_binary_io(self):
assert float(config['learner']['objective'][
'reg_loss_param']['scale_pos_weight']) == 0.5

buf = bst.save_raw()
from_raw = xgb.Booster()
from_raw.load_model(buf)

buf_from_raw = from_raw.save_raw()
assert buf == buf_from_raw

def test_model_json_io(self):
loc = locale.getpreferredencoding(False)
model_path = 'test_model_json_io.json'
Expand Down

0 comments on commit 0fd455e

Please sign in to comment.