{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/mlabonne/how-to-data-science/blob/main/Improve_ChatGPT_with_Knowledge_Graphs.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Improve ChatGPT with Knowledge Graphs\n",
        "\n",
        "❤️ Created by [@maximelabonne](https://twitter.com/maximelabonne)."
      ],
      "metadata": {
        "id": "vM5VDCJyAWnC"
      },
      "id": "vM5VDCJyAWnC"
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install -q openai langchain"
      ],
      "metadata": {
        "id": "uWz0SMuX5FFp"
      },
      "id": "uWz0SMuX5FFp",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import os\n",
        "import openai\n",
        "\n",
        "os.environ['OPENAI_API_KEY'] = \"your OpenAI key\"\n",
        "openai.api_key = os.environ['OPENAI_API_KEY']\n",
        "\n",
        "question = \"When did apple announced the Vision Pro?\"\n",
        "completion = openai.ChatCompletion.create(model=\"gpt-3.5-turbo\",\n",
        "                                          temperature=0,\n",
        "                                          messages=[{\"role\": \"user\",\n",
        "                                                     \"content\": question}])\n",
        "print(completion[\"choices\"][0][\"message\"][\"content\"])"
      ],
      "metadata": {
        "id": "XpvHDkzNnFLj",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "3d634edc-dbba-4294-8dcc-2ea998e7928e"
      },
      "id": "XpvHDkzNnFLj",
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "As an AI language model, I do not have access to current events or real-time information. However, as of my last training data, Apple has not announced any product called \"Vision Pro.\" It is possible that this product does not exist or has not been announced yet.\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "88a9f2e3-c729-455a-a338-2f83776c1d4c",
      "metadata": {
        "id": "88a9f2e3-c729-455a-a338-2f83776c1d4c",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "f6c9f4d8-387b-4ffc-e200-0c35ed449d8a"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "[('Apple', 'Vision Pro', 'announced'),\n",
              " ('Vision Pro', '2023', 'was announced in')]"
            ]
          },
          "metadata": {},
          "execution_count": 3
        }
      ],
      "source": [
        "from langchain.llms import OpenAI\n",
        "from langchain.indexes import GraphIndexCreator\n",
        "from langchain.chains import GraphQAChain\n",
        "from langchain.prompts import PromptTemplate\n",
        "\n",
        "text = \"Apple announced the Vision Pro in 2023.\"\n",
        "\n",
        "index_creator = GraphIndexCreator(llm=OpenAI(temperature=0))\n",
        "graph = index_creator.from_text(text)\n",
        "graph.get_triples()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "eecf2d57-3efa-4b0d-941a-95438d42893c",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        },
        "id": "eecf2d57-3efa-4b0d-941a-95438d42893c",
        "outputId": "1db33465-8f43-4294-84e0-0ad481b005aa"
      },
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<Figure size 2400x1500 with 1 Axes>"
            ],
            "image/png": "[encoded data, 116104 characters]\n"
          },
          "metadata": {}
        }
      ],
      "source": [
        "import networkx as nx\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "# Create graph\n",
        "G = nx.DiGraph()\n",
        "G.add_edges_from((source, target, {'relation': relation}) for source, relation, target in graph.get_triples())\n",
        "\n",
        "# Plot the graph\n",
        "plt.figure(figsize=(8,5), dpi=300)\n",
        "pos = nx.spring_layout(G, k=3, seed=0)\n",
        "\n",
        "nx.draw_networkx_nodes(G, pos, node_size=2000)\n",
        "nx.draw_networkx_edges(G, pos, edge_color='gray')\n",
        "nx.draw_networkx_labels(G, pos, font_size=12)\n",
        "edge_labels = nx.get_edge_attributes(G, 'relation')\n",
        "nx.draw_networkx_edge_labels(G, pos, edge_labels=edge_labels, font_size=10)\n",
        "\n",
        "# Display the plot\n",
        "plt.axis('off')\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "chain = GraphQAChain.from_llm(OpenAI(temperature=0), graph=graph, verbose=True)\n",
        "chain.run(question)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 199
        },
        "id": "KZIs4N5S8e8S",
        "outputId": "af6ba50c-e5f6-40be-adaf-07911dab948f"
      },
      "id": "KZIs4N5S8e8S",
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "\n",
            "\u001b[1m> Entering new GraphQAChain chain...\u001b[0m\n",
            "Entities Extracted:\n",
            "\u001b[32;1m\u001b[1;3m Apple, Vision Pro\u001b[0m\n",
            "Full Context:\n",
            "\u001b[32;1m\u001b[1;3mApple announced Vision ProVision Pro was announced in 2023\u001b[0m\n",
            "\n",
            "\u001b[1m> Finished chain.\u001b[0m\n"
          ]
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "' Apple announced Vision Pro in 2023.'"
            ],
            "application/vnd.google.colaboratory.intrinsic+json": {
              "type": "string"
            }
          },
          "metadata": {},
          "execution_count": 5
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "from langchain.graphs.networkx_graph import KnowledgeTriple\n",
        "\n",
        "# Knowledge graph\n",
        "kg = [\n",
        "    ('Apple', 'is', 'Company'),\n",
        "    ('Apple', 'created', 'iMac'),\n",
        "    ('Apple', 'created', 'iPhone'),\n",
        "    ('Apple', 'created', 'Apple Watch'),\n",
        "    ('Apple', 'created', 'Vision Pro'),\n",
        "\n",
        "    ('Apple', 'developed', 'macOS'),\n",
        "    ('Apple', 'developed', 'iOS'),\n",
        "    ('Apple', 'developed', 'watchOS'),\n",
        "\n",
        "    ('Apple', 'is located in', 'USA'),\n",
        "    ('Steve Jobs', 'co-founded', 'Apple'),\n",
        "    ('Steve Wozniak', 'co-founded', 'Apple'),\n",
        "    ('Tim Cook', 'is the CEO of', 'Apple'),\n",
        "\n",
        "    ('iOS', 'runs on', 'iPhone'),\n",
        "    ('macOS', 'runs on', 'iMac'),\n",
        "    ('watchOS', 'runs on', 'Apple Watch'),\n",
        "\n",
        "    ('Apple', 'was founded in', '1976'),\n",
        "    ('Apple', 'owns', 'App Store'),\n",
        "    ('App Store', 'sells', 'iOS apps'),\n",
        "\n",
        "    ('iPhone', 'announced in', '2007'),\n",
        "    ('iMac', 'announced in', '1998'),\n",
        "    ('Apple Watch', 'announced in', '2014'),\n",
        "    ('Vision Pro', 'announced in', '2023'),\n",
        "]\n",
        "\n",
        "graph = index_creator.from_text('')\n",
        "for (node1, relation, node2) in kg:\n",
        "    graph.add_triple(KnowledgeTriple(node1, relation, node2))"
      ],
      "metadata": {
        "id": "aa1R_DL-vTgv"
      },
      "id": "aa1R_DL-vTgv",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Create directed graph\n",
        "G = nx.DiGraph()\n",
        "for node1, relation, node2 in kg:\n",
        "    G.add_edge(node1, node2, label=relation)\n",
        "\n",
        "# Plot the graph\n",
        "plt.figure(figsize=(25, 25), dpi=300)\n",
        "pos = nx.spring_layout(G, k=2, iterations=50, seed=0)\n",
        "\n",
        "nx.draw_networkx_nodes(G, pos, node_size=5000)\n",
        "nx.draw_networkx_edges(G, pos, edge_color='gray', edgelist=G.edges(), width=2)\n",
        "nx.draw_networkx_labels(G, pos, font_size=12)\n",
        "edge_labels = nx.get_edge_attributes(G, 'label')\n",
        "nx.draw_networkx_edge_labels(G, pos, edge_labels=edge_labels, font_size=12)\n",
        "\n",
        "# Display the plot\n",
        "plt.axis('off')\n",
        "plt.show()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        },
        "id": "GhVuIxbUwJ1a",
        "outputId": "a3c0e3a8-b4a3-4baa-93ea-2daf03e352a2"
      },
      "id": "GhVuIxbUwJ1a",
      "execution_count": null,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<Figure size 7500x7500 with 1 Axes>"
            ],
            "image/png": "[encoded data, 1857040 characters]\n"
          },
          "metadata": {}
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "chain = GraphQAChain.from_llm(OpenAI(temperature=0), graph=graph, verbose=True)\n",
        "chain.run(question)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 380
        },
        "id": "7oGBfjuSuRNm",
        "outputId": "f4b03c71-c05e-492a-c9ba-21accfb994c8"
      },
      "id": "7oGBfjuSuRNm",
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "\n",
            "\u001b[1m> Entering new GraphQAChain chain...\u001b[0m\n",
            "Entities Extracted:\n",
            "\u001b[32;1m\u001b[1;3m Apple, Vision Pro\u001b[0m\n",
            "Full Context:\n",
            "\u001b[32;1m\u001b[1;3mApple is Company\n",
            "Apple created iMac\n",
            "Apple created iPhone\n",
            "Apple created Apple Watch\n",
            "Apple created Vision Pro\n",
            "Apple developed macOS\n",
            "Apple developed iOS\n",
            "Apple developed watchOS\n",
            "Apple is located in USA\n",
            "Apple was founded in 1976\n",
            "Apple owns App StoreVision Pro announced in 2023\u001b[0m\n",
            "\n",
            "\u001b[1m> Finished chain.\u001b[0m\n"
          ]
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "' Apple announced the Vision Pro in 2023.'"
            ],
            "application/vnd.google.colaboratory.intrinsic+json": {
              "type": "string"
            }
          },
          "metadata": {},
          "execution_count": 8
        }
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3 (ipykernel)",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.9.16"
    },
    "colab": {
      "provenance": [],
      "include_colab_link": true
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}