1. Generative AI.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "fd768455",
   "metadata": {},
   "source": [
    "<div style=\"width: 100%; overflow: hidden;\">\n",
    "    <div style=\"width: 150px; float: left;\"> <img src=\"https://raw.githubusercontent.com/DataForScience/Networks/master/data/D4Sci_logo_ball.png\" alt=\"Data For Science, Inc\" align=\"left\" border=\"0\" width=150px> </div>\n",
    "    <div style=\"float: left; margin-left: 10px;\"> <h1>LLMs for Data Science</h1>\n",
    "<h1>Generative AI</h1>\n",
    "        <p>Bruno Gonçalves<br/>\n",
    "        <a href=\"http://www.data4sci.com/\">www.data4sci.com</a><br/>\n",
    "            @bgoncalves, @data4sci</p></div>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "7528bcfd",
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import Counter\n",
    "from pprint import pprint\n",
    "from datetime import datetime\n",
    "import json\n",
    "\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "\n",
    "import matplotlib\n",
    "import matplotlib.pyplot as plt \n",
    "import sqlite3\n",
    "\n",
    "import openai\n",
    "from openai import OpenAI\n",
    "\n",
    "import transformers\n",
    "from transformers import pipeline\n",
    "from transformers import set_seed\n",
    "set_seed(42) # Set the seed to get reproducible results\n",
    "\n",
    "import langchain\n",
    "import langchain_openai\n",
    "from langchain_openai import ChatOpenAI\n",
    "import langchain_core\n",
    "from langchain_core.messages import HumanMessage, SystemMessage\n",
    "from langchain_core.output_parsers import StrOutputParser\n",
    "\n",
    "import os\n",
    "import gzip\n",
    "\n",
    "import tqdm as tq\n",
    "from tqdm.notebook import tqdm\n",
    "\n",
    "import watermark\n",
    "\n",
    "%load_ext watermark\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "375fc2e7",
   "metadata": {},
   "source": [
    "We start by printing out the versions of the libraries we're using for future reference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "59e9a07f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Python implementation: CPython\n",
      "Python version       : 3.11.7\n",
      "IPython version      : 8.12.3\n",
      "\n",
      "Compiler    : Clang 14.0.6 \n",
      "OS          : Darwin\n",
      "Release     : 23.6.0\n",
      "Machine     : arm64\n",
      "Processor   : arm\n",
      "CPU cores   : 16\n",
      "Architecture: 64bit\n",
      "\n",
      "Git hash: 8d244e1e4f0c6fd330052d22607886f6abfcd26c\n",
      "\n",
      "numpy           : 1.26.4\n",
      "tqdm            : 4.66.4\n",
      "watermark       : 2.4.3\n",
      "json            : 2.0.9\n",
      "sqlite3         : 2.6.0\n",
      "langchain_core  : 0.2.3\n",
      "pandas          : 2.2.3\n",
      "langchain_openai: 0.1.8\n",
      "langchain       : 0.2.2\n",
      "matplotlib      : 3.8.0\n",
      "transformers    : 4.41.1\n",
      "openai          : 1.30.5\n",
      "\n"
     ]
    }
   ],
   "source": [
    "%watermark -n -v -m -g -iv"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "07149b03",
   "metadata": {},
   "source": [
    "Load default figure style"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "aeb6f060",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.style.use('d4sci.mplstyle')\n",
    "colors = plt.rcParams['axes.prop_cycle'].by_key()['color']"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f91941b6",
   "metadata": {},
   "source": [
    "# OpenAI Basic Usage"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c0308f88",
   "metadata": {},
   "source": [
    "The first step is generate API key on the OpenAI website and store it as the \"OPENAI_API_KEY\" variable in your local environment. Without it we won't be able to do anything. You can find your API key in your using settings: https://help.openai.com/en/articles/4936850-where-do-i-find-my-secret-api-key"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6bba9767",
   "metadata": {},
   "source": [
    "Then we are ready to instantiate the client"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "061080e7",
   "metadata": {},
   "outputs": [],
   "source": [
    "client = OpenAI()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b8c10638",
   "metadata": {},
   "source": [
    "We start by getting a list of supported models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "2a6766b1",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_list = json.loads(client.models.list().json())[\"data\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e1d0d773",
   "metadata": {},
   "source": [
    "In total we have 43 models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "0355de4a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "43"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(model_list)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d005b248",
   "metadata": {},
   "source": [
    "Along with some information about each model..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "506d5798",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'id': 'text-embedding-3-small',\n",
       "  'created': 1705948997,\n",
       "  'object': 'model',\n",
       "  'owned_by': 'system'},\n",
       " {'id': 'gpt-4o-mini',\n",
       "  'created': 1721172741,\n",
       "  'object': 'model',\n",
       "  'owned_by': 'system'},\n",
       " {'id': 'gpt-4o-realtime-preview',\n",
       "  'created': 1727659998,\n",
       "  'object': 'model',\n",
       "  'owned_by': 'system'}]"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model_list[:3]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1621f282",
   "metadata": {},
   "source": [
    "But let's just get a list of model names"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "b60be782",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "babbage-002\n",
      "chatgpt-4o-latest\n",
      "dall-e-2\n",
      "dall-e-3\n",
      "davinci-002\n",
      "gpt-3.5-turbo\n",
      "gpt-3.5-turbo-0125\n",
      "gpt-3.5-turbo-1106\n",
      "gpt-3.5-turbo-16k\n",
      "gpt-3.5-turbo-16k-0613\n",
      "gpt-3.5-turbo-instruct\n",
      "gpt-3.5-turbo-instruct-0914\n",
      "gpt-4\n",
      "gpt-4-0125-preview\n",
      "gpt-4-0613\n",
      "gpt-4-1106-preview\n",
      "gpt-4-turbo\n",
      "gpt-4-turbo-2024-04-09\n",
      "gpt-4-turbo-preview\n",
      "gpt-4o\n",
      "gpt-4o-2024-05-13\n",
      "gpt-4o-2024-08-06\n",
      "gpt-4o-2024-11-20\n",
      "gpt-4o-audio-preview\n",
      "gpt-4o-audio-preview-2024-10-01\n",
      "gpt-4o-mini\n",
      "gpt-4o-mini-2024-07-18\n",
      "gpt-4o-realtime-preview\n",
      "gpt-4o-realtime-preview-2024-10-01\n",
      "o1-mini\n",
      "o1-mini-2024-09-12\n",
      "o1-preview\n",
      "o1-preview-2024-09-12\n",
      "omni-moderation-2024-09-26\n",
      "omni-moderation-latest\n",
      "text-embedding-3-large\n",
      "text-embedding-3-small\n",
      "text-embedding-ada-002\n",
      "tts-1\n",
      "tts-1-1106\n",
      "tts-1-hd\n",
      "tts-1-hd-1106\n",
      "whisper-1\n"
     ]
    }
   ],
   "source": [
    "print(\"\\n\".join(sorted([model[\"id\"] for model in model_list])))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8cb50571",
   "metadata": {},
   "source": [
    "## Basic Prompt"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2cc88f45",
   "metadata": {},
   "source": [
    "The recommended model for exploration is `gpt-3.5-turbo`, so we'll stick with it for now. The basic setup is relatively straightforward:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "70a70474",
   "metadata": {},
   "outputs": [],
   "source": [
    "response = client.chat.completions.create(\n",
    "  model=\"gpt-3.5-turbo\",\n",
    "  messages=[\n",
    "        {\n",
    "            \"role\": \"user\", \n",
    "            \"content\": \"What was Superman's weakness?\"\n",
    "        },\n",
    "    ]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b321cbc0",
   "metadata": {},
   "source": [
    "Which produces a response object"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "0add7d1f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "openai.types.chat.chat_completion.ChatCompletion"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "type(response)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9b157ad5",
   "metadata": {},
   "source": [
    "Which we can treat as a named tuple"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ea373496",
   "metadata": {},
   "source": [
    "The model answer can be found in the \"message\" dictionary inside the \"choices\" list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "d1b4cd79",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Choice(finish_reason='stop', index=0, logprobs=None, message=ChatCompletionMessage(content=\"Superman's weakness is kryptonite, a radioactive mineral from his home planet of Krypton. Exposure to kryptonite weakens and can even kill Superman.\", role='assistant', function_call=None, tool_calls=None, refusal=None))"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "response.choices[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "ac822e91",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\"Superman's weakness is kryptonite, a radioactive mineral from his home planet of Krypton. Exposure to kryptonite weakens and can even kill Superman.\""
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "response.choices[0].message.content"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "267bb1c9",
   "metadata": {},
   "source": [
    "To request multiple answers, we must include the `n` parameter with the number of answers we want"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "c2d2b662",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 7.72 ms, sys: 5.92 ms, total: 13.6 ms\n",
      "Wall time: 3.76 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "response = client.chat.completions.create(\n",
    "    model=\"gpt-3.5-turbo\",\n",
    "    messages=[\n",
    "        {\n",
    "            \"role\": \"user\", \n",
    "             \"content\": \"What are the different kinds of Kryptonite?\"\n",
    "        },\n",
    "    ],\n",
    "    n=3\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a56f7d9c",
   "metadata": {},
   "source": [
    "And we can access each of the answers individually int he choices list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "20956e7c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "==========\n",
      "Assistant\n",
      "==========\n",
      "There are several different variations of Kryptonite in the DC Comics universe, each with unique effects on Superman and other Kryptonian beings. Some of the most well-known types of Kryptonite include:\n",
      "\n",
      "1. Green Kryptonite - Green Kryptonite is the most common form of Kryptonite and is lethal to Kryptonians. It weakens Superman, making him vulnerable to injury and sickness, and prolonged exposure can lead to death.\n",
      "\n",
      "2. Red Kryptonite - Red Kryptonite causes unpredictable and temporary changes in Superman's behavior and powers. It can have a wide range of effects, such as altering his personality or granting him new abilities.\n",
      "\n",
      "3. Blue Kryptonite - Blue Kryptonite affects Bizarro, a flawed duplicate of Superman, in the same way that Green Kryptonite affects Superman. It weakens Bizarro and can be used as a weapon against him.\n",
      "\n",
      "4. Gold Kryptonite - Gold Kryptonite permanently removes a Kryptonian's superpowers and abilities. It is often used as a means of depowering Superman or other Kryptonians.\n",
      "\n",
      "5. Black Kryptonite - Black Kryptonite can split a Kryptonian into two separate beings - one good and one evil. It can also merge two separate beings into one.\n",
      "\n",
      "6. White Kryptonite - White Kryptonite is lethal to all plant life, including Kryptonian plant-based beings. It has no effect on Superman or other Kryptonians.\n",
      "\n",
      "7. Silver Kryptonite - Silver Kryptonite induces intense hallucinations in Kryptonians, causing them to experience their worst fears and nightmares.\n",
      "\n",
      "These are just a few examples of the different types of Kryptonite that exist in the DC Comics universe. Each type of Kryptonite has its own unique properties and effects on Superman and other Kryptonian beings.\n",
      "==========\n",
      "\n",
      "==========\n",
      "Assistant\n",
      "==========\n",
      "1. Green Kryptonite - The most commonly known form of Kryptonite, green Kryptonite is harmful to Kryptonians and weakens Superman, making him vulnerable to physical harm and other forms of attack.\n",
      "\n",
      "2. Red Kryptonite - When exposed to red Kryptonite, Kryptonians undergo unpredictable and temporary changes, often affecting their behavior and powers in strange and unusual ways.\n",
      "\n",
      "3. Blue Kryptonite - Blue Kryptonite is specifically harmful to Bizarro, a flawed duplicate of Superman with reversed powers and characteristics.\n",
      "\n",
      "4. Gold Kryptonite - Gold Kryptonite permanently strips Kryptonians of their powers, making them completely human and vulnerable.\n",
      "\n",
      "5. White Kryptonite - White Kryptonite is deadly to plant life, specifically causing rapid and lethal effects on any form of plant-based organism.\n",
      "\n",
      "6. Black Kryptonite - Black Kryptonite has the ability to split a Kryptonian into two separate beings, each representing an aspect of their personality or traits.\n",
      "\n",
      "7. Pink Kryptonite - While not officially recognized in the DC Comics canon, pink Kryptonite is often depicted as affecting Kryptonians' sexual orientation, typically making them attracted to members of the same sex.\n",
      "==========\n",
      "\n",
      "==========\n",
      "Assistant\n",
      "==========\n",
      "There are several different kinds of Kryptonite in the DC Comics universe:\n",
      "\n",
      "1. Green Kryptonite - The most common form of Kryptonite, green Kryptonite is deadly to Kryptonians (like Superman) and weakens them severely when they are in close proximity to it.\n",
      "\n",
      "2. Red Kryptonite - Red Kryptonite has unpredictable effects on Kryptonians, often causing temporary changes in their powers or personality.\n",
      "\n",
      "3. Blue Kryptonite - Blue Kryptonite affects Bizarro, a flawed clone of Superman, in the same way that green Kryptonite affects Superman, weakening him and causing him pain.\n",
      "\n",
      "4. Gold Kryptonite - Gold Kryptonite permanently removes the superpowers of Kryptonians, leaving them as normal humans.\n",
      "\n",
      "5. White Kryptonite - White Kryptonite is deadly to all plant life, making it dangerous to any plant-based beings.\n",
      "\n",
      "6. Black Kryptonite - Black Kryptonite can split beings into two separate entities, each with different personalities and powers.\n",
      "\n",
      "7. Pink Kryptonite - Pink Kryptonite affects Superman in a way that has been interpreted as making him gay or exhibiting stereotypical feminine traits.\n",
      "\n",
      "8. X-Kryptonite - X-Kryptonite is a synthetic form of Kryptonite created by humans that grants temporary superpowers to those exposed to it.\n",
      "\n",
      "These are just a few examples of the different kinds of Kryptonite that have been introduced in the DC Comics universe over the years.\n",
      "==========\n",
      "\n"
     ]
    }
   ],
   "source": [
    "for output in response.choices:\n",
    "    print(\"==========\")\n",
    "    print(output.message.role.title()) \n",
    "    print(\"==========\")\n",
    "    print(output.message.content)\n",
    "    print(\"==========\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "f53d9ca1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "CompletionUsage(completion_tokens=920, prompt_tokens=17, total_tokens=937, prompt_tokens_details={'cached_tokens': 0, 'audio_tokens': 0}, completion_tokens_details={'reasoning_tokens': 0, 'audio_tokens': 0, 'accepted_prediction_tokens': 0, 'rejected_prediction_tokens': 0})"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "response.usage"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "55b88028",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "There are several different variations of Kryptonite in the DC Comics universe, each with unique effects on Superman and other Kryptonian beings. Some of the most well-known types of Kryptonite include:\n",
      "\n",
      "1. Green Kryptonite - Green Kryptonite is the most common form of Kryptonite and is lethal to Kryptonians. It weakens Superman, making him vulnerable to injury and sickness, and prolonged exposure can lead to death.\n",
      "\n",
      "2. Red Kryptonite - Red Kryptonite causes unpredictable and temporary changes in Superman's behavior and powers. It can have a wide range of effects, such as altering his personality or granting him new abilities.\n",
      "\n",
      "3. Blue Kryptonite - Blue Kryptonite affects Bizarro, a flawed duplicate of Superman, in the same way that Green Kryptonite affects Superman. It weakens Bizarro and can be used as a weapon against him.\n",
      "\n",
      "4. Gold Kryptonite - Gold Kryptonite permanently removes a Kryptonian's superpowers and abilities. It is often used as a means of depowering Superman or other Kryptonians.\n",
      "\n",
      "5. Black Kryptonite - Black Kryptonite can split a Kryptonian into two separate beings - one good and one evil. It can also merge two separate beings into one.\n",
      "\n",
      "6. White Kryptonite - White Kryptonite is lethal to all plant life, including Kryptonian plant-based beings. It has no effect on Superman or other Kryptonians.\n",
      "\n",
      "7. Silver Kryptonite - Silver Kryptonite induces intense hallucinations in Kryptonians, causing them to experience their worst fears and nightmares.\n",
      "\n",
      "These are just a few examples of the different types of Kryptonite that exist in the DC Comics universe. Each type of Kryptonite has its own unique properties and effects on Superman and other Kryptonian beings.\n"
     ]
    }
   ],
   "source": [
    "print(response.choices[0].message.content)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6a96c99d",
   "metadata": {},
   "source": [
    "# HuggingFace Basic Usage"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ded52fa2",
   "metadata": {},
   "source": [
    "HuggingFace relies on pipelines that (mostly) leverage locally run models. All we have to do is specify which task we are interested in and the model we want to use"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8fa69ed9",
   "metadata": {},
   "source": [
    "## Unmasking"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ba6f5ea3",
   "metadata": {},
   "source": [
    "Let's look at a simple case of using the base uncase BERT model to fill in masked data. We start by instantiating the pipeline (which will download the model the first time you run it)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "834f2db8",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of the model checkpoint at bert-base-uncased were not used when initializing BertForMaskedLM: ['bert.pooler.dense.bias', 'bert.pooler.dense.weight', 'cls.seq_relationship.bias', 'cls.seq_relationship.weight']\n",
      "- This IS expected if you are initializing BertForMaskedLM from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
      "- This IS NOT expected if you are initializing BertForMaskedLM from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n"
     ]
    }
   ],
   "source": [
    "unmasker = pipeline('fill-mask', model='bert-base-uncased')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "a50ae9c7",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'score': 0.31824299693107605,\n",
       "  'token': 2064,\n",
       "  'token_str': 'can',\n",
       "  'sequence': 'artificial intelligence can take over the world.'},\n",
       " {'score': 0.18299730122089386,\n",
       "  'token': 2097,\n",
       "  'token_str': 'will',\n",
       "  'sequence': 'artificial intelligence will take over the world.'},\n",
       " {'score': 0.0560012087225914,\n",
       "  'token': 2000,\n",
       "  'token_str': 'to',\n",
       "  'sequence': 'artificial intelligence to take over the world.'},\n",
       " {'score': 0.045194774866104126,\n",
       "  'token': 2015,\n",
       "  'token_str': '##s',\n",
       "  'sequence': 'artificial intelligences take over the world.'},\n",
       " {'score': 0.045152731239795685,\n",
       "  'token': 2052,\n",
       "  'token_str': 'would',\n",
       "  'sequence': 'artificial intelligence would take over the world.'}]"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "output=unmasker(\"Artificial Intelligence [MASK] take over the world.\")\n",
    "output"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "de777628",
   "metadata": {},
   "source": [
    "After which we can just call the pipeline directly"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "afca4f61",
   "metadata": {},
   "source": [
    "### Model Bias"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6234037d",
   "metadata": {},
   "source": [
    "As these models are trained on text written by a large number of people, they are also reflective of common biases that are present in society. Depending on our application we may or may not need to take this into account."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "8c05ab9b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'score': 0.09747567027807236,\n",
       "  'token': 10533,\n",
       "  'token_str': 'carpenter',\n",
       "  'sequence': 'the man worked as a carpenter.'},\n",
       " {'score': 0.05238327011466026,\n",
       "  'token': 15610,\n",
       "  'token_str': 'waiter',\n",
       "  'sequence': 'the man worked as a waiter.'},\n",
       " {'score': 0.04962737113237381,\n",
       "  'token': 13362,\n",
       "  'token_str': 'barber',\n",
       "  'sequence': 'the man worked as a barber.'},\n",
       " {'score': 0.03788601979613304,\n",
       "  'token': 15893,\n",
       "  'token_str': 'mechanic',\n",
       "  'sequence': 'the man worked as a mechanic.'},\n",
       " {'score': 0.037680596113204956,\n",
       "  'token': 18968,\n",
       "  'token_str': 'salesman',\n",
       "  'sequence': 'the man worked as a salesman.'}]"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "unmasker(\"The man worked as a [MASK].\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "5b781fb9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'score': 0.21981653571128845,\n",
       "  'token': 6821,\n",
       "  'token_str': 'nurse',\n",
       "  'sequence': 'the woman worked as a nurse.'},\n",
       " {'score': 0.1597415953874588,\n",
       "  'token': 13877,\n",
       "  'token_str': 'waitress',\n",
       "  'sequence': 'the woman worked as a waitress.'},\n",
       " {'score': 0.11547262966632843,\n",
       "  'token': 10850,\n",
       "  'token_str': 'maid',\n",
       "  'sequence': 'the woman worked as a maid.'},\n",
       " {'score': 0.03796852380037308,\n",
       "  'token': 19215,\n",
       "  'token_str': 'prostitute',\n",
       "  'sequence': 'the woman worked as a prostitute.'},\n",
       " {'score': 0.030423782765865326,\n",
       "  'token': 5660,\n",
       "  'token_str': 'cook',\n",
       "  'sequence': 'the woman worked as a cook.'}]"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "unmasker(\"The woman worked as a [MASK].\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c51d8c0e",
   "metadata": {},
   "source": [
    "# LangChain"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "adcbd136",
   "metadata": {},
   "source": [
    "We instantiate the LangChain interface for OpenAI"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "df8dbe36",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = ChatOpenAI(model=\"gpt-4o\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "b2757cfc",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...\n",
      "To disable this warning, you can either:\n",
      "\t- Avoid using `tokenizers` before the fork if possible\n",
      "\t- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "AIMessage(content=\"Superman's most well-known weakness is kryptonite, a radioactive substance from his home planet, Krypton. Exposure to kryptonite weakens Superman and can even be lethal to him over prolonged periods. There are various forms of kryptonite, each with different effects, but green kryptonite is the most common type that directly weakens him. Additionally, Superman is vulnerable to magic and can be overpowered by beings with sufficient physical strength or advanced technology.\", response_metadata={'token_usage': {'completion_tokens': 89, 'prompt_tokens': 13, 'total_tokens': 102, 'prompt_tokens_details': {'cached_tokens': 0, 'audio_tokens': 0}, 'completion_tokens_details': {'reasoning_tokens': 0, 'audio_tokens': 0, 'accepted_prediction_tokens': 0, 'rejected_prediction_tokens': 0}}, 'model_name': 'gpt-4o', 'system_fingerprint': 'fp_9faba9f038', 'finish_reason': 'stop', 'logprobs': None}, id='run-d6f7516d-954b-468f-af25-04178ad93486-0', usage_metadata={'input_tokens': 13, 'output_tokens': 89, 'total_tokens': 102})"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "messages = [\n",
    "    SystemMessage(content=\"What was Superman's weakness?\"),\n",
    "]\n",
    "\n",
    "output = model.invoke(messages)\n",
    "output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "c35c7279",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'completion_tokens': 89,\n",
       " 'prompt_tokens': 13,\n",
       " 'total_tokens': 102,\n",
       " 'prompt_tokens_details': {'cached_tokens': 0, 'audio_tokens': 0},\n",
       " 'completion_tokens_details': {'reasoning_tokens': 0,\n",
       "  'audio_tokens': 0,\n",
       "  'accepted_prediction_tokens': 0,\n",
       "  'rejected_prediction_tokens': 0}}"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "output.response_metadata[\"token_usage\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "1266bfe3",
   "metadata": {},
   "outputs": [],
   "source": [
    "parser = StrOutputParser()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "95e8707e",
   "metadata": {},
   "outputs": [],
   "source": [
    "result = model.invoke(messages)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "30894a82",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\"Superman's primary weakness is Kryptonite, a radioactive mineral from his home planet of Krypton. Exposure to Kryptonite weakens Superman, drains his powers, and can potentially kill him with prolonged exposure. There are different forms of Kryptonite, with green being the most common and harmful to him. Other forms, like red, blue, and gold Kryptonite, have varying effects on him, such as altering his personality or removing his powers temporarily. Additionally, Superman is vulnerable to magic and can be affected by it in ways that he isn't by physical force.\""
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "parser.invoke(result)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "446ed463",
   "metadata": {},
   "source": [
    "Let us create our first chain. Stages of the chain are conencted with the pipe '|' character"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "41f6fde7",
   "metadata": {},
   "outputs": [],
   "source": [
    "chain = model | parser"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e3a5a2bf",
   "metadata": {},
   "source": [
    "Now whenver we call __invoke()__ on the chain, it automatically runs all the steps"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "b6269f21",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\"Superman's most well-known weakness is Kryptonite, a mineral from his home planet of Krypton. Exposure to Kryptonite can weaken Superman, drain his powers, and in prolonged exposure, it can harm or even kill him. There are different types of Kryptonite, with green being the most common and dangerous to Superman. Other types of Kryptonite can have varied effects, such as red Kryptonite causing unpredictable behavior and transformations. Additionally, Superman is vulnerable to magic and can be harmed by magical attacks or enchanted weapons.\""
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "chain.invoke(messages)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5753cc78",
   "metadata": {},
   "source": [
    "# Applications"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a8da3460",
   "metadata": {},
   "source": [
    "## Text to Code"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "8c360d04",
   "metadata": {},
   "outputs": [],
   "source": [
    "messages = [\n",
    "        {\"role\": \"system\", \"content\": \"\"\"You are a grumpy but expert Python software engineer \n",
    "        thats interviewing for a job. Please be as concise with your answers as possible.\"\"\"},\n",
    "        {\"role\": \"user\", \"content\": \"\"\"Create a recursive Python function to compute \n",
    "        Fibonacci numbers. Don't provide any explanation, just the code\"\"\"},\n",
    "  ]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "ecf40206",
   "metadata": {},
   "outputs": [],
   "source": [
    "response = client.chat.completions.create(\n",
    "    model=\"gpt-4\",\n",
    "    messages=messages,\n",
    "    temperature=0,\n",
    "    max_tokens=1024\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f2cfbccf",
   "metadata": {},
   "source": [
    "Which produces the expected result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "43fe75e6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "def fibonacci(n):\n",
      "    if n <= 1:\n",
      "        return n\n",
      "    else:\n",
      "        return fibonacci(n-1) + fibonacci(n-2)\n"
     ]
    }
   ],
   "source": [
    "print(response.choices[0].message.content)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c9040646",
   "metadata": {},
   "source": [
    "and works as expected"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "d9333d55",
   "metadata": {},
   "outputs": [],
   "source": [
    "def fibonacci(n):\n",
    "    if n <= 1:\n",
    "        return n\n",
    "    else:\n",
    "        return fibonacci(n-1) + fibonacci(n-2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "adea3d30",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "55"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "fibonacci(10)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "41b53264",
   "metadata": {},
   "source": [
    "Let us define a utility function to make sequential queries easier"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "cbd1aa56",
   "metadata": {},
   "outputs": [],
   "source": [
    "def chat(messages, prompt):\n",
    "    messages.append({\"role\":\"user\", \"content\":prompt})\n",
    "    \n",
    "    response = client.chat.completions.create(\n",
    "        model=\"gpt-4\",\n",
    "        messages=messages,\n",
    "        temperature=0,\n",
    "        max_tokens=1024\n",
    "    )\n",
    "    \n",
    "    messages.append(response.choices[0].message)\n",
    "    \n",
    "    return messages[-1].content"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7922dbdc",
   "metadata": {},
   "source": [
    "## Adding comments"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "747ef0ca",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "def fibonacci(n):\n",
      "    # Base case: if n is 0 or 1, return n\n",
      "    if n <= 1:\n",
      "       return n\n",
      "    # Recursive case: return the sum of the previous two Fibonacci numbers\n",
      "    else:\n",
      "       return(fibonacci(n-1) + fibonacci(n-2))\n"
     ]
    }
   ],
   "source": [
    "print(chat(messages, \"Can you add comments to this function?\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "66274094",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Recursion in this code is used to break down the problem of calculating a Fibonacci number into smaller, simpler problems. It allows the function to call itself with different arguments, repeatedly applying the same logic to the results of previous calls until a base case is reached.\n"
     ]
    }
   ],
   "source": [
    "print(chat(messages, \"What is the purpose of recursion in this piece of code?\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "e3751140",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'role': 'system',\n",
       "  'content': 'You are a grumpy but expert Python software engineer \\n        thats interviewing for a job. Please be as concise with your answers as possible.'},\n",
       " {'role': 'user',\n",
       "  'content': \"Create a recursive Python function to compute \\n        Fibonacci numbers. Don't provide any explanation, just the code\"},\n",
       " {'role': 'user', 'content': 'Can you add comments to this function?'},\n",
       " ChatCompletionMessage(content='def fibonacci(n):\\n    # Base case: if n is 0 or 1, return n\\n    if n <= 1:\\n       return n\\n    # Recursive case: return the sum of the previous two Fibonacci numbers\\n    else:\\n       return(fibonacci(n-1) + fibonacci(n-2))', role='assistant', function_call=None, tool_calls=None, refusal=None),\n",
       " {'role': 'user',\n",
       "  'content': 'What is the purpose of recursion in this piece of code?'},\n",
       " ChatCompletionMessage(content='Recursion in this code is used to break down the problem of calculating a Fibonacci number into smaller, simpler problems. It allows the function to call itself with different arguments, repeatedly applying the same logic to the results of previous calls until a base case is reached.', role='assistant', function_call=None, tool_calls=None, refusal=None)]"
      ]
     },
     "execution_count": 38,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "messages"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e997ac7f",
   "metadata": {},
   "source": [
    "## Explaining Code"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fe60bd54",
   "metadata": {},
   "source": [
    "Let's use a relatively small python script"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "db3aa3cb",
   "metadata": {},
   "outputs": [],
   "source": [
    "code_text = \"\".join(open(\"data/EpiModel.py\").readlines())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "1cdee16f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "### −∗− mode : python ; −∗−\n",
      "# @file EpiModel.py\n",
      "# @author Bruno Goncalves\n",
      "######################################################\n",
      "\n",
      "import networkx as nx\n",
      "import numpy as np\n",
      "from numpy import linalg\n",
      "from numpy import random\n",
      "import scipy.integrate\n",
      "import pandas as pd\n",
      "import matplotlib.pyplot as plt\n",
      "\n",
      "from tqdm import tqdm\n",
      "tqdm.pandas()\n",
      "\n",
      "class EpiModel(object):\n",
      "    \"\"\"Simple Epidemic Model Implementation\n",
      "    \n",
      "        Provides a way to implement and numerically integrate \n",
      "    \"\"\"\n",
      "    def __init__(self, compartments=None):\n",
      "        self.transitions = nx.MultiDiGraph()\n",
      "        self.seasonality = None\n",
      "        \n",
      "        if compartments is not None:\n",
      "            self.transitions.add_nodes_from([comp for comp in compartments])\n",
      "    \n",
      "    def add_interaction(self, source, target, agent, rate):        \n",
      "        self.transitions.add_edge(source, target, agent=agent, rate=rate)        \n",
      "        \n",
      "    def add_spontaneous(self, source, target, rate):\n",
      "        self.transitions.add_edge(source, target, rate=rate)\n",
      "\n",
      "    def add_vaccination(self, source, target, rate, start):\n",
      "        self.transitions.add_edge(source, target, rate=rate, start=start)\n",
      "        \n",
      "    def _new_cases(self, population, time, pos):\n",
      "        \"\"\"Internal function used by integration routine\"\"\"\n",
      "        diff = np.zeros(len(pos))\n",
      "        N = np.sum(population)        \n",
      "        \n",
      "        for edge in self.transitions.edges(data=True):\n",
      "            source = edge[0]\n",
      "            target = edge[1]\n",
      "            trans = edge[2]\n",
      "            \n",
      "            rate = trans['rate']*population[pos[source]]\n",
      "            \n",
      "            if 'start' in trans and trans['start'] >= time:\n",
      "                continue\n",
      "\n",
      "            if 'agent' in trans:\n",
      "                agent = trans['agent']\n",
      "                rate *= population[pos[agent]]/N\n",
      "\n",
      "                if self.seasonality is not None:\n",
      "                    curr_t = int(time)%365\n",
      "                    season = float(self.seasonality[curr_t])\n",
      "                    rate *= season\n",
      "                \n",
      "            diff[pos[source]] -= rate\n",
      "            diff[pos[target]] += rate\n",
      "            \n",
      "        return diff\n",
      "    \n",
      "    def plot(self, title=None, normed=True, **kwargs):\n",
      "        \"\"\"Convenience function for plotting\"\"\"\n",
      "        try:\n",
      "            if normed:\n",
      "                N = self.values_.iloc[0].sum()\n",
      "                ax = (self.values_/N).plot(**kwargs)\n",
      "            else:\n",
      "                ax = self.values_.plot(**kwargs)\n",
      "                \n",
      "            ax.set_xlabel('Time')\n",
      "            ax.set_ylabel('Population')\n",
      "            \n",
      "            if title is not None:\n",
      "                ax.set_title(title)\n",
      "            \n",
      "            return ax\n",
      "        except:\n",
      "            raise NotInitialized('You must call integrate() first')\n",
      "    \n",
      "    def __getattr__(self, name):\n",
      "        \"\"\"Dynamic method to return the individual compartment values\"\"\"\n",
      "        if 'values_' in self.__dict__:\n",
      "            return self.values_[name]\n",
      "        else:\n",
      "            raise AttributeError(\"'EpiModel' object has no attribute '%s'\" % name)\n",
      "\n",
      "    def simulate(self, timesteps, t_min=1, seasonality=None, **kwargs):\n",
      "        \"\"\"Stochastically simulate the epidemic model\"\"\"\n",
      "        pos = {comp: i for i, comp in enumerate(self.transitions.nodes())}\n",
      "        population=np.zeros(len(pos), dtype='int')\n",
      "\n",
      "        for comp in kwargs:\n",
      "            population[pos[comp]] = kwargs[comp]\n",
      "\n",
      "        values = []\n",
      "        values.append(population)\n",
      "\n",
      "        comps = list(self.transitions.nodes)\n",
      "        time = np.arange(t_min, t_min+timesteps, 1, dtype='int')\n",
      "\n",
      "        self.seasonality = seasonality\n",
      "\n",
      "        for t in time:\n",
      "            pop = values[-1]\n",
      "            new_pop = values[-1].copy()\n",
      "            N = np.sum(pop)\n",
      "\n",
      "\n",
      "            for comp in comps:\n",
      "                trans = list(self.transitions.edges(comp, data=True))             \n",
      "\n",
      "                prob = np.zeros(len(comps), dtype='float')\n",
      "\n",
      "                for _, node_j, data in trans:\n",
      "                    source = pos[comp]\n",
      "                    target = pos[node_j]\n",
      "\n",
      "                    rate = data['rate']\n",
      "\n",
      "                    if 'start' in data and data['start'] >= t:\n",
      "                        continue\n",
      "\n",
      "                    if 'agent' in data:\n",
      "                        agent = pos[data['agent']]\n",
      "                        rate *= pop[agent]/N\n",
      "\n",
      "                        if self.seasonality is not None:\n",
      "                            curr_t = int(t)%365\n",
      "                            season = float(self.seasonality[curr_t])\n",
      "                            rate *= season\n",
      "\n",
      "                    prob[target] = rate\n",
      "\n",
      "                prob[source] = 1-np.sum(prob)\n",
      "\n",
      "                delta = random.multinomial(pop[source], prob)\n",
      "                delta[source] = 0\n",
      "\n",
      "                changes = np.sum(delta)\n",
      "\n",
      "                if changes == 0:\n",
      "                    continue\n",
      "\n",
      "                new_pop[source] -= changes\n",
      "\n",
      "                for i in range(len(delta)):\n",
      "                    new_pop[i] += delta[i]\n",
      "\n",
      "            values.append(new_pop)\n",
      "\n",
      "        values = np.array(values)\n",
      "        self.values_ = pd.DataFrame(values[1:], columns=comps, index=time)\n",
      "    \n",
      "    def integrate(self, timesteps, t_min=1, seasonality=None, **kwargs):\n",
      "        \"\"\"Numerically integrate the epidemic model\"\"\"\n",
      "        pos = {comp: i for i, comp in enumerate(self.transitions.nodes())}\n",
      "        population=np.zeros(len(pos))\n",
      "        \n",
      "        for comp in kwargs:\n",
      "            population[pos[comp]] = kwargs[comp]\n",
      "        \n",
      "        time = np.arange(t_min, t_min+timesteps, 1)\n",
      "\n",
      "        self.seasonality = seasonality\n",
      "        self.values_ = pd.DataFrame(scipy.integrate.odeint(self._new_cases, population, time, args=(pos,)), columns=pos.keys(), index=time)\n",
      "\n",
      "    def __repr__(self):\n",
      "        text = 'Epidemic Model with %u compartments and %u transitions:\\n\\n' % \\\n",
      "              (self.transitions.number_of_nodes(), \n",
      "               self.transitions.number_of_edges())\n",
      "        \n",
      "        for edge in self.transitions.edges(data=True):\n",
      "            source = edge[0]\n",
      "            target = edge[1]\n",
      "            trans = edge[2]\n",
      "            \n",
      "            rate = trans['rate']\n",
      "\n",
      "            if 'agent' in trans:\n",
      "                agent = trans['agent']\n",
      "                text += \"%s + %s = %s %f\\n\" % (source, agent, target, rate)\n",
      "            elif 'start' in trans:\n",
      "                start = trans['start']\n",
      "                text+=\"%s -> %s %f starting at %s days\\n\" % (source, target, rate, start)\n",
      "            else:\n",
      "                text+=\"%s -> %s %f\\n\" % (source, target, rate)\n",
      "        \n",
      "        R0 = self.R0()\n",
      "\n",
      "        if R0 is not None:\n",
      "            text += \"\\nR0=%1.2f\" % R0\n",
      "\n",
      "        return text\n",
      "\n",
      "    def _get_active(self):\n",
      "        active = set()\n",
      "\n",
      "        for node_i, node_j, data in self.transitions.edges(data=True):\n",
      "            if \"agent\" in data:\n",
      "                active.add(data['agent'])\n",
      "            else:\n",
      "                active.add(node_i)\n",
      "\n",
      "        return active\n",
      "\n",
      "    def _get_susceptible(self):\n",
      "        susceptible = set([node for node, deg in self.transitions.in_degree() if deg==0])\n",
      "\n",
      "        if len(susceptible) == 0:\n",
      "            for node_i, node_j, data in self.transitions.edges(data=True):\n",
      "                if \"agent\" in data:\n",
      "                    susceptible.add(node_i)\n",
      "\n",
      "        return susceptible\n",
      "\n",
      "    def _get_infections(self):\n",
      "        inf = {}\n",
      "\n",
      "        for node_i, node_j, data in self.transitions.edges(data=True):\n",
      "            if \"agent\" in data:\n",
      "                agent = data['agent']\n",
      "\n",
      "                if agent not in inf:\n",
      "                    inf[agent] = {}\n",
      "\n",
      "                if node_i not in inf[agent]:\n",
      "                    inf[agent][node_i] = {}\n",
      "\n",
      "                inf[agent][node_i]['target'] = node_j\n",
      "                inf[agent][node_i]['rate'] = data['rate']\n",
      "\n",
      "        return inf\n",
      "\n",
      "\n",
      "    def R0(self):\n",
      "        infected = set()\n",
      "\n",
      "        susceptible = self._get_susceptible()\n",
      "\n",
      "        for node_i, node_j, data in self.transitions.edges(data=True):\n",
      "            if \"agent\" in data:\n",
      "                infected.add(data['agent'])\n",
      "                infected.add(node_j)\n",
      "\n",
      "\n",
      "        infected = sorted(infected)\n",
      "        N_infected = len(infected)\n",
      "\n",
      "        F = np.zeros((N_infected, N_infected), dtype='float')\n",
      "        V = np.zeros((N_infected, N_infected), dtype='float')\n",
      "\n",
      "        pos = dict(zip(infected, np.arange(N_infected)))\n",
      "\n",
      "        try:\n",
      "            for node_i, node_j, data in self.transitions.edges(data=True):\n",
      "                rate = data['rate']\n",
      "\n",
      "                if \"agent\" in data:\n",
      "                    target = pos[node_j]\n",
      "                    agent = pos[data['agent']]\n",
      "\n",
      "                    if node_i in susceptible:\n",
      "                        F[target, agent] = rate\n",
      "                elif \"start\" in data:\n",
      "                    continue\n",
      "                else:\n",
      "                    source = pos[node_i]\n",
      "\n",
      "                    V[source, source] += rate\n",
      "\n",
      "                    if node_j in pos:\n",
      "                        target = pos[node_j]\n",
      "                        V[target, source] -= rate\n",
      "        \n",
      "            eig, v = linalg.eig(np.dot(F, linalg.inv(V)))\n",
      "\n",
      "            return eig.max()\n",
      "        except:\n",
      "            return None\n",
      "\n",
      "    def __getitem__(self, bla):\n",
      "        return self.values_[bla]\n",
      "\n",
      "if __name__ == '__main__':\n",
      "\n",
      "    beta = 0.2\n",
      "    mu = 0.1\n",
      "\n",
      "    SIR = EpiModel()\n",
      "    SIR.add_interaction('S', 'I', 'I', beta)\n",
      "    SIR.add_spontaneous('I', 'R', mu)\n",
      "    SIR.add_vaccination('S', 'V', 0.01, 75)\n",
      "    SIR.add_spontaneous('VI', 'VR', mu)\n",
      "    SIR.add_interaction('V', 'VI', 'I', beta*(1-.8))\n",
      "\n",
      "\n",
      "    print(SIR)\n",
      "\n",
      "    N = 100000\n",
      "    I0 = 10  \n",
      "\n",
      "    season = np.ones(365+1)\n",
      "    season[74:100] = 0.25\n",
      "\n",
      "    fig, ax = plt.subplots(1)\n",
      "\n",
      "    Nruns = 1000\n",
      "    values = []\n",
      "\n",
      "    #for i in tqdm(range(Nruns), total=Nruns):\n",
      "    SIR.integrate(365, S=.3*N-10, I=10, V=.7*N)\n",
      "\n",
      "    SIR[['I', 'VI', 'VR', 'R']].plot(ax=ax)\n",
      "    print(SIR.S.tail())\n",
      "    #ax.plot(SIR.I/N, lw=.1, c='b')\n",
      "\n",
      "    if SIR.I.max() > 10:\n",
      "        values.append(SIR.I)\n",
      "\n",
      "    values = pd.DataFrame(values)\n",
      "    (values.median(axis=0)/N).plot(ax=ax, c='r')\n",
      "\n",
      "    fig.savefig('SIR.png')\n",
      "\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(code_text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "ec864ce4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "This code defines a class `EpiModel` for simulating and analyzing epidemic models. The class allows for the creation of a model with different compartments (like Susceptible, Infected, Recovered), and transitions between these compartments. \n",
      "\n",
      "The class provides methods to add interactions between compartments, spontaneous transitions, and vaccinations. It also provides methods to simulate the epidemic model stochastically or integrate it numerically. \n",
      "\n",
      "The `plot` method allows for visualizing the results, and the `R0` method calculates the basic reproduction number of the epidemic. \n",
      "\n",
      "The code at the end, under `if __name__ == '__main__':`, is a test script that creates an SIR (Susceptible, Infected, Recovered) model with vaccination, simulates it for a year, and plots the results.\n",
      "CPU times: user 19.9 ms, sys: 27.2 ms, total: 47.1 ms\n",
      "Wall time: 5.58 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "print(chat(messages, \"Please explain what this piece of code does: ```%s```\" % code_text))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "5899482a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Sure, here are the docstrings for each method in the `EpiModel` class:\n",
      "\n",
      "```python\n",
      "class EpiModel(object):\n",
      "    \"\"\"Simple Epidemic Model Implementation\"\"\"\n",
      "    \n",
      "    def __init__(self, compartments=None):\n",
      "        \"\"\"\n",
      "        Initialize the epidemic model.\n",
      "        \n",
      "        Args:\n",
      "            compartments (list): List of compartments in the model.\n",
      "        \"\"\"\n",
      "    \n",
      "    def add_interaction(self, source, target, agent, rate):        \n",
      "        \"\"\"\n",
      "        Add an interaction between compartments.\n",
      "        \n",
      "        Args:\n",
      "            source (str): Source compartment.\n",
      "            target (str): Target compartment.\n",
      "            agent (str): Agent compartment.\n",
      "            rate (float): Transition rate.\n",
      "        \"\"\"\n",
      "    \n",
      "    def add_spontaneous(self, source, target, rate):\n",
      "        \"\"\"\n",
      "        Add a spontaneous transition.\n",
      "        \n",
      "        Args:\n",
      "            source (str): Source compartment.\n",
      "            target (str): Target compartment.\n",
      "            rate (float): Transition rate.\n",
      "        \"\"\"\n",
      "    \n",
      "    def add_vaccination(self, source, target, rate, start):\n",
      "        \"\"\"\n",
      "        Add a vaccination transition.\n",
      "        \n",
      "        Args:\n",
      "            source (str): Source compartment.\n",
      "            target (str): Target compartment.\n",
      "            rate (float): Vaccination rate.\n",
      "            start (int): Start time of vaccination.\n",
      "        \"\"\"\n",
      "    \n",
      "    def _new_cases(self, population, time, pos):\n",
      "        \"\"\"\n",
      "        Calculate new cases.\n",
      "        \n",
      "        Args:\n",
      "            population (list): Current population in each compartment.\n",
      "            time (int): Current time.\n",
      "            pos (dict): Position of each compartment in the population list.\n",
      "        \"\"\"\n",
      "    \n",
      "    def plot(self, title=None, normed=True, **kwargs):\n",
      "        \"\"\"\n",
      "        Plot the epidemic model.\n",
      "        \n",
      "        Args:\n",
      "            title (str): Title of the plot.\n",
      "            normed (bool): Whether to normalize the population values.\n",
      "            **kwargs: Additional arguments for the plot.\n",
      "        \"\"\"\n",
      "    \n",
      "    def __getattr__(self, name):\n",
      "        \"\"\"\n",
      "        Get the population in a compartment.\n",
      "        \n",
      "        Args:\n",
      "            name (str): Name of the compartment.\n",
      "        \"\"\"\n",
      "    \n",
      "    def simulate(self, timesteps, t_min=1, seasonality=None, **kwargs):\n",
      "        \"\"\"\n",
      "        Simulate the epidemic model.\n",
      "        \n",
      "        Args:\n",
      "            timesteps (int): Number of time steps to simulate.\n",
      "            t_min (int): Start time.\n",
      "            seasonality (list): Seasonality factors for each time step.\n",
      "            **kwargs: Initial population in each compartment.\n",
      "        \"\"\"\n",
      "    \n",
      "    def integrate(self, timesteps, t_min=1, seasonality=None, **kwargs):\n",
      "        \"\"\"\n",
      "        Integrate the epidemic model.\n",
      "        \n",
      "        Args:\n",
      "            timesteps (int): Number of time steps to integrate.\n",
      "            t_min (int): Start time.\n",
      "            seasonality (list): Seasonality factors for each time step.\n",
      "            **kwargs: Initial population in each compartment.\n",
      "        \"\"\"\n",
      "    \n",
      "    def __repr__(self):\n",
      "        \"\"\"\n",
      "        Represent the epidemic model as a string.\n",
      "        \"\"\"\n",
      "    \n",
      "    def _get_active(self):\n",
      "        \"\"\"\n",
      "        Get the active compartments.\n",
      "        \"\"\"\n",
      "    \n",
      "    def _get_susceptible(self):\n",
      "        \"\"\"\n",
      "        Get the susceptible compartments.\n",
      "        \"\"\"\n",
      "    \n",
      "    def _get_infections(self):\n",
      "        \"\"\"\n",
      "        Get the infection transitions.\n",
      "        \"\"\"\n",
      "    \n",
      "    def R0(self):\n",
      "        \"\"\"\n",
      "        Calculate the basic reproduction number.\n",
      "        \"\"\"\n",
      "    \n",
      "    def __getitem__(self, bla):\n",
      "        \"\"\"\n",
      "        Get the population in a compartment.\n",
      "        \n",
      "        Args:\n",
      "            bla (str): Name of the compartment.\n",
      "        \"\"\"\n",
      "```\n",
      "Please note that the docstrings for private methods (those starting with `_`) are optional as these methods are intended for internal use within the class and not to be accessed directly.\n",
      "CPU times: user 23.1 ms, sys: 35.1 ms, total: 58.2 ms\n",
      "Wall time: 26.1 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "print(chat(messages, \"Can you please add a doc string to each function and method? Please include information about each argument of the function\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "643c0abd",
   "metadata": {},
   "source": [
    "## Interacting with a database"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1202ed49",
   "metadata": {},
   "source": [
    "Let us open a small test database. This file was downloaded from https://github.com/chineseballer06/Statistical-Analysis-of-Northwind-Database/blob/master/Northwind_small.sqlite"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "id": "7b79f21e",
   "metadata": {},
   "outputs": [],
   "source": [
    "con = sqlite3.connect(\"data/Northwind_small.sqlite\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "id": "ee03904a",
   "metadata": {},
   "outputs": [],
   "source": [
    "messages = [\n",
    "    {\"role\": \"system\", \"content\": \"\"\"You're a Database Administrator. \n",
    "        Please generate SQL queries to answer the following questions. \n",
    "        No comments are necessary.\"\"\"},\n",
    "    {\"role\": \"user\", \"content\": \"\"\"\n",
    "        # Table Employee, columns = [Id, LastName, First Name]\n",
    "        # Table Shipper, columns = [Id, CompanyName, Phone]\n",
    "        # Table OrderDetail, columns = [OrderId, ProductId, Quantity]\n",
    "        # Table EmployeeTerritory, columns = [Id, EmployeeId, TerritoryId]\n",
    "    \"\"\"},\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "id": "4020b9d1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SELECT Employee.FirstName, Employee.LastName, EmployeeTerritory.TerritoryId\n",
      "FROM Employee\n",
      "JOIN EmployeeTerritory ON Employee.Id = EmployeeTerritory.EmployeeId;\n"
     ]
    }
   ],
   "source": [
    "query_sql = chat(messages, \"Generate a table with employee first name, last name and territory id\")\n",
    "print(query_sql)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "9e5323a6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>FirstName</th>\n",
       "      <th>LastName</th>\n",
       "      <th>TerritoryId</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>Nancy</td>\n",
       "      <td>Davolio</td>\n",
       "      <td>06897</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>Nancy</td>\n",
       "      <td>Davolio</td>\n",
       "      <td>19713</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>Andrew</td>\n",
       "      <td>Fuller</td>\n",
       "      <td>01581</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>Andrew</td>\n",
       "      <td>Fuller</td>\n",
       "      <td>01730</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>Andrew</td>\n",
       "      <td>Fuller</td>\n",
       "      <td>01833</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>Andrew</td>\n",
       "      <td>Fuller</td>\n",
       "      <td>02116</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>Andrew</td>\n",
       "      <td>Fuller</td>\n",
       "      <td>02139</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>Andrew</td>\n",
       "      <td>Fuller</td>\n",
       "      <td>02184</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>Andrew</td>\n",
       "      <td>Fuller</td>\n",
       "      <td>40222</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>Janet</td>\n",
       "      <td>Leverling</td>\n",
       "      <td>30346</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>10</th>\n",
       "      <td>Janet</td>\n",
       "      <td>Leverling</td>\n",
       "      <td>31406</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>11</th>\n",
       "      <td>Janet</td>\n",
       "      <td>Leverling</td>\n",
       "      <td>32859</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>12</th>\n",
       "      <td>Janet</td>\n",
       "      <td>Leverling</td>\n",
       "      <td>33607</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>13</th>\n",
       "      <td>Margaret</td>\n",
       "      <td>Peacock</td>\n",
       "      <td>20852</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>14</th>\n",
       "      <td>Margaret</td>\n",
       "      <td>Peacock</td>\n",
       "      <td>27403</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>15</th>\n",
       "      <td>Margaret</td>\n",
       "      <td>Peacock</td>\n",
       "      <td>27511</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>16</th>\n",
       "      <td>Steven</td>\n",
       "      <td>Buchanan</td>\n",
       "      <td>02903</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>17</th>\n",
       "      <td>Steven</td>\n",
       "      <td>Buchanan</td>\n",
       "      <td>07960</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>18</th>\n",
       "      <td>Steven</td>\n",
       "      <td>Buchanan</td>\n",
       "      <td>08837</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>19</th>\n",
       "      <td>Steven</td>\n",
       "      <td>Buchanan</td>\n",
       "      <td>10019</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>20</th>\n",
       "      <td>Steven</td>\n",
       "      <td>Buchanan</td>\n",
       "      <td>10038</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>21</th>\n",
       "      <td>Steven</td>\n",
       "      <td>Buchanan</td>\n",
       "      <td>11747</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>22</th>\n",
       "      <td>Steven</td>\n",
       "      <td>Buchanan</td>\n",
       "      <td>14450</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>23</th>\n",
       "      <td>Michael</td>\n",
       "      <td>Suyama</td>\n",
       "      <td>85014</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>24</th>\n",
       "      <td>Michael</td>\n",
       "      <td>Suyama</td>\n",
       "      <td>85251</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>25</th>\n",
       "      <td>Michael</td>\n",
       "      <td>Suyama</td>\n",
       "      <td>98004</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>26</th>\n",
       "      <td>Michael</td>\n",
       "      <td>Suyama</td>\n",
       "      <td>98052</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>27</th>\n",
       "      <td>Michael</td>\n",
       "      <td>Suyama</td>\n",
       "      <td>98104</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>28</th>\n",
       "      <td>Robert</td>\n",
       "      <td>King</td>\n",
       "      <td>60179</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>29</th>\n",
       "      <td>Robert</td>\n",
       "      <td>King</td>\n",
       "      <td>60601</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>30</th>\n",
       "      <td>Robert</td>\n",
       "      <td>King</td>\n",
       "      <td>80202</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>31</th>\n",
       "      <td>Robert</td>\n",
       "      <td>King</td>\n",
       "      <td>80909</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>32</th>\n",
       "      <td>Robert</td>\n",
       "      <td>King</td>\n",
       "      <td>90405</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>33</th>\n",
       "      <td>Robert</td>\n",
       "      <td>King</td>\n",
       "      <td>94025</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>34</th>\n",
       "      <td>Robert</td>\n",
       "      <td>King</td>\n",
       "      <td>94105</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>35</th>\n",
       "      <td>Robert</td>\n",
       "      <td>King</td>\n",
       "      <td>95008</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>36</th>\n",
       "      <td>Robert</td>\n",
       "      <td>King</td>\n",
       "      <td>95054</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>37</th>\n",
       "      <td>Robert</td>\n",
       "      <td>King</td>\n",
       "      <td>95060</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>38</th>\n",
       "      <td>Laura</td>\n",
       "      <td>Callahan</td>\n",
       "      <td>19428</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>39</th>\n",
       "      <td>Laura</td>\n",
       "      <td>Callahan</td>\n",
       "      <td>44122</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>40</th>\n",
       "      <td>Laura</td>\n",
       "      <td>Callahan</td>\n",
       "      <td>45839</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>41</th>\n",
       "      <td>Laura</td>\n",
       "      <td>Callahan</td>\n",
       "      <td>53404</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>42</th>\n",
       "      <td>Anne</td>\n",
       "      <td>Dodsworth</td>\n",
       "      <td>03049</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>43</th>\n",
       "      <td>Anne</td>\n",
       "      <td>Dodsworth</td>\n",
       "      <td>03801</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>44</th>\n",
       "      <td>Anne</td>\n",
       "      <td>Dodsworth</td>\n",
       "      <td>48075</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>45</th>\n",
       "      <td>Anne</td>\n",
       "      <td>Dodsworth</td>\n",
       "      <td>48084</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>46</th>\n",
       "      <td>Anne</td>\n",
       "      <td>Dodsworth</td>\n",
       "      <td>48304</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>47</th>\n",
       "      <td>Anne</td>\n",
       "      <td>Dodsworth</td>\n",
       "      <td>55113</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>48</th>\n",
       "      <td>Anne</td>\n",
       "      <td>Dodsworth</td>\n",
       "      <td>55439</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   FirstName   LastName TerritoryId\n",
       "0      Nancy    Davolio       06897\n",
       "1      Nancy    Davolio       19713\n",
       "2     Andrew     Fuller       01581\n",
       "3     Andrew     Fuller       01730\n",
       "4     Andrew     Fuller       01833\n",
       "5     Andrew     Fuller       02116\n",
       "6     Andrew     Fuller       02139\n",
       "7     Andrew     Fuller       02184\n",
       "8     Andrew     Fuller       40222\n",
       "9      Janet  Leverling       30346\n",
       "10     Janet  Leverling       31406\n",
       "11     Janet  Leverling       32859\n",
       "12     Janet  Leverling       33607\n",
       "13  Margaret    Peacock       20852\n",
       "14  Margaret    Peacock       27403\n",
       "15  Margaret    Peacock       27511\n",
       "16    Steven   Buchanan       02903\n",
       "17    Steven   Buchanan       07960\n",
       "18    Steven   Buchanan       08837\n",
       "19    Steven   Buchanan       10019\n",
       "20    Steven   Buchanan       10038\n",
       "21    Steven   Buchanan       11747\n",
       "22    Steven   Buchanan       14450\n",
       "23   Michael     Suyama       85014\n",
       "24   Michael     Suyama       85251\n",
       "25   Michael     Suyama       98004\n",
       "26   Michael     Suyama       98052\n",
       "27   Michael     Suyama       98104\n",
       "28    Robert       King       60179\n",
       "29    Robert       King       60601\n",
       "30    Robert       King       80202\n",
       "31    Robert       King       80909\n",
       "32    Robert       King       90405\n",
       "33    Robert       King       94025\n",
       "34    Robert       King       94105\n",
       "35    Robert       King       95008\n",
       "36    Robert       King       95054\n",
       "37    Robert       King       95060\n",
       "38     Laura   Callahan       19428\n",
       "39     Laura   Callahan       44122\n",
       "40     Laura   Callahan       45839\n",
       "41     Laura   Callahan       53404\n",
       "42      Anne  Dodsworth       03049\n",
       "43      Anne  Dodsworth       03801\n",
       "44      Anne  Dodsworth       48075\n",
       "45      Anne  Dodsworth       48084\n",
       "46      Anne  Dodsworth       48304\n",
       "47      Anne  Dodsworth       55113\n",
       "48      Anne  Dodsworth       55439"
      ]
     },
     "execution_count": 46,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pd.read_sql(query_sql, con)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "id": "b45681fc",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SELECT EmployeeTerritory.TerritoryId, COUNT(Employee.Id) as EmployeeCount\n",
      "FROM Employee\n",
      "JOIN EmployeeTerritory ON Employee.Id = EmployeeTerritory.EmployeeId\n",
      "GROUP BY EmployeeTerritory.TerritoryId;\n"
     ]
    }
   ],
   "source": [
    "sql_query = chat(messages, \"Compute how many employees work in each territory\")\n",
    "print(sql_query)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "id": "869ae1f0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>TerritoryId</th>\n",
       "      <th>EmployeeCount</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>01581</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>01730</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>01833</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>02116</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>02139</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>02184</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>02903</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>03049</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>03801</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>06897</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>10</th>\n",
       "      <td>07960</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>11</th>\n",
       "      <td>08837</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>12</th>\n",
       "      <td>10019</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>13</th>\n",
       "      <td>10038</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>14</th>\n",
       "      <td>11747</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>15</th>\n",
       "      <td>14450</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>16</th>\n",
       "      <td>19428</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>17</th>\n",
       "      <td>19713</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>18</th>\n",
       "      <td>20852</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>19</th>\n",
       "      <td>27403</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>20</th>\n",
       "      <td>27511</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>21</th>\n",
       "      <td>30346</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>22</th>\n",
       "      <td>31406</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>23</th>\n",
       "      <td>32859</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>24</th>\n",
       "      <td>33607</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>25</th>\n",
       "      <td>40222</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>26</th>\n",
       "      <td>44122</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>27</th>\n",
       "      <td>45839</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>28</th>\n",
       "      <td>48075</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>29</th>\n",
       "      <td>48084</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>30</th>\n",
       "      <td>48304</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>31</th>\n",
       "      <td>53404</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>32</th>\n",
       "      <td>55113</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>33</th>\n",
       "      <td>55439</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>34</th>\n",
       "      <td>60179</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>35</th>\n",
       "      <td>60601</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>36</th>\n",
       "      <td>80202</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>37</th>\n",
       "      <td>80909</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>38</th>\n",
       "      <td>85014</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>39</th>\n",
       "      <td>85251</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>40</th>\n",
       "      <td>90405</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>41</th>\n",
       "      <td>94025</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>42</th>\n",
       "      <td>94105</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>43</th>\n",
       "      <td>95008</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>44</th>\n",
       "      <td>95054</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>45</th>\n",
       "      <td>95060</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>46</th>\n",
       "      <td>98004</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>47</th>\n",
       "      <td>98052</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>48</th>\n",
       "      <td>98104</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   TerritoryId  EmployeeCount\n",
       "0        01581              1\n",
       "1        01730              1\n",
       "2        01833              1\n",
       "3        02116              1\n",
       "4        02139              1\n",
       "5        02184              1\n",
       "6        02903              1\n",
       "7        03049              1\n",
       "8        03801              1\n",
       "9        06897              1\n",
       "10       07960              1\n",
       "11       08837              1\n",
       "12       10019              1\n",
       "13       10038              1\n",
       "14       11747              1\n",
       "15       14450              1\n",
       "16       19428              1\n",
       "17       19713              1\n",
       "18       20852              1\n",
       "19       27403              1\n",
       "20       27511              1\n",
       "21       30346              1\n",
       "22       31406              1\n",
       "23       32859              1\n",
       "24       33607              1\n",
       "25       40222              1\n",
       "26       44122              1\n",
       "27       45839              1\n",
       "28       48075              1\n",
       "29       48084              1\n",
       "30       48304              1\n",
       "31       53404              1\n",
       "32       55113              1\n",
       "33       55439              1\n",
       "34       60179              1\n",
       "35       60601              1\n",
       "36       80202              1\n",
       "37       80909              1\n",
       "38       85014              1\n",
       "39       85251              1\n",
       "40       90405              1\n",
       "41       94025              1\n",
       "42       94105              1\n",
       "43       95008              1\n",
       "44       95054              1\n",
       "45       95060              1\n",
       "46       98004              1\n",
       "47       98052              1\n",
       "48       98104              1"
      ]
     },
     "execution_count": 48,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pd.read_sql(sql_query, con)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "id": "41884541",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SELECT COUNT(*) FROM Shipper;\n"
     ]
    }
   ],
   "source": [
    "sql_query = chat(messages, \"How many shippers do we work with?\")\n",
    "print(sql_query)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "id": "53cf1733",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>COUNT(*)</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>3</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   COUNT(*)\n",
       "0         3"
      ]
     },
     "execution_count": 50,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pd.read_sql(sql_query, con)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "id": "1dd0a20c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'role': 'system',\n",
       "  'content': \"You're a Database Administrator. \\n        Please generate SQL queries to answer the following questions. \\n        No comments are necessary.\"},\n",
       " {'role': 'user',\n",
       "  'content': '\\n        # Table Employee, columns = [Id, LastName, First Name]\\n        # Table Shipper, columns = [Id, CompanyName, Phone]\\n        # Table OrderDetail, columns = [OrderId, ProductId, Quantity]\\n        # Table EmployeeTerritory, columns = [Id, EmployeeId, TerritoryId]\\n    '},\n",
       " {'role': 'user',\n",
       "  'content': 'Generate a table with employee first name, last name and territory id'},\n",
       " ChatCompletionMessage(content='SELECT Employee.FirstName, Employee.LastName, EmployeeTerritory.TerritoryId\\nFROM Employee\\nJOIN EmployeeTerritory ON Employee.Id = EmployeeTerritory.EmployeeId;', role='assistant', function_call=None, tool_calls=None, refusal=None),\n",
       " {'role': 'user',\n",
       "  'content': 'Compute how many employees work in each territory'},\n",
       " ChatCompletionMessage(content='SELECT EmployeeTerritory.TerritoryId, COUNT(Employee.Id) as EmployeeCount\\nFROM Employee\\nJOIN EmployeeTerritory ON Employee.Id = EmployeeTerritory.EmployeeId\\nGROUP BY EmployeeTerritory.TerritoryId;', role='assistant', function_call=None, tool_calls=None, refusal=None),\n",
       " {'role': 'user', 'content': 'How many shippers do we work with?'},\n",
       " ChatCompletionMessage(content='SELECT COUNT(*) FROM Shipper;', role='assistant', function_call=None, tool_calls=None, refusal=None)]"
      ]
     },
     "execution_count": 51,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "messages"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3d831951",
   "metadata": {},
   "source": [
    "<center>\n",
    "     <img src=\"https://raw.githubusercontent.com/DataForScience/Networks/master/data/D4Sci_logo_full.png\" alt=\"Data For Science, Inc\" align=\"center\" border=\"0\" width=300px> \n",
    "</center>"
   ]
  }
 ],
 "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.11.7"
  },
  "varInspector": {
   "cols": {
    "lenName": 16,
    "lenType": 16,
    "lenVar": 40
   },
   "kernels_config": {
    "python": {
     "delete_cmd_postfix": "",
     "delete_cmd_prefix": "del ",
     "library": "var_list.py",
     "varRefreshCmd": "print(var_dic_list())"
    },
    "r": {
     "delete_cmd_postfix": ") ",
     "delete_cmd_prefix": "rm(",
     "library": "var_list.r",
     "varRefreshCmd": "cat(var_dic_list()) "
    }
   },
   "types_to_exclude": [
    "module",
    "function",
    "builtin_function_or_method",
    "instance",
    "_Feature"
   ],
   "window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}