{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Hyrax Getting Started\n",
    "\n",
    "In this getting started notebook we'll create an instance of a Hyrax object, train a builtin model on the CiFAR training dataset, and then use that trained model to run inference on the CiFAR testing dataset."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create a Hyrax instance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[2025-08-25 21:14:19,112 hyrax:INFO] Runtime Config read from: /home/drew/code/hyrax/src/hyrax/hyrax_default_config.toml\n"
     ]
    }
   ],
   "source": [
    "import hyrax\n",
    "\n",
    "h = hyrax.Hyrax()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Update the configuration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "h.config[\"model\"][\"name\"] = \"HyraxAutoencoder\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For this demo, we'll make a few adjustments to the default configuration settings that the `hyrax` object was instantiated with.\n",
    "By accessing the `.config` attribute of the hyrax instance, we can modify any configuration value.\n",
    "There are many configuration values that can be set, but here, we update only the model to train."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train a model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[2025-08-25 21:14:33,430 hyrax.models.model_registry:INFO] Using criterion: torch.nn.CrossEntropyLoss with default arguments.\n",
      "2025-08-25 21:14:33,503 ignite.distributed.auto.auto_dataloader INFO: Use data loader kwargs for dataset '<hyrax.data_sets.hyr': \n",
      "\t{'sampler': <hyrax.pytorch_ignite.SubsetSequentialSampler object at 0x70da50623980>, 'batch_size': 512, 'shuffle': False, 'pin_memory': True}\n",
      "2025-08-25 21:14:33,504 ignite.distributed.auto.auto_dataloader INFO: Use data loader kwargs for dataset '<hyrax.data_sets.hyr': \n",
      "\t{'sampler': <hyrax.pytorch_ignite.SubsetSequentialSampler object at 0x70da5046eea0>, 'batch_size': 512, 'shuffle': False, 'pin_memory': True}\n",
      "/home/drew/miniconda3/envs/hyrax/lib/python3.12/site-packages/ignite/handlers/tqdm_logger.py:127: TqdmExperimentalWarning: Using `tqdm.autonotebook.tqdm` in notebook mode. Use `tqdm.tqdm` instead to force console mode (e.g. in jupyter console)\n",
      "  from tqdm.autonotebook import tqdm\n",
      "2025/08/25 21:14:33 INFO mlflow.system_metrics.system_metrics_monitor: Started monitoring system metrics.\n",
      "[2025-08-25 21:14:33,874 hyrax.pytorch_ignite:INFO] Training model on device: cuda\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7222baf5bfc74a17aeddbfc1f3dcf866",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  2%|1         | 1/59 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "acb5ef1fd60a4856806c805b894e6d4e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  2%|1         | 1/59 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "475afa618df94f47af756ba53740ba50",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  2%|1         | 1/59 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "47bfd5bd82974c8c9a914bf1ae431f7e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  2%|1         | 1/59 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "952a16b49bc744d0b3a981797cbd9a6d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  2%|1         | 1/59 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6e8722711f5f4beebec1067b41872131",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  2%|1         | 1/59 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a59e059d9ad24f68b4695cf9ebcc90a9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  2%|1         | 1/59 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "70d792681c8b45a28590272ce23356d8",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  2%|1         | 1/59 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "35df480d01224d5bbace42cc8aaf73c1",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  2%|1         | 1/59 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "af99379ce2b14fc99e233ca92c5988a3",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  2%|1         | 1/59 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[2025-08-25 21:16:09,062 hyrax.pytorch_ignite:INFO] Total training time: 95.19[s]\n",
      "[2025-08-25 21:16:09,063 hyrax.pytorch_ignite:INFO] Latest checkpoint saved as: /home/drew/code/hyrax/docs/pre_executed/results/20250825-211424-train-0D0a/checkpoint_epoch_10.pt\n",
      "[2025-08-25 21:16:09,063 hyrax.pytorch_ignite:INFO] Best metric checkpoint saved as: /home/drew/code/hyrax/docs/pre_executed/results/20250825-211424-train-0D0a/checkpoint_9_loss=-126.9350.pt\n",
      "2025/08/25 21:16:09 INFO mlflow.system_metrics.system_metrics_monitor: Stopping system metrics monitoring...\n",
      "2025/08/25 21:16:09 INFO mlflow.system_metrics.system_metrics_monitor: Successfully terminated system metrics monitoring!\n",
      "[2025-08-25 21:16:09,079 hyrax.verbs.train:INFO] Finished Training\n",
      "[2025-08-25 21:16:09,364 hyrax.model_exporters:INFO] Exported model to ONNX format: /home/drew/code/hyrax/docs/pre_executed/results/20250825-211424-train-0D0a/example_model_opset_20.onnx\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "HyraxAutoencoder(\n",
       "  (encoder): Sequential(\n",
       "    (0): Conv2d(3, 32, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
       "    (1): GELU(approximate='none')\n",
       "    (2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "    (3): GELU(approximate='none')\n",
       "    (4): Conv2d(32, 64, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
       "    (5): GELU(approximate='none')\n",
       "    (6): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "    (7): GELU(approximate='none')\n",
       "    (8): Conv2d(64, 64, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
       "    (9): GELU(approximate='none')\n",
       "    (10): Flatten(start_dim=1, end_dim=-1)\n",
       "    (11): Linear(in_features=1024, out_features=64, bias=True)\n",
       "  )\n",
       "  (dec_linear): Sequential(\n",
       "    (0): Linear(in_features=64, out_features=1024, bias=True)\n",
       "    (1): GELU(approximate='none')\n",
       "  )\n",
       "  (decoder): Sequential(\n",
       "    (0): ConvTranspose2d(64, 64, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), output_padding=(1, 1))\n",
       "    (1): GELU(approximate='none')\n",
       "    (2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "    (3): GELU(approximate='none')\n",
       "    (4): ConvTranspose2d(64, 32, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), output_padding=(1, 1))\n",
       "    (5): GELU(approximate='none')\n",
       "    (6): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "    (7): GELU(approximate='none')\n",
       "    (8): ConvTranspose2d(32, 3, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), output_padding=(1, 1))\n",
       "    (9): Tanh()\n",
       "  )\n",
       "  (criterion): CrossEntropyLoss()\n",
       ")"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "h.train()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The output of the training will be stored in a time-stamped directory under the `./results/`.\n",
    "By default, a copy of the final configuration used in training is persisted as `runtime_config.toml`.\n",
    "To train again with the same configuration, you can reference this runtime_config.toml file.\n",
    "\n",
    "If running in another notebook, instantiate a hyrax object like so:\n",
    "```\n",
    "new_hyrax_instance = hyrax.Hyrax(config_file='./results/<timestamped_directory>/runtime_config.toml')\n",
    "```\n",
    "\n",
    "Or from the command line:\n",
    "```\n",
    ">> hyrax train --runtime-config ./results/<timestamped_directory>/runtime_config.toml\n",
    "```\n",
    "\n",
    "Note here we're training on only a small handful of CiFAR data, but Hyrax has demonstrated that it can scale up to training sets with >1M samples."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[2025-08-25 21:16:17,243 hyrax.models.model_registry:INFO] Using criterion: torch.nn.CrossEntropyLoss with default arguments.\n",
      "[2025-08-25 21:16:17,244 hyrax.verbs.infer:INFO] data set has length 50000\n",
      "2025-08-25 21:16:17,246 ignite.distributed.auto.auto_dataloader INFO: Use data loader kwargs for dataset '<hyrax.data_sets.hyr': \n",
      "\t{'sampler': None, 'batch_size': 128, 'shuffle': False, 'pin_memory': True}\n",
      "[2025-08-25 21:16:17,263 hyrax.verbs.infer:INFO] Saving inference results at: /home/drew/code/hyrax/docs/pre_executed/results/20250825-211609-infer-SY49\n",
      "[2025-08-25 21:16:17,602 hyrax.pytorch_ignite:INFO] Evaluating model on device: cuda\n",
      "[2025-08-25 21:16:17,605 hyrax.pytorch_ignite:INFO] Total epochs: 1\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b214578af62349308fadc484cc9a359d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 1/391 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[2025-08-25 21:16:30,310 hyrax.pytorch_ignite:INFO] Total evaluation time: 12.71[s]\n",
      "[2025-08-25 21:16:30,434 hyrax.verbs.infer:INFO] Inference Complete.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<hyrax.data_sets.inference_dataset.InferenceDataSet at 0x70db7a0bafc0>"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "h.config[\"data_set\"][\"test_size\"] = 1.0\n",
    "h.config[\"data_set\"][\"train_size\"] = 0.0\n",
    "h.config[\"data_set\"][\"validate_size\"] = 0.0\n",
    "h.config[\"data_loader\"][\"batch_size\"] = 128\n",
    "\n",
    "h.infer()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Once a model has been trained, we can use the model weights file to run inference.\n",
    "By default running `infer` will look for the latest available model weights file.\n",
    "A specific model weights file can be specified with `h.config['infer']['model_weights_file'] = <path_to_model_weights_file>`.\n",
    "\n",
    "Here we'll make use of the last trained model weights file, and update the data set splits so that 100% of the data will be used for inference.\n",
    "\n",
    "With the configuration updated, we can run inference by calling `h.infer()`.\n",
    "\n",
    "The results of running inference are saved in the output directory noted in the last log line.\n",
    "The default output format is batched .npy files.\n",
    "Additionally a ChromaDB vector database will be populated with the inference results to enable efficient similarity search."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "fibad",
   "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.12.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
