{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# AI Code Assistant at NERSC\n",
        "\n",
        "If you are used to Github Copilot on VS Code or on Google Colaboratory\n",
        "as I am, coding without an AI assistant is so slow!\n",
        "\n",
        "In this tutorial we will see how to activate Google AI Gemini on\n",
        "[Jupyter@NERSC](https://jupyter.nersc.gov/), this will work as well in\n",
        "any other Jupyter environment. I choose Google Gemini because it has a\n",
        "generous free tier allowance, ChatGPT instead requires to buy credits\n",
        "for using it through the API.\n",
        "\n",
        "Notice that using the Gemini API you are sending your interactions with\n",
        "AI to Google, see the [Gemini API privacy\n",
        "policy](https://support.google.com/gemini/answer/13594961) for more\n",
        "details.\n",
        "\n",
        "The good news is that we do not need to be in control of the environment\n",
        "running JupyterHub, we can successfully install what is necessary in a\n",
        "Jupyter kernel we control.\n",
        "\n",
        "For example we could use our own [conda environment at\n",
        "NERSC](./2024-05-02-python-nersc-conda.md).\n",
        "\n",
        "This web page was generated from a Notebook, from the sidebar you can\n",
        "directly download the source Notebook, upload it to Jupyter@NERSC and\n",
        "directly execute it there, no need to copy-paste.\n",
        "\n",
        "## Install packages\n",
        "\n",
        "First make sure you are running the right Jupyter Kernel, which should\n",
        "be running in your own conda environment, see above on instructions on\n",
        "how to create it. You need to install the packages just once, either\n",
        "from this notebook or from the command line after having activated the\n",
        "environment. Once the packages are installed, no need to run this cell\n",
        "anymore.\n",
        "\n",
        "Jupyter AI supports many model providers, each with different required\n",
        "packages:"
      ],
      "id": "f1f3ba44-2e85-4bb0-8cf7-20f11ac0531a"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "%pip install jupyter-ai langchain-google-genai"
      ],
      "id": "c38ec60f"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Configure Google API key\n",
        "\n",
        "Next we need a Google API key to authenticate, go to:\n",
        "\n",
        "<https://aistudio.google.com/app/apikey>\n",
        "\n",
        "create a new secret key, and paste it below:"
      ],
      "id": "d626a218-6fd7-40e0-b7f0-ca0ba341e833"
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {
        "tags": []
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "env: GOOGLE_API_KEY=xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx"
          ]
        }
      ],
      "source": [
        "%env GOOGLE_API_KEY=xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx"
      ],
      "id": "31f2e432-d414-4354-ace1-182b0d2735d4"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Load and test Jupyter AI\n",
        "\n",
        "Finally we can test it"
      ],
      "id": "a2f0f19a-2188-4e32-ad26-ab0514b8c02a"
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "tags": []
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "/global/common/software/cmb/zonca/conda/pycmb/lib/python3.10/site-packages/tqdm/auto.py:22: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
            "  from .autonotebook import tqdm as notebook_tqdm"
          ]
        }
      ],
      "source": [
        "%load_ext jupyter_ai"
      ],
      "id": "ee6a0457-333e-4fdd-b11f-e635cf4e9e6b"
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "tags": []
      },
      "outputs": [],
      "source": [
        "%ai list gemini"
      ],
      "id": "6682d616-b67f-41fe-884c-1f19c9b75586"
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "tags": []
      },
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {
            "text/html": [
              "AI generated code inserted below &#11015;&#65039;"
            ]
          }
        }
      ],
      "source": [
        "%%ai gemini:gemini-pro -f code\n",
        "\n",
        "find indices of the largest 10 values in a numpy array"
      ],
      "id": "d62d0ae9-68c5-4717-a21d-145c9ca5fff2"
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {},
      "outputs": [],
      "source": [
        "import numpy as np\n",
        "\n",
        "def find_indices_of_largest_10_values(array):\n",
        "  \"\"\"Finds the indices of the largest 10 values in a numpy array.\n",
        "\n",
        "  Args:\n",
        "    array: A numpy array.\n",
        "\n",
        "  Returns:\n",
        "    A list of the indices of the largest 10 values in the array.\n",
        "  \"\"\"\n",
        "\n",
        "  # Find the indices of the largest 10 values in the array.\n",
        "  indices = np.argsort(array)[-10:]\n",
        "\n",
        "  # Return the list of indices.\n",
        "  return indices"
      ],
      "id": "c7ec34b2"
    }
  ],
  "nbformat": 4,
  "nbformat_minor": 5,
  "metadata": {
    "@webio": {
      "lastCommId": "",
      "lastKernelId": ""
    },
    "kernelspec": {
      "name": "pycmb",
      "display_name": "pycmb",
      "language": "python"
    },
    "language_info": {
      "name": "python",
      "codemirror_mode": {
        "name": "ipython",
        "version": "3"
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.10.0"
    }
  }
}