0.1.13

erdogant · Aug 28, 2020 · 5e75854 · 5e75854
1 parent 53301f9
commit 5e75854
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/treeplot/__init__.py b/treeplot/__init__.py
@@ -8,7 +8,7 @@
 
 __author__ = 'Erdogan Tasksen'
 __email__ = 'erdogant@gmail.com'
-__version__ = '0.1.12'
+__version__ = '0.1.13'
 
 # module level doc-string
 __doc__ = """

diff --git a/treeplot/treeplot.py b/treeplot/treeplot.py
@@ -50,15 +50,15 @@ def plot(model, featnames=None, num_trees=None, plottype='horizontal', figsize=(
         Figure axis of the input model.
 
     """
-    modelname = str(model).lower()
+    modelname = str(model).lower()[0:30]
     if ('xgb' in modelname):
         if verbose>=4: print('xgboost plotting pipeline.')
         ax = xgboost(model, featnames=featnames, num_trees=num_trees, figsize=figsize, plottype=plottype, verbose=verbose)
+    elif ('lgb' in modelname):
+        ax = lgbm(model, featnames=featnames, num_trees=num_trees, figsize=figsize, verbose=verbose)
     elif ('tree' in modelname) or ('forest' in modelname) or ('gradientboosting' in modelname):
         if verbose>=4: print('tree plotting pipeline.')
         ax = randomforest(model, featnames=featnames, num_trees=num_trees, figsize=figsize, verbose=verbose)
-    elif ('lgb' in modelname):
-        ax = lgbm(model, featnames=featnames, num_trees=num_trees, figsize=figsize, verbose=verbose)
     else:
         print('[treeplot] >Model not recognized: %s' %(modelname))
         ax = None