{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.utils.data import TensorDataset, random_split\n",
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## prep data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def prep_iris_data(\n",
    "\tpath: str = \"../demos/training_data/iris.csv\",\n",
    "\tp_test: float = 0.2,\n",
    ") -> dict:\n",
    "\t# Load data\n",
    "\tiris_df: pd.DataFrame = pd.read_csv(path)\n",
    "\n",
    "\t# Create a mapping from integer labels to species names\n",
    "\ttarget_list: list[str] = list(iris_df[\"species\"].unique())\n",
    "\ttarget_map: dict[str, int] = {name: i for i, name in enumerate(target_list)}\n",
    "\n",
    "\t# Convert species to integer labels\n",
    "\tiris_df[\"species\"] = iris_df[\"species\"].map(target_map)\n",
    "\n",
    "\t# Create tensors\n",
    "\tinputs: torch.Tensor = torch.tensor(\n",
    "\t\tiris_df.iloc[:, :-1].values, dtype=torch.float32\n",
    "\t)\n",
    "\ttargets: torch.Tensor = torch.tensor(iris_df[\"species\"].values, dtype=torch.int64)\n",
    "\n",
    "\t# Create a TensorDataset\n",
    "\tdataset: TensorDataset = TensorDataset(inputs, targets)\n",
    "\n",
    "\t# Determine train and test sizes\n",
    "\ttotal_size: int = len(dataset)\n",
    "\ttest_size: int = int(p_test * total_size)\n",
    "\ttrain_size: int = total_size - test_size\n",
    "\n",
    "\t# Split dataset into train and test sets\n",
    "\ttrain_dataset, test_dataset = random_split(dataset, [train_size, test_size])\n",
    "\n",
    "\t# Prepare the final dictionary\n",
    "\tresult = {\n",
    "\t\t\"train\": train_dataset,\n",
    "\t\t\"test\": test_dataset,\n",
    "\t\t\"column_names\": iris_df.columns[:-1].tolist(),\n",
    "\t\t\"target_list\": target_list,\n",
    "\t\t\"target_map\": target_map,\n",
    "\t}\n",
    "\n",
    "\treturn result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'train': <torch.utils.data.dataset.Subset object at 0x0000026BB6A83740>, 'test': <torch.utils.data.dataset.Subset object at 0x0000026BB45C1460>, 'column_names': ['sepal_length', 'sepal_width', 'petal_length', 'petal_width'], 'target_list': ['setosa', 'versicolor', 'virginica'], 'target_map': {'setosa': 0, 'versicolor': 1, 'virginica': 2}}\n",
      "(tensor([5.9000, 3.0000, 4.2000, 1.5000]), tensor(1))\n",
      "(tensor([5.0000, 3.5000, 1.6000, 0.6000]), tensor(0))\n"
     ]
    }
   ],
   "source": [
    "IRIS_DATA: dict = prep_iris_data()\n",
    "print(IRIS_DATA)\n",
    "print(IRIS_DATA[\"train\"][0])\n",
    "print(IRIS_DATA[\"test\"][0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## define DNN"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DNN(nn.Module):\n",
    "\tdef __init__(self, d_input: int, d_hidden: int, d_output: int):\n",
    "\t\tsuper().__init__()\n",
    "\t\tself.net: nn.Module = nn.Sequential(\n",
    "\t\t\tnn.Linear(d_input, d_hidden),\n",
    "\t\t\tnn.ReLU(),\n",
    "\t\t\tnn.Linear(d_hidden, d_output),\n",
    "\t\t\tnn.Softmax(dim=1),\n",
    "\t\t)\n",
    "\n",
    "\tdef forward(self, x: torch.Tensor) -> torch.Tensor:\n",
    "\t\treturn self.net(x)\n",
    "\n",
    "\tdef predict(self, x: torch.Tensor) -> torch.Tensor:\n",
    "\t\treturn torch.argmax(self.forward(x), dim=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Init model & dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize the DNN model\n",
    "model: DNN = DNN(\n",
    "\td_input=4,\n",
    "\td_hidden=10,\n",
    "\td_output=3,\n",
    ")\n",
    "model.train()\n",
    "\n",
    "# Loss function and optimizer\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=0.001)\n",
    "\n",
    "# Prepare the Iris dataset\n",
    "iris_data = prep_iris_data()  # Assuming this calls the function we defined earlier\n",
    "TRAIN_LOADER = torch.utils.data.DataLoader(\n",
    "\tiris_data[\"train\"], batch_size=32, shuffle=True\n",
    ")\n",
    "TEST_LOADER = torch.utils.data.DataLoader(\n",
    "\tiris_data[\"test\"], batch_size=32, shuffle=False\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# init logger"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "from trnbl.loggers.local import LocalLogger\n",
    "\n",
    "# from trnbl.loggers.wandb import WandbLogger\n",
    "from trnbl.training_manager import TrainingManager"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "# starting logger with id hf4264-250116_2136-daqujo\n"
     ]
    }
   ],
   "source": [
    "train_config = dict(\n",
    "\tmodel=str(model),\n",
    "\tdataset=\"iris\",\n",
    "\toptimizer=str(optimizer),\n",
    "\tcriterion=str(criterion),\n",
    ")\n",
    "\n",
    "logger = LocalLogger(\n",
    "\tproject=\"iris-demo\",\n",
    "\tmetric_names=[\"train/loss\", \"train/acc\", \"val/loss\", \"val/acc\"],\n",
    "\ttrain_config=train_config,\n",
    "\tbase_path=\"../demos/local\",\n",
    ")\n",
    "\n",
    "# alternate loggers:\n",
    "# logger_tb = TensorBoardLogger(log_dir=\"demos/tensorboard\")\n",
    "# logger_wandb = WandbLogger.create(train_config, project=\"iris-demo\", job_type=\"train\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# define evaluation function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "def eval_func(model):\n",
    "\tlosses = torch.full((len(TEST_LOADER),), fill_value=torch.nan)\n",
    "\taccuracies = torch.full((len(TEST_LOADER),), fill_value=torch.nan)\n",
    "\tfor idx, (inputs, targets) in enumerate(TEST_LOADER):\n",
    "\t\toutputs = model(inputs)\n",
    "\t\tloss = criterion(outputs, targets)\n",
    "\t\taccuracy = torch.sum(torch.argmax(outputs, dim=1) == targets).item() / len(\n",
    "\t\t\ttargets\n",
    "\t\t)\n",
    "\t\tlosses[idx] = loss.item()\n",
    "\t\taccuracies[idx] = accuracy\n",
    "\n",
    "\treturn {\n",
    "\t\t\"val/loss\": losses.mean().item(),\n",
    "\t\t\"val/acc\": accuracies.mean().item(),\n",
    "\t}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# run training loop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "# starting training manager initialization\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "training run:   0%|          | 0/120 [00:00<?, ? epochs/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "# initialized training manager\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "training run: 100%|██████████| 120/120 [00:41<00:00,  2.87 epochs/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "# training complete\n",
      "# closing logger\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "with TrainingManager(\n",
    "\tmodel=model,\n",
    "\tlogger=logger,\n",
    "\tevals={\n",
    "\t\t\"1 epochs\": eval_func,\n",
    "\t\t\"1/2 epochs\": lambda model: logger.get_mem_usage(),\n",
    "\t}.items(),\n",
    "\tcheckpoint_interval=\"50 epochs\",\n",
    ") as tr:\n",
    "\t# Training loop\n",
    "\tfor epoch in tr.epoch_loop(range(120)):\n",
    "\t\tfor inputs, targets in tr.batch_loop(TRAIN_LOADER):\n",
    "\t\t\toptimizer.zero_grad()\n",
    "\t\t\toutputs = model(inputs)\n",
    "\t\t\tloss = criterion(outputs, targets)\n",
    "\t\t\tloss.backward()\n",
    "\t\t\toptimizer.step()\n",
    "\n",
    "\t\t\taccuracy = torch.sum(torch.argmax(outputs, dim=1) == targets).item() / len(\n",
    "\t\t\t\ttargets\n",
    "\t\t\t)\n",
    "\n",
    "\t\t\ttr.batch_update(\n",
    "\t\t\t\tsamples=len(targets),\n",
    "\t\t\t\t**{\"train/loss\": loss.item(), \"train/acc\": accuracy},\n",
    "\t\t\t)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.0"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
