{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5342505d-4a76-404a-9011-7b6af316c6d3",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import time\n",
    "import urllib3\n",
    "import requests\n",
    "\n",
    "import yaml\n",
    "import mlflow\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "%matplotlib inline\n",
    "urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ffe8a279-4e33-4f65-86aa-d50114df871f",
   "metadata": {},
   "outputs": [],
   "source": [
    "TRAIN_EPOCHS = 10 # Number of epochs to train\n",
    "OUTPUTS_MIDDLE_LAYER = 1024 # Number of units for Dense layer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b3c6c972-a2a4-4ade-8374-e76fc7f03431",
   "metadata": {},
   "outputs": [],
   "source": [
    "def mnist_datasets():\n",
    "    (x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data(path='/mnt/shared/aie-tutorials/current-release/Data-Science/Kubeflow-GPU/mnist.npz')\n",
    "    x_train, x_test = x_train / np.float32(255), x_test / np.float32(255)\n",
    "    y_train, y_test = y_train.astype(np.int64), y_test.astype(np.int64)\n",
    "    return x_train, x_test, y_train, y_test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "794d7272-d629-4a1f-8e17-941be1b5ca37",
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_model():\n",
    "    model = tf.keras.models.Sequential([\n",
    "        tf.keras.layers.Flatten(input_shape=(28, 28)),\n",
    "        tf.keras.layers.Dense(OUTPUTS_MIDDLE_LAYER, activation='relu'),\n",
    "        tf.keras.layers.Dense(10, activation='softmax'),\n",
    "    ])\n",
    "\n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9d4aa20c-f992-4d04-8d88-aef8529730fa",
   "metadata": {},
   "outputs": [],
   "source": [
    "def compile_model(model):\n",
    "    model.compile(optimizer='adam',\n",
    "                  loss='sparse_categorical_crossentropy',\n",
    "                  metrics=['accuracy'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eea5e48c-c1ce-42d5-bc02-d106ea9f6abe",
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_model(model, x_train, x_test, y_test, y_train):\n",
    "    print(\"Num GPUs Available: \", len(tf.config.experimental.list_physical_devices('GPU')))\n",
    "    print()\n",
    "    # model fitting\n",
    "    start = time.time()\n",
    "    history = model.fit(\n",
    "        x_train, y_train, epochs=TRAIN_EPOCHS,\n",
    "        validation_data=(x_test, y_test)\n",
    "    )\n",
    "    duration_total = time.time() - start\n",
    "    print()\n",
    "    print('Total time %f sec' % (duration_total))\n",
    "    return history, duration_total"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2b238b1f-8ec6-4bf5-a53b-3d51cb3429db",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load dataset\n",
    "x_train, x_test, y_train, y_test = mnist_datasets()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a95b44d5-0f72-4da3-a01d-45fa057e5916",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create the model\n",
    "model_mnist = create_model()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "53c5a324-7b2c-4660-89da-b6406c750c7b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Train and evaluate for a set number of epochs.\n",
    "compile_model(model_mnist)\n",
    "history, duration = train_model(model_mnist, x_train, x_test, y_test, y_train)\n",
    "acc = history.history['accuracy'][-1]\n",
    "print(f\"Accuracy: {acc * 100}%\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b0a462d4-b03f-4967-a7ca-496ae7229402",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Save model in s3 via mlflow\n",
    "mlflow.set_experiment(\"trained_mnist_model\")\n",
    "signature = mlflow.models.signature.infer_signature(x_test, y_test)\n",
    "mlflow_tf_model_save_path_backup = mlflow.tensorflow._MODEL_SAVE_PATH\n",
    "mlflow.tensorflow._MODEL_SAVE_PATH += \"/1\"\n",
    "model_info = mlflow.tensorflow.log_model(model=model_mnist, artifact_path=\"model\", signature=signature, registered_model_name=\"tf-mnist-model\")\n",
    "mlflow.tensorflow._MODEL_SAVE_PATH = mlflow_tf_model_save_path_backup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9c078390-7980-4107-8785-09cae69995c9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get URI of saved model in s3 object store\n",
    "saved_model_path = os.path.join(mlflow.get_run(model_info.run_id).info.artifact_uri, model_info.artifact_path, \"data\", mlflow.tensorflow._MODEL_SAVE_PATH)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cc5c20bd-f07a-4581-bfe8-fd5637a77a5e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set current JWT token as access and secret key for local minio\n",
    "!sed -e \"s/\\$AUTH_TOKEN/$AUTH_TOKEN/\" /mnt/shared/aie-tutorials/current-release/Data-Science/Kubeflow-GPU/object_store_sa_secret.yaml.tpl > /tmp/object_store_sa_secret.yaml"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1dde54e3-ed30-448e-a1d8-e6238c9ab722",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create service account and secret in order to allow kserve access object store\n",
    "!kubectl apply -f /tmp/object_store_sa_secret.yaml"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "27f78480-5632-4bf5-a561-7548a8424f88",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create inferenceservice for trained mnist model\n",
    "with open('/mnt/shared/aie-tutorials/current-release/Data-Science/Kubeflow-GPU/gpu_mnist_inferenceservice.yaml', 'r') as file:\n",
    "    inferenceservice_yaml_data = yaml.safe_load(file)\n",
    "\n",
    "inferenceservice_yaml_data[\"spec\"][\"predictor\"][\"tensorflow\"][\"storageUri\"] = saved_model_path\n",
    "\n",
    "with open('/tmp/gpu_mnist_inferenceservice.yaml', 'w') as yaml_file:\n",
    "    yaml.dump(inferenceservice_yaml_data, yaml_file)\n",
    "\n",
    "!kubectl apply -f /tmp/gpu_mnist_inferenceservice.yaml"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2fd0abad-17c5-4ce8-87dc-c7cbf07746fe",
   "metadata": {},
   "outputs": [],
   "source": [
    "!until kubectl get pods -l serving.kserve.io/inferenceservice=gpu-mnist-inferenceservice | grep -q \"Running\"; do sleep 10; done\n",
    "!echo \"Inferenceservice is running\"\n",
    "!kubectl wait --for=condition=ready pod -l serving.kserve.io/inferenceservice=gpu-mnist-inferenceservice --timeout=600s"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "76cc876c-b237-4f18-9657-72ca84d2370f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Make a call to inferenceservice\n",
    "URL = f\"http://gpu-mnist-inferenceservice-predictor-default.{os.getenv('JUPYTERHUB_USER')}.svc.cluster.local/v1/models/gpu-mnist-inferenceservice:predict\"\n",
    "count_input_images = 10\n",
    "test_images_data = x_test[0:count_input_images]\n",
    "inputs = {\"inputs\": test_images_data[np.newaxis, ...].tolist()}\n",
    "headers = {\"Authorization\": f\"Bearer {os.getenv('AUTH_TOKEN')}\"}\n",
    "response = requests.post(URL, json=inputs, headers=headers)\n",
    "\n",
    "print(f\"Status: {response.reason}\")\n",
    "print(f\"JSON data: {response.json()}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "44e4fc23-a5a5-4b3d-98be-c0ed36c4a078",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Parse predicted numbers from inferenceservice output\n",
    "outputs = response.json()[\"outputs\"]\n",
    "predicted_numbers = [np.argmax(x) for x in outputs]\n",
    "print(f\"Numbers predicted by inferenceservice from input: {predicted_numbers}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "53abdf54-22a3-48c6-8ee9-bb84ab2dea92",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Show input images with corresponding predicted numbers\n",
    "num_col = 5\n",
    "num_row = int(np.ceil(count_input_images / num_col))\n",
    "\n",
    "# plot images\n",
    "fig, axes = plt.subplots(num_row, num_col, figsize=(1.5 * num_col, 2 * num_row))\n",
    "for i in range(count_input_images):\n",
    "    ax = axes[i // num_col, i % num_col]\n",
    "    ax.imshow(test_images_data[i], cmap='gray')\n",
    "    ax.set_title('Predicted: {}'.format(predicted_numbers[i]))\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  }
 ],
 "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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}