Fix TeamHG-MemexGH-248

sharthZ23 · Dec 7, 2017 · 38ae164 · 38ae164
1 parent eeb3735
commit 38ae164
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/eli5/xgboost.py b/eli5/xgboost.py
@@ -356,18 +356,19 @@ def _parse_tree_dump(text_dump):
 def _parse_dump_line(line):
     # type: (str) -> Tuple[int, Dict[str, Any]]
     branch_match = re.match(
-        '^(\t*)(\d+):\[([^<]+)<([^\]]+)\] '
+        '^(\t*)(\d+):\[([^<]+)<?([^\]]*)\] '
         'yes=(\d+),no=(\d+),missing=(\d+),'
         'gain=([^,]+),cover=(.+)$', line)
     if branch_match:
         tabs, node_id, feature, condition, yes, no, missing, gain, cover = \
             branch_match.groups()
         depth = len(tabs)
+        split_condition = float(condition) if condition else 0.
         return depth, {
             'depth': depth,
             'nodeid': int(node_id),
             'split': feature,
-            'split_condition': float(condition),
+            'split_condition': split_condition,
             'yes': int(yes),
             'no': int(no),
             'missing': int(missing),