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 ...@@ -15,6 +15,7 @@ from alphamind.model.treemodel import RandomForestClassifier
from alphamind.model.treemodel import XGBRegressor from alphamind.model.treemodel import XGBRegressor
from alphamind.model.treemodel import XGBClassifier from alphamind.model.treemodel import XGBClassifier
from alphamind.model.treemodel import XGBTrainer from alphamind.model.treemodel import XGBTrainer
from alphamind.model.svm import NvSVRModel
def load_model(model_desc: dict) -> ModelBase: def load_model(model_desc: dict) -> ModelBase:
...@@ -40,5 +41,7 @@ def load_model(model_desc: dict) -> ModelBase: ...@@ -40,5 +41,7 @@ def load_model(model_desc: dict) -> ModelBase:
return XGBClassifier.load(model_desc) return XGBClassifier.load(model_desc)
elif 'XGBTrainer' in model_name_parts: elif 'XGBTrainer' in model_name_parts:
return XGBTrainer.load(model_desc) return XGBTrainer.load(model_desc)
elif 'NvSVR' in model_name_parts:
return NvSVRModel.load(model_desc)
else: else:
raise ValueError('{0} is not currently supported in model loader.'.format(model_name)) 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