Skip to content

A
alpha-mind
Commit 4dcce5d3 authored by wegamekinglc's avatar wegamekinglc
Browse files
Options

update loader

parent f57d5587
Show whitespace changes
Inline Side-by-side
Showing with 3 additions and 0 deletions
alphamind/model/loader.py
View file @ 4dcce5d3
......@@ -15,6 +15,7 @@ from alphamind.model.treemodel import RandomForestClassifier
from alphamind.model.treemodel import XGBRegressor
from alphamind.model.treemodel import XGBClassifier
from alphamind.model.treemodel import XGBTrainer
from alphamind.model.svm import NvSVRModel
def load_model(model_desc: dict) -> ModelBase:
......@@ -40,5 +41,7 @@ def load_model(model_desc: dict) -> ModelBase:
return XGBClassifier.load(model_desc)
elif 'XGBTrainer' in model_name_parts:
return XGBTrainer.load(model_desc)
elif 'NvSVR' in model_name_parts:
return NvSVRModel.load(model_desc)
else:
raise ValueError('{0} is not currently supported in model loader.'.format(model_name))
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment