AI Code Assistant at NERSC

python
nersc
Published

May 28, 2024

If you are used to Github Copilot on VS Code or on Google Colaboratory as I am, coding without an AI assistant is so slow!

In this tutorial we will see how to activate Google AI Gemini on Jupyter@NERSC, this will work as well in any other Jupyter environment. I choose Google Gemini because it has a generous free tier allowance, ChatGPT instead requires to buy credits for using it through the API.

Notice that using the Gemini API you are sending your interactions with AI to Google, see the Gemini API privacy policy for more details.

The good news is that we do not need to be in control of the environment running JupyterHub, we can successfully install what is necessary in a Jupyter kernel we control.

For example we could use our own conda environment at NERSC.

This web page was generated from a Notebook, from the sidebar you can directly download the source Notebook, upload it to Jupyter@NERSC and directly execute it there, no need to copy-paste.

Install packages

First make sure you are running the right Jupyter Kernel, which should be running in your own conda environment, see above on instructions on how to create it. You need to install the packages just once, either from this notebook or from the command line after having activated the environment. Once the packages are installed, no need to run this cell anymore.

Jupyter AI supports many model providers, each with different required packages:

%pip install jupyter-ai langchain-google-genai

Configure Google API key

Next we need a Google API key to authenticate, go to:

https://aistudio.google.com/app/apikey

create a new secret key, and paste it below:

%env GOOGLE_API_KEY=xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx
env: GOOGLE_API_KEY=xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx

Load and test Jupyter AI

Finally we can test it

%load_ext jupyter_ai
/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
  from .autonotebook import tqdm as notebook_tqdm
%ai list gemini
Provider Environment variable Set? Models
gemini GOOGLE_API_KEY ✅ 
%%ai gemini:gemini-pro -f code

find indices of the largest 10 values in a numpy array
AI generated code inserted below ⬇️ 
import numpy as np

def find_indices_of_largest_10_values(array):
  """Finds the indices of the largest 10 values in a numpy array.

  Args:
    array: A numpy array.

  Returns:
    A list of the indices of the largest 10 values in the array.
  """

  # Find the indices of the largest 10 values in the array.
  indices = np.argsort(array)[-10:]

  # Return the list of indices.
  return indices