Update README.md

README.md

@@ -1,3 +1,104 @@
 ---
 license: mit
+pipeline_tag: tabular-classification
 ---
+
+<h1 align="center"> 🌲 MetaTree 🌲 </h1>
+<p align="center"> <b>Learning a Decision Tree Algorithm with Transformers</b>  (<a href="https://arxiv.org/abs/2402.03774">Zhuang et al. 2024</a>). 
+</p>
+
+<p align="center"> MetaTree is a transformer-based decision tree algorithm. It learns from classical decision tree algorithms (greedy algorithm CART, optimal algorithm GOSDT), for better generalization capabilities.
+</p>
+
+## Quickstart -- use MetaTree to generate decision tree models
+
+Model is avaliable at https://huggingface.co/yzhuang/MetaTree
+
+1. Install `metatreelib`:
+
+```bash
+pip install metatreelib
+# Alternatively,  
+# clone then pip install -e .
+# pip install git+https://github.com/EvanZhuang/MetaTree
+```
+
+2. Use MetaTree on your datasets to generate a decision tree model
+ 
+```python
+from metatree.model_metatree import LlamaForMetaTree as MetaTree
+from metatree.decision_tree_class import DecisionTree, DecisionTreeForest
+from metatree.run_train import preprocess_dimension_patch
+from transformers import AutoConfig
+import imodels # pip install imodels 
+
+# Initialize Model
+model_name_or_path = "yzhuang/MetaTree"
+
+config = AutoConfig.from_pretrained(model_name_or_path)
+model = MetaTree.from_pretrained(
+    model_name_or_path,
+    config=config,
+)   
+
+# Load Datasets
+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 Subsampling
+feature_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 Data
+subset_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 = 2
+outputs = 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) for x in outputs.metatree_dimensions])
+print("Decision Tree Threasholds: ", outputs.tentative_splits)
+```
+
+3. Inference with the decision tree model
+
+```python
+tree_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)
+```
+
+## Example Usage
+
+We show a complete example of using MetaTree at [notebook](examples/example_usage.ipynb)
+
+## Questions?
+
+If you have any questions related to the code or the paper, feel free to reach out to us at [email protected].
+
+
+## Citation
+
+If you find our paper and code useful, please cite us:
+```r
+@misc{zhuang2024learning,
+      title={Learning a Decision Tree Algorithm with Transformers}, 
+      author={Yufan Zhuang and Liyuan Liu and Chandan Singh and Jingbo Shang and Jianfeng Gao},
+      year={2024},
+      eprint={2402.03774},
+      archivePrefix={arXiv},
+      primaryClass={cs.LG}
+}
+```