Colab Magic

This notebook introduces Colab magic commands for PaLM. Magics make it easy to develop, test, compare, and evaluate prompts from within a Colab notebook.

View on ai.google.dev Try a Colab notebook View notebook on GitHub

Setup

Follow the steps below to install and test the magics.

Installing the PaLM magic

To use the PaLM magic commands in Colab or other IPython environment, you will first need to download and install the google-generativeai Python package.

%pip install -q google-generativeai

Loading the PaLM magic

Next, load the %%palm magic by using the %load_ext magic:

%load_ext google.generativeai.notebook

Test the installation

To test for correct installation of the magic commands, run %%palm --help. Note that you will also need a PaLM API key, if you don't have one already (see next step).

%%palm --help
usage: palm [-h] {run,compile,compare,eval} ...

A system for interacting with LLMs.

positional arguments:
  {run,compile,compare,eval}

options:
  -h, --help            show this help message and exit

Getting a PaLM API key

To use the PaLM API, you will need to create an API key. (You only need to do this step once.)

Set the API key in the notebook

Set your API key by running the cell below.

%env GOOGLE_API_KEY=YOUR PALM KEY

PaLM magic commands: run, compile, compare, and evaluate

PaLM magics provide four different commands:

  1. run
  2. compile
  3. compare
  4. evaluate

Command: palm run

The run command sends the contents of the cell to the model.

Because running prompts is so common, PaLM magics defaults to the run command if no command is given. For example, the next two cells are identical.

%%palm run
The opposite of hot is
%%palm
The opposite of hot is

Understanding the output

The Prompt column shows the text that was sent to the model, and the text_result column shows the result. The other columns will be introduced as you progress through this guide.

Prompt templates

Prompts do not have to be fixed strings. You can inject values into a prompt using template placeholders by using {curly braces}.

english_words = {
    # Each value here (hot, cold) will be substituted in for {word} in the prompt
    'word': ['hot', 'cold']
}
%%palm --inputs english_words
The opposite of {word} is

Understanding the output

The Input Num column tracks the index of the input word in the list(s). In these examples, Input Num of 0 is 'hot', and 1 is 'cold'.

Specifying multiple sets of inputs

You can also specify multiple sets of inputs at one time.

extreme_temperatures = {
    'word': ['hot', 'cold']
}
minor_temperatures = {
    'word': ['warm', 'chilly']
}
%%palm --inputs extreme_temperatures minor_temperatures
The opposite of {word} is

Reading data from Google Sheets

The PaLM magic can also read and write to Google Sheets. You will need to be logged in to access Sheets data. This section focuses on reading data from Sheets; a later section shows how you can write output to a Google Sheet.

Log in and authorize access to Sheets

Formatting a spreadsheet for use with the PaLM magic

Pass the ID or URL of a Google Sheet to the --sheets_input_names flag to load it up as template data.

Use the following format in your spreadsheet to use the data in a prompt template:

  1. Put the names of the variables (of your prompt template) in the first row of the sheet.
  2. Put the data to substitute for each variable in the rows below.

For example, if your prompt template has two variables to substitute, name and temperament, you would write your spreadsheet like this:

name temperament
Milo cheeky
Bigsly relaxed
Subra shy
%%palm --sheets_input_names https://docs.google.com/spreadsheets/d/1UHfpkmBqIX5RjeJcGXOevIEhMmEoKlf5f9teqwQyHqc/edit
Create a single sentence description of a monkey's personality. The monkey's name is {name} and it has a {temperament} temperament.

Try it yourself!

To try this out using your own data, create a new Sheet and pass the ID to --sheets_input_names. As well as ID and URL, you can also search your sheets by title, e.g. %%palm --sheets_input_names "Animal adjectives".

Combining Sheets inputs with Python inputs

Sheets inputs can also be combined with --inputs:

new_monkeys = {
    'name': ['Hackerella'],
    'temperament': ['clever'],
}
%%palm --inputs new_monkeys --sheets_input_names 1UHfpkmBqIX5RjeJcGXOevIEhMmEoKlf5f9teqwQyHqc 1UHfpkmBqIX5RjeJcGXOevIEhMmEoKlf5f9teqwQyHqc
Create a single sentence description of a monkey's personality. The monkey's name is {name} and it has a {temperament} temperament.

Command: palm eval

Use %%palm eval to compare the output of a prompt with known ground-truth data.

test_data = {
    "word": ["dog", "cat", "house"]
}
ground_truth = ["chien", "chat", "maison"]
%%palm eval --inputs test_data --ground_truth ground_truth
English: Hello
French: Bonjour
English: {word}
French:

Post processing model outputs

To perform ground-truth testing, you may need to post-process the model output.

Post-processing functions allow you to define a function that processes the model output. In the case of the eval command, only the result column is used in the final equality check.

Use the post_process_replace_fn decorator to define a function to post-process results:

from google.generativeai.notebook import magics

# Define a function to extract only the first response.
@magics.post_process_replace_fn
def extract_and_normalize(input):
  first_line, *unused = input.split('English:')
  return first_line.strip().lower()

The extract_and_normalize function defined above will take the output from the model and trim any repeated language pairs, leaving just the first response. Check out the post-processing section to learn more about post-processing.

%%palm eval --inputs test_data --ground_truth ground_truth | extract_and_normalize
English: Hello
French: Bonjour
English: {word}
French:

Command: palm compile

Use the %%palm compile command to convert a prompt with placeholders to a function callable from within Python.

All flags and post-processing are "compiled" into the function and will be used when invoked.

In this example, a function called translate_en_to_fr is created, using the extract_and_normalize post-processing function from before.

%%palm compile translate_en_to_fr | extract_and_normalize
English: Hello
French: Bonjour
English: {word}
French:
'Saved function to Python variable: translate_en_to_fr'
en_words = ['cat', 'dog']
translate_en_to_fr({'word': en_words})

Output formats

By default, a "compiled" function returns its output as an object that will be displayed as Pandas DataFrame. However, you can convert the results object to a DataFrame or dictionary with .as_dict() or .as_dataframe(), respectively.

For more information, see the --outputs flag.

results = translate_en_to_fr({'word': en_words}).as_dict()

fr_words = results['text_result']

for en, fr in zip(en_words, fr_words):
  print(f'{fr} is French for {en}')
chat is French for cat
chien is French for dog

Command: palm compare

%%palm compare runs compiled prompts and produces a table with the comparison results side-by-side, so you can inspect the differences.

%%palm compile few_shot_prompt
English: Hello
French: Bonjour
English: {word}
French:
'Saved function to Python variable: few_shot_prompt'
%%palm compile zero_shot_prompt
{word} translated to French is:
'Saved function to Python variable: zero_shot_prompt'
words = {
    "word": ["dog", "cat", "house"]
}
%%palm compare few_shot_prompt zero_shot_prompt --inputs words

Custom comparison functions

By default, compare just checks for equalilty in the returned results. However, you can specify one or more custom functions with the --compare_fn flag:.

def average_word_length(lhs, rhs):
  """Count the average number of words used across prompts."""
  return (len(lhs.split(' ')) + len(rhs.split(' '))) / 2

def shortest_answer(lhs, rhs):
  """Label the prompt that generated the shortest output."""
  if len(lhs) < len(rhs):
    return 'first'
  elif len(lhs) > len(rhs):
    return 'second'
  else:
    return 'same'
%%palm compare few_shot_prompt zero_shot_prompt --inputs words --compare_fn average_word_length shortest_answer

Other commands

Help

The --help flag displays the supported commands that you can pass directly to %%palm

Append --help to view detailed documentation for each command. For example,

%%palm run --help
usage: palm run [-h] [--model_type {echo,text}] [--temperature TEMPERATURE]
                [--model MODEL] [--candidate_count CANDIDATE_COUNT] [--unique]
                [--inputs INPUTS [INPUTS ...]]
                [--sheets_input_names SHEETS_INPUT_NAMES [SHEETS_INPUT_NAMES ...]]
                [--outputs OUTPUTS [OUTPUTS ...]]
                [--sheets_output_names SHEETS_OUTPUT_NAMES [SHEETS_OUTPUT_NAMES ...]]

options:
  -h, --help            show this help message and exit
  --model_type {echo,text}, -mt {echo,text}
                        The type of model to use.
  --temperature TEMPERATURE, -t TEMPERATURE
                        Controls the randomness of the output. Must be
                        positive. Typical values are in the range: [0.0, 1.0].
                        Higher values produce a more random and varied
                        response. A temperature of zero will be deterministic.
  --model MODEL, -m MODEL
                        The name of the model to use. If not provided, a
                        default model will be used.
  --candidate_count CANDIDATE_COUNT, -cc CANDIDATE_COUNT
                        The number of candidates to produce.
  --unique              Whether to dedupe candidates returned by the model.
  --inputs INPUTS [INPUTS ...], -i INPUTS [INPUTS ...]
                        Optional names of Python variables containing inputs
                        to use to instantiate a prompt. The variable must be
                        either: a dictionary {'key1': ['val1', 'val2'] ...},
                        or an instance of LLMFnInputsSource such as
                        SheetsInput.
  --sheets_input_names SHEETS_INPUT_NAMES [SHEETS_INPUT_NAMES ...], -si SHEETS_INPUT_NAMES [SHEETS_INPUT_NAMES ...]
                        Optional names of Google Sheets to read inputs from.
                        This is equivalent to using --inputs with the names of
                        variables that are instances of SheetsInputs, just
                        more convenient to use.
  --outputs OUTPUTS [OUTPUTS ...], -o OUTPUTS [OUTPUTS ...]
                        Optional names of Python variables to output to. If
                        the Python variable has not already been defined, it
                        will be created. If the variable is defined and is an
                        instance of LLMFnOutputsSink, the outputs will be
                        written through the sink's write_outputs() method.
  --sheets_output_names SHEETS_OUTPUT_NAMES [SHEETS_OUTPUT_NAMES ...], -so SHEETS_OUTPUT_NAMES [SHEETS_OUTPUT_NAMES ...]
                        Optional names of Google Sheets to write inputs to.
                        This is equivalent to using --outputs with the names
                        of variables that are instances of SheetsOutputs, just
                        more convenient to use.

Models

Use the --model flag to specify the PaLM model variant you wish to use.

See the list_models() method to retrieve the supported models. The PaLM magic can be used with any model supporting the generateText method.

%%palm run --model models/text-bison-001
My favourite color is

Model parameters

You can also configure model parameters, such as --candidate_count and --temperature.

%%palm run --model models/text-bison-001 --temperature 0.5
My favourite color is

Debugging: The echo model

An echo model is also available that will echo the prompt back to you. It does not make any API calls or consume your quota so it can be a fast and simple way to test output or post-processing.

%%palm --model_type echo
A duck's quack does not echo.

Export output to Python

In addition to displaying tabular output, the PaLM magic can save model output to Python variables, allowing you to manipulate them further or to export your results.

In this example, the output is saved to a Python variable: fave_colors

%%palm --outputs fave_colors
The best colors to wear in spring-time are

Output variables are custom objects that will display as Pandas DataFrames by default. They can be coerced into a Python dictionary or dataframe explicitly by calling as_dict() or as_pandas_dataframe().

from pprint import pprint

pprint(fave_colors.as_dict())
{'Input Num': [0],
 'Prompt': ['The best colors to wear in spring-time are'],
 'Prompt Num': [0],
 'Result Num': [0],
 'text_result': ['* Pastels: These soft, muted colors are perfect for the '
                 'springtime, as they are fresh and airy. Some popular pastel '
                 'colors include baby blue, mint green, and pale pink.\n'
                 '* Brights: If you want to make a statement, bright colors '
                 'are a great option for spring. Some popular bright colors '
                 'include fuchsia, cobalt blue, and yellow.\n'
                 '* Neutrals: Neutral colors are always a good choice, as they '
                 'can be easily dressed up or down. Some popular neutrals '
                 'include beige, gray, and white.\n'
                 '\n'
                 'When choosing colors to wear in the spring, it is important '
                 'to consider the occasion and your personal style. For '
                 'example, if you are attending a formal event, you may want '
                 'to choose a more muted color palette, such as pastels or '
                 'neutrals. If you are going for a more casual look, you may '
                 'want to choose brighter colors, such as brights or pastels.']}

Write to Google Sheets

You can save output back to Google Sheets, using --sheets_output_names. You must be logged in, and you must have the appropriate permissions to access private Sheets.

To try this out, create a new Sheet and name it Translation results. Like the input flag, the --sheets_output_names flag also accepts the sheet URL or ID in place of the textual name.

%%palm --inputs english_words --sheets_output_names "Translation results"
English: Hello
French: Bonjour
English: {word}
French:

The results are saved to a new tab and contain the same data you see here in Colab.

Example of a saved sheet

Generating multiple candidates

To generate more than one output for a single prompt, you can pass --candidate_count to the model. This is set to 1 by default, which outputs only the top result.

Sometimes the model will generate the same output across candidates. These can be filtered with the --unique flag, which de-duplicates results out of the candidate batch (but not across multiple prompts).

%%palm run --temperature 1.0 --candidate_count 8 --unique
In a single word, my favourite color is

The Result Num column distinguishes multiple candidates generated from the same prompt.

Post-processing model output

The broad range of possible outputs and structures can make it difficult to adapt the model's output to your problem domain. The PaLM magic provides post-processing options that allow you to modify or process model output using Python code.

Post-processing functions can either add a new column to the output, or modify the text_result column. The text_result column is the last column, and is used by the eval and compare commands to determine the final output.

Here are some sample functions to use in post-processing. One adds a new column and the other updates the result column, using the post_process_replace_fn decorator.

import re
from google.generativeai.notebook import magics

# Add a new column.
def word_count(result):
  return len(result.split(' '))

# Modify the text_result column
@magics.post_process_replace_fn
def extract_first_sentence(result):
  """Extracts the first word from the raw result."""
  first, *_ = re.split(r'\.\s*', result)
  return first

To use these functions, append them to the %%palm command using the pipe (|) operator, like so.

%%palm run | word_count | extract_first_sentence
The happiest thing I can imagine is

Order matters here. When word_count is invoked, the original model output is used to calculate the number of words. If you swap these around, the word count would instead be the number of words in the extracted first sentence.

Further reading