Package: metatreelib
🌲 MetaTree 🌲Learning a Decision Tree Algorithm with TransformersMetaTree is a transformer-based decision tree algorithm. It learns from classical decision tree algorithms (greedy algorithm CART, optimal algorithm GOSDT), for better generalization capabilities.Quickstart -- use MetaTree to generate decision tree modelsModel is avaliable athttps://huggingface.co/yzhuang/MetaTreeInstallmetatreelib:pipinstallmetatreelib# Alternatively,# clone then pip install -e .# pip install git+https://github.com/EvanZhuang/MetaTreeUse MetaTree on your datasets to generate a decision tree modelfrommetatree.model_metatreeimportLlamaForMetaTreeasMetaTreefrommetatree.decision_tree_classimportDecisionTree,DecisionTreeForestfrommetatree.run_trainimportpreprocess_dimension_patchfromtransformersimportAutoConfigimportimodels# pip install imodels# Initialize Modelmodel_name_or_path="yzhuang/MetaTree"config=AutoConfig.from_pretrained(model_name_or_path)model=MetaTree.from_pretrained(model_name_or_path,config=config,)X,y,feature_names=imodels.get_clean_dataset('fico',data_source='imodels')print("Dataset Shapes X={}, y={}, Num of Classes={}".format(X.shape,y.shape,len(set(y))))train_idx,test_idx=sklearn.model_selection.train_test_split(range(X.shape[0]),test_size=0.3,random_state=seed)# Dimension Subsamplingfeature_idx=np.random.choice(X.shape[1],10,replace=False)X=X[:,feature_idx]test_X,test_y=X[test_idx],y[test_idx]# Sample Train and Test Datasubset_idx=random.sample(train_idx,256)train_X,train_y=X[subset_idx],y[subset_idx]input_x=torch.tensor(train_X,dtype=torch.float32)input_y=torch.nn.functional.one_hot(torch.tensor(train_y)).float()batch={"input_x":input_x,"input_y":input_y,"input_y_clean":input_y}batch=preprocess_dimension_patch(batch,n_feature=10,n_class=10)model.depth=2outputs=model.generate_decision_tree(batch['input_x'],batch['input_y'],depth=model.depth)decision_tree_forest.add_tree(DecisionTree(auto_dims=outputs.metatree_dimensions,auto_thresholds=outputs.tentative_splits,input_x=batch['input_x'],input_y=batch['input_y'],depth=model.depth))print("Decision Tree Features: ",[x.argmax(dim=-1)forxinoutputs.metatree_dimensions])print("Decision Tree Threasholds: ",outputs.tentative_splits)Inference with the decision tree modeltree_pred=decision_tree_forest.predict(torch.tensor(test_X,dtype=torch.float32))accuracy=accuracy_score(test_y,tree_pred.argmax(dim=-1).squeeze(0))print("MetaTree Test Accuracy: ",accuracy)To reproduce the evaluation resultsDatasets used in evaluation are avaliable on huggingfacefordataset_nameinmetatree_mfeat_fouriermetatree_mfeat_zernikemetatree_mfeat_morphologicalmetatree_mfeat_karhunenmetatree_page_blocksmetatree_optdigitsmetatree_pendigitsmetatree_waveform_5000metatree_Hyperplane_10_1E_3metatree_Hyperplane_10_1E_4metatree_pokerhandmetatree_RandomRBF_0_0metatree_RandomRBF_10_1E_3metatree_RandomRBF_50_1E_3metatree_RandomRBF_10_1E_4metatree_RandomRBF_50_1E_4metatree_SEA_50_metatree_SEA_50000_metatree_satimagemetatree_BNG_labor_metatree_BNG_breast_w_metatree_BNG_mfeat_karhunen_metatree_BNG_bridges_version1_metatree_BNG_mfeat_zernike_metatree_BNG_cmc_metatree_BNG_colic_ORIG_metatree_BNG_colic_metatree_BNG_credit_a_metatree_BNG_page_blocks_metatree_BNG_credit_g_metatree_BNG_pendigits_metatree_BNG_cylinder_bands_metatree_BNG_dermatology_metatree_BNG_sonar_metatree_BNG_glass_metatree_BNG_heart_c_metatree_BNG_heart_statlog_metatree_BNG_vehicle_metatree_BNG_hepatitis_metatree_BNG_waveform_5000_metatree_BNG_zoo_metatree_vehicle_sensITmetatree_UNIX_user_datametatree_fri_c3_1000_25metatree_rmftsa_sleepdatametatree_JapaneseVowelsmetatree_fri_c4_1000_100metatree_abalonemetatree_fri_c4_1000_25metatree_bank8FMmetatree_analcatdata_suprememetatree_aileronsmetatree_cpu_smallmetatree_space_gametatree_fri_c1_1000_5metatree_puma32Hmetatree_fri_c3_1000_10metatree_cpu_actmetatree_fri_c4_1000_10metatree_quakemetatree_fri_c4_1000_50metatree_fri_c0_1000_5metatree_delta_aileronsmetatree_fri_c3_1000_50metatree_kin8nmmetatree_fri_c3_1000_5metatree_puma8NHmetatree_delta_elevatorsmetatree_housesmetatree_bank32nhmetatree_fri_c1_1000_50metatree_house_8Lmetatree_fri_c0_1000_10metatree_elevatorsmetatree_windmetatree_fri_c0_1000_25metatree_fri_c2_1000_50metatree_pollenmetatree_mvmetatree_friedmetatree_fri_c2_1000_25metatree_fri_c0_1000_50metatree_fri_c1_1000_10metatree_fri_c2_1000_5metatree_fri_c2_1000_10metatree_fri_c1_1000_25metatree_visualizing_soilmetatree_socmobmetatree_mozilla4metatree_pc3metatree_pc1dofortree_sizein15102030405060708090100doPYTHONPATH="./metatree":"$PYTHONPATH"python./metatree/eval_generalization.py\--model_name_or_path"yzhuang/MetaTree"\--normalize\--dataset_nameyzhuang/$dataset_name\--max_train_steps$tree_size\--backward_window1\--n_feature10\--n_class10\--block_size256\--depth2\--per_device_train_batch_size256\--per_device_eval_batch_size512\--preprocessing_num_workers8\--inference_only\#--with_tracking \#--report_to "wandb" # In case you want to use wandb to log the resultsdonedone