{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "5872c787",
   "metadata": {},
   "source": [
    "# Providing Data - Level 2\n",
    "## Fine grained control over requested data\n",
    "\n",
    "In this notebook we’ll explore a few more options you can put in ``model_inputs`` to control exactly what data your model receives.\n",
    "\n",
    "We’ll continue with the built-in ``HyraxRandomDataset`` from Level 1 to demonstrate passing dataset-specific options via ``dataset_config``.\n",
    "After that we’ll switch to ``HyraxCifar`` to show how to use ``data_location`` and ``fields`` so you can request only the pieces you need.\n",
    "\n",
    "As always, start by creating an instance of the Hyrax class, then preparing the same dataset that was used at the end of Level 1."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "744b9f7f",
   "metadata": {},
   "outputs": [],
   "source": [
    "from hyrax import Hyrax\n",
    "\n",
    "h = Hyrax()\n",
    "\n",
    "model_inputs_definition = {\n",
    "    \"train\": {\n",
    "        \"data\": {\n",
    "            \"dataset_class\": \"HyraxRandomDataset\",\n",
    "            \"primary_id_field\": \"object_id\",\n",
    "        },\n",
    "    },\n",
    "}\n",
    "\n",
    "h.set_config(\"model_inputs\", model_inputs_definition)\n",
    "\n",
    "# Prepare \"model_inputs\"\n",
    "d = h.prepare()\n",
    "\n",
    "# Print a sample of the data\n",
    "d[\"train\"].sample_data()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4a8e9601",
   "metadata": {},
   "source": [
    "## Dataset-specific configurations\n",
    "\n",
    "The HyraxRandomDataset is very flexible, and has several configuration settings.\n",
    "Let's see how we can set the shape of the returned \"image\" using \"model_inputs\"."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23be8128",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_inputs_definition = {\n",
    "    \"train\": {\n",
    "        \"data\": {\n",
    "            \"dataset_class\": \"HyraxRandomDataset\",\n",
    "            \"primary_id_field\": \"object_id\",\n",
    "            \"dataset_config\": {\n",
    "                \"shape\": (3, 2, 4),  # <- Change the shape\n",
    "            },\n",
    "        },\n",
    "    }\n",
    "}\n",
    "\n",
    "h.set_config(\"model_inputs\", model_inputs_definition)\n",
    "\n",
    "d = h.prepare()\n",
    "\n",
    "sample = d[\"train\"].sample_data()\n",
    "print(f\"Image shape: {sample['data']['image'].shape}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4d2b3d68",
   "metadata": {},
   "source": [
    "Any other dataset configuration parameters can be set in the `\"dataset_config\"` dictionary.\n",
    "While this may seem redundant (since you can set these same values elsewhere in the config)\n",
    "the real power of this show in Providing Data - Level 3 when we request data from multiple datasets at once."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4db3ddf2",
   "metadata": {},
   "source": [
    "## Defining the data location\n",
    "So far, the built-in HyraxRandomDataset has been used as a lightweight stand‑in for real data — it lives only in memory and is great for quick experiments.\n",
    "Now we'll switch to the HyraxCifarDataset and learn how to tell Hyrax where the data lives using the \"data_location\" parameter. This allows Hyrax to load examples from disk, instead of generating them on the fly, which is what you’ll do in real projects.\n",
    "\n",
    "Note - We'll need to download the CiFAR dataset for this, it's about 170Mb."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "81aeb1cf",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_inputs_definition = {\n",
    "    \"train\": {\n",
    "        \"data\": {\n",
    "            \"dataset_class\": \"HyraxCifarDataSet\",\n",
    "            \"data_location\": \"./data\",  # <- Define where to find the data\n",
    "            \"primary_id_field\": \"object_id\",\n",
    "        },\n",
    "    },\n",
    "}\n",
    "\n",
    "h.set_config(\"model_inputs\", model_inputs_definition)\n",
    "\n",
    "d = h.prepare()\n",
    "\n",
    "d[\"train\"].sample_data()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cbb28731",
   "metadata": {},
   "source": [
    "## If you only want _some_ fields\n",
    "\n",
    "In the minimal setup, Hyrax grabs all the fields that a dataset can provide.\n",
    "That’s handy for a quick demo, but in real projects you usually don’t want everything.\n",
    "For example, you might only need images and labels, while ignoring extra metadata.\n",
    "\n",
    "Fortunately, you’re in control. The ``DataProvider`` can show you exactly which fields are available,\n",
    "and the ``model_inputs`` configuration lets you pick and choose the ones you actually want."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "173e4687",
   "metadata": {},
   "outputs": [],
   "source": [
    "d[\"train\"].fields()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "93954670",
   "metadata": {},
   "source": [
    "Note that `fields()` will always return all the fields that the dataset exposes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5abd1a0d",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_inputs_definition = {\n",
    "    \"train\": {\n",
    "        \"data\": {\n",
    "            \"dataset_class\": \"HyraxCifarDataSet\",\n",
    "            \"data_location\": \"./data\",  # <- Define where to find the data\n",
    "            \"primary_id_field\": \"object_id\",\n",
    "            \"fields\": [\"label\"],  # <- Request only this specific field.\n",
    "        },\n",
    "    },\n",
    "}\n",
    "\n",
    "h.set_config(\"model_inputs\", model_inputs_definition)\n",
    "\n",
    "d = h.prepare()\n",
    "\n",
    "d[\"train\"].sample_data()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0e44a224",
   "metadata": {},
   "source": [
    "Huzzah! 🎉 We just looked at all the fields that ``HyraxCifarDataSet`` can provide,\n",
    "and then updated ``model_inputs`` to request only \"label\".\n",
    "Now our ``DataProvider`` returns exactly the fields we want — no extras."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9ca1c096",
   "metadata": {},
   "source": [
    "## Recap\n",
    "\n",
    "In this notebook, we explored how to control the data provided to your models using Hyrax's flexible configuration system.\n",
    "You learned how to specify exactly what data you want, where it comes from, and which fields are included, giving you fine-grained control over your data pipeline.\n",
    "In this notebook, we:\n",
    "\n",
    "- Used `HyraxRandomDataset` and set dataset-specific options via `dataset_config`\n",
    "- Switched to `HyraxCifarDataSet` and specified the data location with `data_location`\n",
    "- Inspected available fields with `fields()`\n",
    "- Selected only the fields needed for your workflow using the `fields` parameter in `model_inputs`\n",
    "\n",
    "Next, work through the \"Providing Data - Level 3\" notebook to learn how to use `model_inputs` and `DataProvider` to specify dataset splits, provide data for inference and request multimodal data for your models!"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "hyrax",
   "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.13.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
