Add Odte model

This commit is contained in:
2021-11-10 11:12:10 +01:00
parent d55546e56b
commit 4aeaf0fb08

View File

@@ -2,6 +2,7 @@ from sklearn.tree import DecisionTreeClassifier, ExtraTreeClassifier
from sklearn.svm import SVC
from stree import Stree
from wodt import TreeClassifier
from odte import Odte
class Models:
@@ -17,10 +18,14 @@ class Models:
return TreeClassifier
elif name == "SVC":
return SVC
elif name == "ODTE":
return Odte
else:
msg = f"No model recognized {name}"
if name == "Stree" or name == "stree":
msg += ", did you mean STree?"
elif name == "odte" or name == "Odte":
msg += ", did you mean ODTE?"
raise ValueError(msg)
@staticmethod