{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# MNIST Example\n",
    "A classic example of MNIST digit classification using `phitodeep`:\n",
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from datasets import load_dataset\n",
    "\n",
    "from phitodeep.model import SequentialBuilder\n",
    "from phitodeep.loss import CategoricalCrossEntropy\n",
    "from phitodeep.optimization.optimizers import Adam\n",
    "from phitodeep.optimization.initialization import Xavier, InitType\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load MNIST dataset\n",
    "Using the `datasets` library from Hugging Face, we can easily load the MNIST dataset:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Warning: You are sending unauthenticated requests to the HF Hub. Please set a HF_TOKEN to enable higher rate limits and faster downloads.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(60000, 28, 28) (60000,)\n"
     ]
    }
   ],
   "source": [
    "train_dataset = load_dataset(\"ylecun/mnist\", split=\"train\")\n",
    "test_dataset = load_dataset(\"ylecun/mnist\", split=\"test\")\n",
    "\n",
    "X_train = train_dataset[\"image\"]\n",
    "y_train = train_dataset[\"label\"]\n",
    "X_test = test_dataset[\"image\"]\n",
    "y_test = test_dataset[\"label\"]\n",
    "\n",
    "X_train = np.array(X_train).astype(np.float32) / 255.0\n",
    "y_train = np.array(y_train)\n",
    "X_test = np.array(X_test).astype(np.float32) / 255.0\n",
    "y_test = np.array(y_test)\n",
    "print(X_train.shape, y_train.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Model Definition"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model Summary:\n",
      "------------------------------------------------------------\n",
      "Optimizer: Adam | Learning Rate: 0.05 | Batch Size: 64 \n",
      "Epochs: 5 | Loss: CategoricalCrossEntropy\n",
      "------------------------------------------------------------\n",
      "Layer 0: FLATTEN   \n",
      "Layer 1: DENSE      | Input: 784   Output: 128  \n",
      "Layer 2: RELU      \n",
      "Layer 3: DENSE      | Input: 128   Output: 64   \n",
      "Layer 4: RELU      \n",
      "Layer 5: DENSE      | Input: 64    Output: 10   \n",
      "Layer 6: SOFTMAX   \n",
      "------------------------------------------------------------\n"
     ]
    }
   ],
   "source": [
    "model = (\n",
    "    SequentialBuilder()\n",
    "    .flatten()\n",
    "    .dense(784, 128, Xavier(InitType.NORMAL))\n",
    "    .relu()\n",
    "    .dense(128, 64, Xavier(InitType.NORMAL))\n",
    "    .relu()\n",
    "    .dense(64, 10, Xavier(InitType.NORMAL))\n",
    "    .softmax()\n",
    "    .optimizer(Adam())\n",
    "    .loss(CategoricalCrossEntropy())\n",
    "    .alpha(0.05)\n",
    "    .epochs(5)\n",
    "    .batch(64)\n",
    "    .build()\n",
    ")\n",
    "\n",
    "model.summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training \n",
    "First we check prediction accuracy before training, then we train the model for 5 epochs, and finally check the accuracy again to see the improvement."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test Accuracy: 9.8000 %\n"
     ]
    }
   ],
   "source": [
    "y_pred = model.predict(X_test)\n",
    "accuracy = np.mean(np.argmax(y_pred, axis=1) == y_test) * 100\n",
    "print(f\"Test Accuracy: {accuracy:.4f} %\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 0, Loss: 0.1890, Test Loss: 0.1935\n",
      "Training complete.\n",
      "------------------------------------------------------------\n",
      "Starting Training Loss: 0.1890 | Starting Test Loss: 0.1935\n",
      "Final Training Loss: 0.1077 | Final Test Loss: 0.1588\n",
      "Training Loss Improvement: 0.0813 | Test Loss Improvement: 0.0347\n",
      "------------------------------------------------------------\n"
     ]
    }
   ],
   "source": [
    "losses = model.train(X_train, y_train, X_test, y_test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test Accuracy: 96.1000 %\n"
     ]
    }
   ],
   "source": [
    "y_pred = model.predict(X_test)\n",
    "accuracy = np.mean(np.argmax(y_pred, axis=1) == y_test) * 100\n",
    "print(f\"Test Accuracy: {accuracy:.4f} %\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.14.4"
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "state": {},
    "version_major": 2,
    "version_minor": 0
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}