Update Entity_extraction_JSON notebook to use genai module (#509)

Refers #446

* Updated module google-generativeai to use google-genai>=1.0.0
* Refactored imports

Signed-off-by: Faakhir30 <[email protected]>

Author: Faakhir30

examples/json_capabilities/Entity_Extraction_JSON.ipynb

@@ -64,28 +64,13 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 8,
+      "execution_count": 5,
       "metadata": {
         "id": "Ne-3gnXqR0hI"
       },
       "outputs": [],
       "source": [
-        "!pip install -U -q \"google-generativeai>=0.7.2\""
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": 9,
-      "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"
+        "%pip install -U -q \"google-genai>=1.0.0\""
       ]
     },
     {
@@ -101,16 +86,17 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 10,
+      "execution_count": 6,
       "metadata": {
         "id": "v-JZzORUpVR2"
       },
       "outputs": [],
       "source": [
         "from google.colab import userdata\n",
-        "GOOGLE_API_KEY=userdata.get('GOOGLE_API_KEY')\n",
+        "from google import genai\n",
         "\n",
-        "genai.configure(api_key=GOOGLE_API_KEY)"
+        "GOOGLE_API_KEY=userdata.get('GOOGLE_API_KEY')\n",
+        "client = genai.Client(api_key=GOOGLE_API_KEY)"
       ]
     },
     {
@@ -124,12 +110,16 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 20,
+      "execution_count": 7,
       "metadata": {
         "id": "QGdJnd0AOKbu"
       },
       "outputs": [],
       "source": [
+        "from enum import Enum\n",
+        "from typing_extensions import TypedDict # in python 3.12 replace typing_extensions with typing\n",
+        "from google.genai import types\n",
+        "\n",
         "entity_recognition_text = \"John Johnson, the CEO of the Oil Inc. and Coal Inc. companies, has unveiled plans to build a new factory in Houston, Texas.\"\n",
         "prompt = f\"\"\"\n",
         "Generate list of entities in text based on the following Python class structure:\n",
@@ -148,13 +138,21 @@
         "  entities: list[Entity]\n",
         "\n",
         "{entity_recognition_text}\"\"\"\n",
-        "model = genai.GenerativeModel(model_name='gemini-2.0-flash', generation_config={\"temperature\": 0})\n",
-        "response = model.generate_content(prompt, generation_config={\"response_mime_type\": \"application/json\"})"
+        "MODEL_ID = \"gemini-2.0-flash\" # @param [\"gemini-2.0-flash-lite\",\"gemini-2.0-flash\",\"gemini-2.0-pro-exp-02-05\"] {\"allow-input\":true, isTemplate: true}\n",
+        "\n",
+        "response = client.models.generate_content(\n",
+        "    model=MODEL_ID,\n",
+        "    contents=prompt,\n",
+        "    config=types.GenerateContentConfig(\n",
+        "        temperature=0,\n",
+        "        response_mime_type=\"application/json\"\n",
+        "    )\n",
+        "  )"
       ]
     },
     {
       "cell_type": "code",
-      "execution_count": 21,
+      "execution_count": 8,
       "metadata": {
         "id": "d5tOgde6ONo3"
       },
@@ -191,6 +189,8 @@
         }
       ],
       "source": [
+        "import json\n",
+        "\n",
         "print(json.dumps(json.loads(response.text), indent=4))"
       ]
     },