From 4aeaf0fb08b65ce593170a7b7eb89bc3c8cf13af Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ricardo=20Montan=CC=83ana?= Date: Wed, 10 Nov 2021 11:12:10 +0100 Subject: [PATCH] Add Odte model --- src/Models.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/src/Models.py b/src/Models.py index 97c0b4e..2c2f4be 100644 --- a/src/Models.py +++ b/src/Models.py @@ -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