Skip to content

Commit

Permalink
Add JSON Schema quickstart. (google-gemini#170)
Browse files Browse the repository at this point in the history 
* Add example for adding context information for prompting

* delete

* Add JSON schema quickstart

* Add info about Gemini 1.5 pro

* Combining JSON mode and schema notebooks

* Reformatted notebook

* replace we with you

* format

* Use the same example in both halves.

---------

Co-authored-by: Mark Daoust <[email protected]>
  • Loading branch information
@shilpakancharla @MarkDaoust
shilpakancharla and MarkDaoust authored May 28, 2024
1 parent 4239715 commit 4957058
Showing 1 changed file with 140 additions and 54 deletions.
194 changes: 140 additions & 54 deletions quickstarts/JSON_mode.ipynb
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@
},
{
"cell_type": "code",
"execution_count": null,
"execution_count": 1,
"metadata": {
"cellView": "form",
"id": "tuOe1ymfHZPu"
Expand All @@ -34,7 +34,7 @@
{
"cell_type": "markdown",
"metadata": {
"id": "893sOzyhJDma"
"id": "GAsiP4mohC2_"
},
"source": [
"# Gemini API: JSON Mode Quickstart\n",
Expand All @@ -49,151 +49,165 @@
{
"cell_type": "markdown",
"metadata": {
"id": "h4LQoYRTJIP9"
"id": "lF6sWVRGQ_bi"
},
"source": [
"This notebook demonstrates how to use JSON mode."
"The Gemini API can be used to generate a JSON output if you set the schema that you would like to use.\n",
"\n",
"**Note**: JSON schemas are only supported by Gemini 1.5 Pro right now."
]
},
{
"cell_type": "code",
"execution_count": 1,
"execution_count": null,
"metadata": {
"id": "_PBH7eR9He0I"
"id": "qLuL9m7KhvxR"
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m146.8/146.8 kB\u001b[0m \u001b[31m1.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
"\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m664.5/664.5 kB\u001b[0m \u001b[31m6.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
"\u001b[?25h"
]
}
],
"outputs": [],
"source": [
"!pip install -qU google-generativeai"
"!pip install -U -q google-generativeai"
]
},
{
"cell_type": "code",
"execution_count": 2,
"metadata": {
"id": "2zwIBNLWJvRf"
"id": "ATIbQM0NHhkj"
},
"outputs": [],
"source": [
"import google.generativeai as genai\n",
"import json"
"\n",
"import json\n",
"import dataclasses\n",
"import typing_extensions as typing"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "F6gHNgcUypVN"
"id": "B-axqBTM8Lbd"
},
"source": [
"To run the following cell, your API key must be stored it in a Colab Secret named `GOOGLE_API_KEY`. If you don't already have an API key, or you're not sure how to create a Colab Secret, see the [Authentication](https://github.com/google-gemini/cookbook/blob/main/quickstarts/Authentication.ipynb) quickstart for an example."
"## Configure your API key\n",
"\n",
"To run the following cell, your API key must be stored in a Colab Secret named `GOOGLE_API_KEY`. If you don't already have an API key, or you're not sure how to create a Colab Secret, see [Authentication](https://github.com/google-gemini/cookbook/blob/main/quickstarts/Authentication.ipynb) for an example."
]
},
{
"cell_type": "code",
"execution_count": 3,
"execution_count": null,
"metadata": {
"id": "t0jy9XWjJwv7"
"id": "d6lYXRcjthKV"
},
"outputs": [],
"source": [
"from google.colab import userdata\n",
"GOOGLE_API_KEY=userdata.get('GOOGLE_API_KEY')\n",
"\n",
"genai.configure(api_key=GOOGLE_API_KEY)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "vf42XN1KLcfV"
"id": "hD3qXcOTRD3z"
},
"source": [
"## Activate JSON Mode\n",
"\n",
" Activate JSON mode by specifying `respose_mime_type` in the `generation_config` parameter:"
]
},
{
"cell_type": "code",
"execution_count": 3,
"metadata": {
"id": "i5Rod-lXRIhf"
},
"outputs": [],
"source": [
"## Activate JSON mode"
"model = genai.GenerativeModel(\"gemini-1.5-flash-latest\",\n",
" generation_config={\"response_mime_type\": \"application/json\"})"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "dC5-79CDMJ3R"
"id": "4071a6143d31"
},
"source": [
"Activate JSON mode by specifying `respose_mime_type` in the `generation_config` parameter."
"For this first example just describe the schema you want back:"
]
},
{
"cell_type": "code",
"execution_count": 4,
"metadata": {
"id": "WWq64FXSLXgr"
"id": "K8ezjNb0RJ6Y"
},
"outputs": [],
"source": [
"model = genai.GenerativeModel(\"gemini-1.5-pro-latest\",\n",
" generation_config={\"response_mime_type\": \"application/json\"})"
"prompt = \"\"\"List a few popular cookie recipes using this JSON schema:\n",
"\n",
"Recipe = {'recipe_name': str}\n",
"Return: list[Recipe]\"\"\""
]
},
{
"cell_type": "code",
"execution_count": 5,
"metadata": {
"id": "Y_djQzyyaCLg"
"id": "ggudoxK8RMlb"
},
"outputs": [],
"source": [
"prompt = \"\"\"List a few popular cookie recipes using this JSON schema:\n",
"{'type': 'object', 'properties': { 'recipe_name': {'type': 'string'}}}\"\"\""
"raw_response = model.generate_content(prompt)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "9TqoNg3VSMYB"
},
"source": [
"Parse the string to JSON:"
]
},
{
"cell_type": "code",
"execution_count": 6,
"metadata": {
"id": "aENeySrWMJN6"
"id": "b99ee66972f5"
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"[\n",
" {\"recipe_name\": \"Chocolate Chip Cookies\"},\n",
" {\"recipe_name\": \"Peanut Butter Cookies\"},\n",
" {\"recipe_name\": \"Oatmeal Raisin Cookies\"},\n",
" {\"recipe_name\": \"Sugar Cookies\"},\n",
" {\"recipe_name\": \"Shortbread Cookies\"}\n",
"]\n",
"\n",
"\n"
"[{'recipe_name': 'Chocolate Chip Cookies'}, {'recipe_name': 'Oatmeal Raisin Cookies'}, {'recipe_name': 'Snickerdoodles'}, {'recipe_name': 'Sugar Cookies'}, {'recipe_name': 'Peanut Butter Cookies'}]\n"
]
}
],
"source": [
"response = model.generate_content(prompt)\n",
"print(response.text)"
"response = json.loads(raw_response.text)\n",
"print(response)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "pqNsOE1YysLc"
"id": "1092c669169a"
},
"source": [
"Just for fun, parse the string to JSON, and then serialize it."
"For readability searialize and print it:"
]
},
{
"cell_type": "code",
"execution_count": 9,
"execution_count": 7,
"metadata": {
"id": "nb9Z9TdHRzTu"
"id": "WLDPREpmSMu5"
},
"outputs": [
{
Expand All @@ -205,23 +219,95 @@
" \"recipe_name\": \"Chocolate Chip Cookies\"\n",
" },\n",
" {\n",
" \"recipe_name\": \"Peanut Butter Cookies\"\n",
" \"recipe_name\": \"Oatmeal Raisin Cookies\"\n",
" },\n",
" {\n",
" \"recipe_name\": \"Oatmeal Raisin Cookies\"\n",
" \"recipe_name\": \"Snickerdoodles\"\n",
" },\n",
" {\n",
" \"recipe_name\": \"Sugar Cookies\"\n",
" },\n",
" {\n",
" \"recipe_name\": \"Shortbread Cookies\"\n",
" \"recipe_name\": \"Peanut Butter Cookies\"\n",
" }\n",
"]\n"
]
}
],
"source": [
"print(json.dumps(json.loads(response.text), indent=4))"
"print(json.dumps(response, indent=4))"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "K9nIks0R-tIa"
},
"source": [
"## Generate JSON from schema\n",
"\n",
"While `gemini-1.5-flash` models only accept a text description of the JSON you want back, `gemini-1.5-pro` models support \"controlled generation\" (aka \"constrained decoding\"). This allows you to pass a schema object (or a python type equivalent) and the output will strictly follow that schema.\n",
"\n",
"Following the same example as the previous section, here's that recipe type:"
]
},
{
"cell_type": "code",
"execution_count": 8,
"metadata": {
"id": "JiIxKaLl4R0f"
},
"outputs": [],
"source": [
"class Recipe(typing.TypedDict):\n",
" recipe_name: str"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "vBlWzt6M-2oM"
},
"source": [
"For this exaple you want a list of `Recipe` objects, so pass `list[Recipe]` to the `response_schema` field of the `generation_config`."
]
},
{
"cell_type": "code",
"execution_count": 9,
"metadata": {
"id": "8oe-tL8MDGtx"
},
"outputs": [],
"source": [
"model = genai.GenerativeModel(model_name=\"models/gemini-1.5-pro-latest\")\n",
"\n",
"result = model.generate_content(\n",
" \"List a few popular cookie recipes\",\n",
" generation_config=genai.GenerationConfig(\n",
" response_mime_type=\"application/json\",\n",
" response_schema = list[Recipe]),\n",
" request_options={\"timeout\": 600},\n",
")"
]
},
{
"cell_type": "code",
"execution_count": 10,
"metadata": {
"id": "slYcVAcqaDQY"
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"[{\"recipe_name\": \"Chocolate Chip Cookies\"}, {\"recipe_name\": \"Peanut Butter Cookies\"}, {\"recipe_name\": \"Oatmeal Raisin Cookies\"}, {\"recipe_name\": \"Sugar Cookies\"}, {\"recipe_name\": \"Snickerdoodles\"}] \n"
]
}
],
"source": [
"print(result.text)"
]
}
],
Expand Down

0 comments on commit 4957058

Please sign in to comment.