geoalgo commited on
Commit
04077ef
·
1 Parent(s): ef09dad

add missing families

Browse files
Files changed (1) hide show
  1. main.py +1 -1
main.py CHANGED
@@ -39,7 +39,7 @@ def get_model_family(model_name: str) -> str:
39
  prefixes_mapping = {
40
  Constants.automl: ["AutoGluon"],
41
  Constants.finetuned: ["REALMLP", "TabM", "FASTAI", "MNCA", "NN_TORCH"],
42
- Constants.tree: ["GBM", "CAT", "EBM", "XGB"],
43
  Constants.foundational: ["TABDPT", "TABICL", "TABPFN"],
44
  Constants.baseline: ["KNN", "LR"]
45
  }
 
39
  prefixes_mapping = {
40
  Constants.automl: ["AutoGluon"],
41
  Constants.finetuned: ["REALMLP", "TabM", "FASTAI", "MNCA", "NN_TORCH"],
42
+ Constants.tree: ["GBM", "CAT", "EBM", "XGB", "XT", "RF"],
43
  Constants.foundational: ["TABDPT", "TABICL", "TABPFN"],
44
  Constants.baseline: ["KNN", "LR"]
45
  }