Add entity extraction example

examples/json_capabilities/Entity_Extraction_JSON.ipynb

@@ -0,0 +1,231 @@
+{
+  "cells": [
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "T47AX_Is2FjB"
+      },
+      "source": [
+        "##### Copyright 2024 Google LLC."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "M_vx0YO92qlR"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
+        "# you may not use this file except in compliance with the License.\n",
+        "# You may obtain a copy of the License at\n",
+        "#\n",
+        "# https://www.apache.org/licenses/LICENSE-2.0\n",
+        "#\n",
+        "# Unless required by applicable law or agreed to in writing, software\n",
+        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
+        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
+        "# See the License for the specific language governing permissions and\n",
+        "# limitations under the License."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "sP8PQnz1QrcF"
+      },
+      "source": [
+        "# Gemini API: Entity Extraction"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "bxGr_x3MRA0z"
+      },
+      "source": [
+        "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
+        "  <td>\n",
+        "    <a target=\"_blank\" href=\"https://colab.research.google.com/github/google-gemini/cookbook/blob/main/quickstarts/examples/json_capabilities/Entity_Extraction.ipynb\"><img src = \"https://www.tensorflow.org/images/colab_logo_32px.png\"/>Run in Google Colab</a>\n",
+        "  </td>\n",
+        "</table>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "ysy--KfNRrCq"
+      },
+      "source": [
+        "You will use Gemini to extract all fields that fit one of the predefined classes and label them."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 1,
+      "metadata": {
+        "id": "Ne-3gnXqR0hI"
+      },
+      "outputs": [],
+      "source": [
+        "!pip install -U -q google-generativeai"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 3,
+      "metadata": {
+        "id": "EconMHePQHGw"
+      },
+      "outputs": [],
+      "source": [
+        "import google.generativeai as genai\n",
+        "\n",
+        "import json\n",
+        "from enum import Enum\n",
+        "from typing_extensions import TypedDict # in python 3.12 replace typing_extensions with typing"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "eomJzCa6lb90"
+      },
+      "source": [
+        "## Configure your API key\n",
+        "\n",
+        "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 [Authentication](https://github.com/google-gemini/cookbook/blob/main/quickstarts/Authentication.ipynb) for an example."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 4,
+      "metadata": {
+        "id": "v-JZzORUpVR2"
+      },
+      "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": "R3EUoLgJNfe7"
+      },
+      "source": [
+        "## Example"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 5,
+      "metadata": {
+        "id": "sShzxm3JNm6M"
+      },
+      "outputs": [],
+      "source": [
+        "class CategoryEnum(str, Enum):\n",
+        "    Person = 'Person'\n",
+        "    Company = 'Company'\n",
+        "    State = 'State'\n",
+        "    City = 'City'\n",
+        "\n",
+        "class Entity(TypedDict):\n",
+        "  name: str\n",
+        "  category: str\n",
+        "\n",
+        "class Entities(TypedDict):\n",
+        "  entities: list[Entity]"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 6,
+      "metadata": {
+        "id": "QGdJnd0AOKbu"
+      },
+      "outputs": [],
+      "source": [
+        "entity_recongnition_text = \"Elon Musk, the CEO of Tesla and SpaceX, has unveiled plans to build a new factory in Austin, Texas. The factory will create 10,000 new jobs and produce electric vehicles.\"\n",
+        "prompt = f\"\"\"\n",
+        "Generate list of entities in text:\n",
+        "\n",
+        "{entity_recongnition_text}\"\"\"\n",
+        "model = genai.GenerativeModel(model_name='gemini-1.5-flash-latest', generation_config={\"temperature\": 0})\n",
+        "response = model.generate_content(prompt, generation_config={\"response_mime_type\": \"application/json\", \"response_schema\": Entities})"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 7,
+      "metadata": {
+        "id": "d5tOgde6ONo3"
+      },
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "{\n",
+            "    \"entities\": [\n",
+            "        {\n",
+            "            \"text\": \"Elon Musk\",\n",
+            "            \"type\": \"PERSON\"\n",
+            "        },\n",
+            "        {\n",
+            "            \"text\": \"Tesla\",\n",
+            "            \"type\": \"ORG\"\n",
+            "        },\n",
+            "        {\n",
+            "            \"text\": \"SpaceX\",\n",
+            "            \"type\": \"ORG\"\n",
+            "        },\n",
+            "        {\n",
+            "            \"text\": \"Austin\",\n",
+            "            \"type\": \"GPE\"\n",
+            "        },\n",
+            "        {\n",
+            "            \"text\": \"Texas\",\n",
+            "            \"type\": \"GPE\"\n",
+            "        }\n",
+            "    ]\n",
+            "}\n"
+          ]
+        }
+      ],
+      "source": [
+        "print(json.dumps(json.loads(response.text), indent=4))"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "2uv9Rikv27tf"
+      },
+      "source": [
+        "## Summary\n",
+        "You have used the Gemini API to extract entities of predifined categories with their labels. You extracted every person, company, state, and country. You are not limited to these categories, as this should work with any category of your choice.\n",
+        "\n",
+        "Please see the other notebooks in this directory to learn more about how you can use the Gemini API for other JSON related tasks.\n"
+      ]
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "name": "Entity_Extraction_JSON.ipynb",
+      "toc_visible": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}