examples/example_usage.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Example Usage of MetaTree"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/yufan/anaconda3/envs/gosdt/lib/python3.9/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n",
      "/home/yufan/anaconda3/envs/gosdt/lib/python3.9/site-packages/transformers/deepspeed.py:23: FutureWarning: transformers.deepspeed module is deprecated and will be removed in a future version. Please import deepspeed modules directly from transformers.integrations\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "import sys\n",
    "sys.path.append('..')\n",
    "\n",
    "from metatree.model_metatree import LlamaForMetaTree as MetaTree\n",
    "from metatree.decision_tree_class import DecisionTree, DecisionTreeForest\n",
    "from metatree.run_train import preprocess_dimension_patch\n",
    "from transformers import AutoConfig\n",
    "\n",
    "from sklearn.metrics import accuracy_score\n",
    "import sklearn\n",
    "\n",
    "import torch\n",
    "import numpy as np\n",
    "import seaborn as sns\n",
    "import matplotlib.pyplot as plt\n",
    "import imodels\n",
    "import random\n",
    "\n",
    "model_name_or_path = \"yzhuang/MetaTree\"\n",
    "\n",
    "config = AutoConfig.from_pretrained(model_name_or_path)\n",
    "model = MetaTree.from_pretrained(\n",
    "    model_name_or_path,\n",
    "    config=config,\n",
    ")   \n",
    "decision_tree_forest = DecisionTreeForest()\n",
    "\n",
    "ensemble_size = 1\n",
    "seed = 42"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Generate the decision tree with MetaTree\n",
    "\n",
    "We use a dataset from imodels (https://github.com/csinva/imodels) as our example, to install imodels, run *pip install imodels*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "fetching fico from imodels\n",
      "Dataset Shapes X=(10459, 23), y=(10459,), Num of Classes=2\n",
      "Decision Tree Features:  [tensor([4]), tensor([8]), tensor([7])]\n",
      "Decision Tree Threasholds:  [tensor([[72.5000]]), tensor([[64.5000]]), tensor([[97.5000]])]\n"
     ]
    }
   ],
   "source": [
    "X, y, feature_names = imodels.get_clean_dataset('fico', data_source='imodels')\n",
    "\n",
    "print(\"Dataset Shapes X={}, y={}, Num of Classes={}\".format(X.shape, y.shape, len(set(y))))\n",
    "\n",
    "train_idx, test_idx = sklearn.model_selection.train_test_split(range(X.shape[0]), test_size=0.3, random_state=seed)\n",
    "\n",
    "# Dimension \n",
    "random.seed(seed)\n",
    "feature_idx = np.random.choice(X.shape[1], 10, replace=False)\n",
    "X = X[:, feature_idx]\n",
    "\n",
    "test_X, test_y = X[test_idx], y[test_idx]\n",
    "\n",
    "for i in range(ensemble_size):\n",
    "    # Sample Train and Test Data\n",
    "    random.seed(seed+i+1)\n",
    "    subset_idx = random.sample(train_idx, 256)\n",
    "    train_X, train_y = X[subset_idx], y[subset_idx]\n",
    "    \n",
    "\n",
    "    input_x = torch.tensor(train_X, dtype=torch.float32)\n",
    "    input_y = torch.nn.functional.one_hot(torch.tensor(train_y)).float()\n",
    "\n",
    "    batch = {\"input_x\": input_x, \"input_y\": input_y, \"input_y_clean\": input_y}\n",
    "    batch = preprocess_dimension_patch(batch, n_feature=10, n_class=10)\n",
    "    model.depth = 2\n",
    "    outputs = model.generate_decision_tree(batch['input_x'], batch['input_y'], depth=model.depth)\n",
    "    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))\n",
    "\n",
    "    print(\"Decision Tree Features: \", [x.argmax(dim=-1) for x in outputs.metatree_dimensions])\n",
    "    print(\"Decision Tree Threasholds: \", outputs.tentative_splits)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Evaluation on the test set"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "MetaTree Test Accuracy:  0.6912045889101338\n"
     ]
    }
   ],
   "source": [
    "tree_pred = decision_tree_forest.predict(torch.tensor(test_X, dtype=torch.float32))\n",
    "\n",
    "accuracy = accuracy_score(test_y, tree_pred.argmax(dim=-1).squeeze(0))\n",
    "print(\"MetaTree Test Accuracy: \", accuracy)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Fitting CART tree as a comparison"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CART Test Accuracy:  0.6883365200764818\n"
     ]
    }
   ],
   "source": [
    "cart_ensemble = []\n",
    "\n",
    "for i in range(ensemble_size):\n",
    "    random.seed(seed+i+1)\n",
    "    subset_idx = random.sample(train_idx, 256)\n",
    "    train_X, train_y = X[subset_idx], y[subset_idx]\n",
    "    \n",
    "    clf = sklearn.tree.DecisionTreeClassifier(max_depth=2, random_state=seed+i+1)\n",
    "    clf.fit(train_X, train_y)\n",
    "    cart_ensemble.append(clf)\n",
    "\n",
    "overall_pred = np.zeros((test_X.shape[0], len(set(test_y))))\n",
    "for clf in cart_ensemble:\n",
    "    overall_pred += clf.predict_proba(test_X)\n",
    "overall_pred = overall_pred / len(cart_ensemble)\n",
    "\n",
    "accuracy = accuracy_score(test_y, overall_pred.argmax(axis=-1))\n",
    "print(\"CART Test Accuracy: \", accuracy)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "gosdt",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.17"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}