06_zephyr_7b_alpha_feature_test.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "gpuType": "T4",
      "machine_shape": "hm",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU",
    "widgets": {
      "application/vnd.jupyter.widget-state+json": {
        "5355afef0aee436f9d1e4781d607d396": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_f465320c53ff4dbbb5cf3378fba1b1f4",
              "IPY_MODEL_7d47537b13a448d1ab7afffd18d9028e",
              "IPY_MODEL_81aec3a9171d444f82237edd2409c18d"
            ],
            "layout": "IPY_MODEL_d7ecdf203fbc489c91b72ff65a775ab5"
          }
        },
        "f465320c53ff4dbbb5cf3378fba1b1f4": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_8cbc6a687b4f4aa18f5979e721a39c9e",
            "placeholder": "​",
            "style": "IPY_MODEL_1caa626fa2e344a88de9206bfeb46abd",
            "value": "Loading checkpoint shards:   0%"
          }
        },
        "7d47537b13a448d1ab7afffd18d9028e": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_0a219e78fc694c469e147905c17394f8",
            "max": 2,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_731bec67a1b9425d9e3b19f5d336f135",
            "value": 0
          }
        },
        "81aec3a9171d444f82237edd2409c18d": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_f2900efc3efb43a8a50c8f166c92e733",
            "placeholder": "​",
            "style": "IPY_MODEL_fc4714e6c50340dfb23448db89284ad9",
            "value": " 0/2 [00:00&lt;?, ?it/s]"
          }
        },
        "d7ecdf203fbc489c91b72ff65a775ab5": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "8cbc6a687b4f4aa18f5979e721a39c9e": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "1caa626fa2e344a88de9206bfeb46abd": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "0a219e78fc694c469e147905c17394f8": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "731bec67a1b9425d9e3b19f5d336f135": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "f2900efc3efb43a8a50c8f166c92e733": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "fc4714e6c50340dfb23448db89284ad9": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        }
      }
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/gulabpatel/LLMs/blob/main/06_zephyr_7b_alpha_feature_test.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Note: Responses from local models can be quite slow, especially with 8-bit quantization.\n",
        "\n",
        "With 4bit quantization, `HuggingFaceH4/zephyr-7b-alpha` uses about 8GB of VRAM and spiked to 14GB of RAM when loading the model, then settled around 5GB. I used a T4 instance for this notebook."
      ],
      "metadata": {
        "id": "mRbrmi4xNkRT"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install llama-index transformers accelerate bitsandbytes"
      ],
      "metadata": {
        "id": "V1DwVkHCNDgT",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "ef4cb38c-449f-4a74-f5bc-234cbab1f82d"
      },
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Requirement already satisfied: llama-index in /usr/local/lib/python3.10/dist-packages (0.8.43)\n",
            "Requirement already satisfied: transformers in /usr/local/lib/python3.10/dist-packages (4.34.0)\n",
            "Requirement already satisfied: accelerate in /usr/local/lib/python3.10/dist-packages (0.23.0)\n",
            "Requirement already satisfied: bitsandbytes in /usr/local/lib/python3.10/dist-packages (0.41.1)\n",
            "Requirement already satisfied: SQLAlchemy[asyncio]>=2.0.15 in /usr/local/lib/python3.10/dist-packages (from llama-index) (2.0.21)\n",
            "Requirement already satisfied: dataclasses-json<0.6.0,>=0.5.7 in /usr/local/lib/python3.10/dist-packages (from llama-index) (0.5.14)\n",
            "Requirement already satisfied: fsspec>=2023.5.0 in /usr/local/lib/python3.10/dist-packages (from llama-index) (2023.6.0)\n",
            "Requirement already satisfied: langchain>=0.0.303 in /usr/local/lib/python3.10/dist-packages (from llama-index) (0.0.312)\n",
            "Requirement already satisfied: nest-asyncio<2.0.0,>=1.5.8 in /usr/local/lib/python3.10/dist-packages (from llama-index) (1.5.8)\n",
            "Requirement already satisfied: nltk<4.0.0,>=3.8.1 in /usr/local/lib/python3.10/dist-packages (from llama-index) (3.8.1)\n",
            "Requirement already satisfied: numpy<2.0.0,>=1.24.4 in /usr/local/lib/python3.10/dist-packages (from llama-index) (1.26.0)\n",
            "Requirement already satisfied: openai>=0.26.4 in /usr/local/lib/python3.10/dist-packages (from llama-index) (0.28.1)\n",
            "Requirement already satisfied: pandas<3.0.0,>=2.0.3 in /usr/local/lib/python3.10/dist-packages (from llama-index) (2.1.1)\n",
            "Requirement already satisfied: tenacity<9.0.0,>=8.2.0 in /usr/local/lib/python3.10/dist-packages (from llama-index) (8.2.3)\n",
            "Requirement already satisfied: tiktoken>=0.3.3 in /usr/local/lib/python3.10/dist-packages (from llama-index) (0.5.1)\n",
            "Requirement already satisfied: typing-extensions>=4.5.0 in /usr/local/lib/python3.10/dist-packages (from llama-index) (4.5.0)\n",
            "Requirement already satisfied: typing-inspect>=0.8.0 in /usr/local/lib/python3.10/dist-packages (from llama-index) (0.9.0)\n",
            "Requirement already satisfied: urllib3<2 in /usr/local/lib/python3.10/dist-packages (from llama-index) (1.26.17)\n",
            "Requirement already satisfied: filelock in /usr/local/lib/python3.10/dist-packages (from transformers) (3.12.4)\n",
            "Requirement already satisfied: huggingface-hub<1.0,>=0.16.4 in /usr/local/lib/python3.10/dist-packages (from transformers) (0.17.3)\n",
            "Requirement already satisfied: packaging>=20.0 in /usr/local/lib/python3.10/dist-packages (from transformers) (23.2)\n",
            "Requirement already satisfied: pyyaml>=5.1 in /usr/local/lib/python3.10/dist-packages (from transformers) (6.0.1)\n",
            "Requirement already satisfied: regex!=2019.12.17 in /usr/local/lib/python3.10/dist-packages (from transformers) (2023.6.3)\n",
            "Requirement already satisfied: requests in /usr/local/lib/python3.10/dist-packages (from transformers) (2.31.0)\n",
            "Requirement already satisfied: tokenizers<0.15,>=0.14 in /usr/local/lib/python3.10/dist-packages (from transformers) (0.14.1)\n",
            "Requirement already satisfied: safetensors>=0.3.1 in /usr/local/lib/python3.10/dist-packages (from transformers) (0.4.0)\n",
            "Requirement already satisfied: tqdm>=4.27 in /usr/local/lib/python3.10/dist-packages (from transformers) (4.66.1)\n",
            "Requirement already satisfied: psutil in /usr/local/lib/python3.10/dist-packages (from accelerate) (5.9.5)\n",
            "Requirement already satisfied: torch>=1.10.0 in /usr/local/lib/python3.10/dist-packages (from accelerate) (2.0.1+cu118)\n",
            "Requirement already satisfied: marshmallow<4.0.0,>=3.18.0 in /usr/local/lib/python3.10/dist-packages (from dataclasses-json<0.6.0,>=0.5.7->llama-index) (3.20.1)\n",
            "Requirement already satisfied: aiohttp<4.0.0,>=3.8.3 in /usr/local/lib/python3.10/dist-packages (from langchain>=0.0.303->llama-index) (3.8.5)\n",
            "Requirement already satisfied: anyio<4.0 in /usr/local/lib/python3.10/dist-packages (from langchain>=0.0.303->llama-index) (3.7.1)\n",
            "Requirement already satisfied: async-timeout<5.0.0,>=4.0.0 in /usr/local/lib/python3.10/dist-packages (from langchain>=0.0.303->llama-index) (4.0.3)\n",
            "Requirement already satisfied: jsonpatch<2.0,>=1.33 in /usr/local/lib/python3.10/dist-packages (from langchain>=0.0.303->llama-index) (1.33)\n",
            "Requirement already satisfied: langsmith<0.1.0,>=0.0.43 in /usr/local/lib/python3.10/dist-packages (from langchain>=0.0.303->llama-index) (0.0.43)\n",
            "Requirement already satisfied: pydantic<3,>=1 in /usr/local/lib/python3.10/dist-packages (from langchain>=0.0.303->llama-index) (1.10.13)\n",
            "Requirement already satisfied: click in /usr/local/lib/python3.10/dist-packages (from nltk<4.0.0,>=3.8.1->llama-index) (8.1.7)\n",
            "Requirement already satisfied: joblib in /usr/local/lib/python3.10/dist-packages (from nltk<4.0.0,>=3.8.1->llama-index) (1.3.2)\n",
            "Requirement already satisfied: python-dateutil>=2.8.2 in /usr/local/lib/python3.10/dist-packages (from pandas<3.0.0,>=2.0.3->llama-index) (2.8.2)\n",
            "Requirement already satisfied: pytz>=2020.1 in /usr/local/lib/python3.10/dist-packages (from pandas<3.0.0,>=2.0.3->llama-index) (2023.3.post1)\n",
            "Requirement already satisfied: tzdata>=2022.1 in /usr/local/lib/python3.10/dist-packages (from pandas<3.0.0,>=2.0.3->llama-index) (2023.3)\n",
            "Requirement already satisfied: charset-normalizer<4,>=2 in /usr/local/lib/python3.10/dist-packages (from requests->transformers) (3.3.0)\n",
            "Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.10/dist-packages (from requests->transformers) (3.4)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.10/dist-packages (from requests->transformers) (2023.7.22)\n",
            "Requirement already satisfied: greenlet!=0.4.17 in /usr/local/lib/python3.10/dist-packages (from SQLAlchemy[asyncio]>=2.0.15->llama-index) (3.0.0)\n",
            "Requirement already satisfied: sympy in /usr/local/lib/python3.10/dist-packages (from torch>=1.10.0->accelerate) (1.12)\n",
            "Requirement already satisfied: networkx in /usr/local/lib/python3.10/dist-packages (from torch>=1.10.0->accelerate) (3.1)\n",
            "Requirement already satisfied: jinja2 in /usr/local/lib/python3.10/dist-packages (from torch>=1.10.0->accelerate) (3.1.2)\n",
            "Requirement already satisfied: triton==2.0.0 in /usr/local/lib/python3.10/dist-packages (from torch>=1.10.0->accelerate) (2.0.0)\n",
            "Requirement already satisfied: cmake in /usr/local/lib/python3.10/dist-packages (from triton==2.0.0->torch>=1.10.0->accelerate) (3.27.6)\n",
            "Requirement already satisfied: lit in /usr/local/lib/python3.10/dist-packages (from triton==2.0.0->torch>=1.10.0->accelerate) (17.0.2)\n",
            "Requirement already satisfied: mypy-extensions>=0.3.0 in /usr/local/lib/python3.10/dist-packages (from typing-inspect>=0.8.0->llama-index) (1.0.0)\n",
            "Requirement already satisfied: attrs>=17.3.0 in /usr/local/lib/python3.10/dist-packages (from aiohttp<4.0.0,>=3.8.3->langchain>=0.0.303->llama-index) (23.1.0)\n",
            "Requirement already satisfied: multidict<7.0,>=4.5 in /usr/local/lib/python3.10/dist-packages (from aiohttp<4.0.0,>=3.8.3->langchain>=0.0.303->llama-index) (6.0.4)\n",
            "Requirement already satisfied: yarl<2.0,>=1.0 in /usr/local/lib/python3.10/dist-packages (from aiohttp<4.0.0,>=3.8.3->langchain>=0.0.303->llama-index) (1.9.2)\n",
            "Requirement already satisfied: frozenlist>=1.1.1 in /usr/local/lib/python3.10/dist-packages (from aiohttp<4.0.0,>=3.8.3->langchain>=0.0.303->llama-index) (1.4.0)\n",
            "Requirement already satisfied: aiosignal>=1.1.2 in /usr/local/lib/python3.10/dist-packages (from aiohttp<4.0.0,>=3.8.3->langchain>=0.0.303->llama-index) (1.3.1)\n",
            "Requirement already satisfied: sniffio>=1.1 in /usr/local/lib/python3.10/dist-packages (from anyio<4.0->langchain>=0.0.303->llama-index) (1.3.0)\n",
            "Requirement already satisfied: exceptiongroup in /usr/local/lib/python3.10/dist-packages (from anyio<4.0->langchain>=0.0.303->llama-index) (1.1.3)\n",
            "Requirement already satisfied: jsonpointer>=1.9 in /usr/local/lib/python3.10/dist-packages (from jsonpatch<2.0,>=1.33->langchain>=0.0.303->llama-index) (2.4)\n",
            "Requirement already satisfied: six>=1.5 in /usr/local/lib/python3.10/dist-packages (from python-dateutil>=2.8.2->pandas<3.0.0,>=2.0.3->llama-index) (1.16.0)\n",
            "Requirement already satisfied: MarkupSafe>=2.0 in /usr/local/lib/python3.10/dist-packages (from jinja2->torch>=1.10.0->accelerate) (2.1.3)\n",
            "Requirement already satisfied: mpmath>=0.19 in /usr/local/lib/python3.10/dist-packages (from sympy->torch>=1.10.0->accelerate) (1.3.0)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Setup"
      ],
      "metadata": {
        "id": "q1_IZvrVM_FQ"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Data"
      ],
      "metadata": {
        "id": "kkjVukvQO_5D"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from llama_index.readers import BeautifulSoupWebReader\n",
        "\n",
        "url = \"https://www.theverge.com/2023/9/29/23895675/ai-bot-social-network-openai-meta-chatbots\"\n",
        "\n",
        "documents = BeautifulSoupWebReader().load_data([url])"
      ],
      "metadata": {
        "id": "quYBhKouOIxM"
      },
      "execution_count": 2,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### LLM\n",
        "\n",
        "This should run on a T4 instance on the free tier"
      ],
      "metadata": {
        "id": "zXOJQa7FOGr5"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lMNaHDzPM68f",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 49,
          "referenced_widgets": [
            "5355afef0aee436f9d1e4781d607d396",
            "f465320c53ff4dbbb5cf3378fba1b1f4",
            "7d47537b13a448d1ab7afffd18d9028e",
            "81aec3a9171d444f82237edd2409c18d",
            "d7ecdf203fbc489c91b72ff65a775ab5",
            "8cbc6a687b4f4aa18f5979e721a39c9e",
            "1caa626fa2e344a88de9206bfeb46abd",
            "0a219e78fc694c469e147905c17394f8",
            "731bec67a1b9425d9e3b19f5d336f135",
            "f2900efc3efb43a8a50c8f166c92e733",
            "fc4714e6c50340dfb23448db89284ad9"
          ]
        },
        "outputId": "43b89e1d-976c-466d-aa57-94768abe61ed"
      },
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "Loading checkpoint shards:   0%|          | 0/2 [00:00<?, ?it/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "5355afef0aee436f9d1e4781d607d396"
            }
          },
          "metadata": {}
        }
      ],
      "source": [
        "import torch\n",
        "from transformers import BitsAndBytesConfig\n",
        "from llama_index.prompts import PromptTemplate\n",
        "from llama_index.llms import HuggingFaceLLM\n",
        "\n",
        "quantization_config = BitsAndBytesConfig(\n",
        "    load_in_4bit=True,\n",
        "    bnb_4bit_compute_dtype=torch.float16,\n",
        "    bnb_4bit_quant_type=\"nf4\",\n",
        "    bnb_4bit_use_double_quant=True,\n",
        ")\n",
        "\n",
        "\n",
        "def messages_to_prompt(messages):\n",
        "  prompt = \"\"\n",
        "  for message in messages:\n",
        "    if message.role == 'system':\n",
        "      prompt += f\"<|system|>\\n{message.content}</s>\\n\"\n",
        "    elif message.role == 'user':\n",
        "      prompt += f\"<|user|>\\n{message.content}</s>\\n\"\n",
        "    elif message.role == 'assistant':\n",
        "      prompt += f\"<|assistant|>\\n{message.content}</s>\\n\"\n",
        "\n",
        "  # ensure we start with a system prompt, insert blank if needed\n",
        "  if not prompt.startswith(\"<|system|>\\n\"):\n",
        "    prompt = \"<|system|>\\n</s>\\n\" + prompt\n",
        "\n",
        "  # add final assistant prompt\n",
        "  prompt = prompt + \"<|assistant|>\\n\"\n",
        "\n",
        "  return prompt\n",
        "\n",
        "\n",
        "llm = HuggingFaceLLM(\n",
        "    model_name=\"HuggingFaceH4/zephyr-7b-alpha\",\n",
        "    tokenizer_name=\"HuggingFaceH4/zephyr-7b-alpha\",\n",
        "    query_wrapper_prompt=PromptTemplate(\"<|system|>\\n</s>\\n<|user|>\\n{query_str}</s>\\n<|assistant|>\\n\"),\n",
        "    context_window=3900,\n",
        "    max_new_tokens=256,\n",
        "    model_kwargs={\"quantization_config\": quantization_config},\n",
        "    # tokenizer_kwargs={},\n",
        "    generate_kwargs={\"temperature\": 0.7, \"top_k\": 50, \"top_p\": 0.95},\n",
        "    messages_to_prompt=messages_to_prompt,\n",
        "    device_map=\"auto\",\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "from llama_index import ServiceContext\n",
        "\n",
        "service_context = ServiceContext.from_defaults(llm=llm, embed_model=\"local:BAAI/bge-small-en-v1.5\")"
      ],
      "metadata": {
        "id": "rPawBPL0SwFj"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Index Setup"
      ],
      "metadata": {
        "id": "_iYeofUlS6_D"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from llama_index import VectorStoreIndex\n",
        "\n",
        "vector_index = VectorStoreIndex.from_documents(documents, service_context=service_context)"
      ],
      "metadata": {
        "id": "Jz173lnJS8cZ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from llama_index import SummaryIndex\n",
        "\n",
        "summary_index = SummaryIndex.from_documents(documents, service_context=service_context)"
      ],
      "metadata": {
        "id": "CuGwYVl5S-GH"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Helpful Imports / Logging"
      ],
      "metadata": {
        "id": "pxFrAm_bTCYY"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from llama_index.response.notebook_utils import display_response"
      ],
      "metadata": {
        "id": "ghqk6C04TD3b"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import logging\n",
        "import sys\n",
        "\n",
        "logging.basicConfig(stream=sys.stdout, level=logging.INFO)\n",
        "logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))"
      ],
      "metadata": {
        "id": "6YVB1ny-UjN4"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Basic Query Engine"
      ],
      "metadata": {
        "id": "K0l3oUS8NO-I"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Compact (default)"
      ],
      "metadata": {
        "id": "TC2Vdq5LNe6H"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "query_engine = vector_index.as_query_engine(response_mode=\"compact\")\n",
        "\n",
        "response = query_engine.query(\"How do OpenAI and Meta differ on AI tools?\")\n",
        "\n",
        "display_response(response)"
      ],
      "metadata": {
        "id": "Lb9CsvvQNR9K"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Refine"
      ],
      "metadata": {
        "id": "p8WNEWO6NeEJ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "query_engine = vector_index.as_query_engine(response_mode=\"refine\")\n",
        "\n",
        "response = query_engine.query(\"How do OpenAI and Meta differ on AI tools?\")\n",
        "\n",
        "display_response(response)"
      ],
      "metadata": {
        "id": "91CZsrtjNi4I"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Tree Summarize"
      ],
      "metadata": {
        "id": "EBr8vyabNSmK"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "query_engine = vector_index.as_query_engine(response_mode=\"tree_summarize\")\n",
        "\n",
        "response = query_engine.query(\"How do OpenAI and Meta differ on AI tools?\")\n",
        "\n",
        "display_response(response)"
      ],
      "metadata": {
        "id": "X3YTKZXVNZ0R"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Router Query Engine"
      ],
      "metadata": {
        "id": "aCdR2_wmNol6"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from llama_index.tools import QueryEngineTool, ToolMetadata\n",
        "\n",
        "vector_tool = QueryEngineTool(\n",
        "    vector_index.as_query_engine(),\n",
        "    metadata=ToolMetadata(\n",
        "        name=\"vector_search\",\n",
        "        description=\"Useful for searching for specific facts.\"\n",
        "    )\n",
        ")\n",
        "\n",
        "summary_tool = QueryEngineTool(\n",
        "    summary_index.as_query_engine(response_mode=\"tree_summarize\"),\n",
        "    metadata=ToolMetadata(\n",
        "        name=\"summary\",\n",
        "        description=\"Useful for summarizing an entire document.\"\n",
        "    )\n",
        ")"
      ],
      "metadata": {
        "id": "axuO04J9SdZ_"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Single Selector"
      ],
      "metadata": {
        "id": "C7YhReZgNuoj"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from llama_index.query_engine import RouterQueryEngine\n",
        "\n",
        "query_engine = RouterQueryEngine.from_defaults(\n",
        "    [vector_tool, summary_tool],\n",
        "    service_context=service_context,\n",
        "    select_multi=False\n",
        ")\n",
        "\n",
        "response = query_engine.query(\"What was mentioned about Meta?\")\n",
        "\n",
        "display_response(response)"
      ],
      "metadata": {
        "id": "S7Dv6ENHNuI8"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Multi Selector"
      ],
      "metadata": {
        "id": "4bfvsADtNweD"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from llama_index.query_engine import RouterQueryEngine\n",
        "\n",
        "query_engine = RouterQueryEngine.from_defaults(\n",
        "    [vector_tool, summary_tool],\n",
        "    service_context=service_context,\n",
        "    select_multi=True,\n",
        ")\n",
        "\n",
        "response = query_engine.query(\"What was mentioned about Meta? Summarize with any other companies mentioned in the entire document.\")\n",
        "\n",
        "display_response(response)"
      ],
      "metadata": {
        "id": "VKqv3-txNwQS"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## SubQuestion Query Engine"
      ],
      "metadata": {
        "id": "0Xb4uNjSNyrQ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from llama_index.tools import QueryEngineTool, ToolMetadata\n",
        "\n",
        "vector_tool = QueryEngineTool(\n",
        "    vector_index.as_query_engine(),\n",
        "    metadata=ToolMetadata(\n",
        "        name=\"vector_search\",\n",
        "        description=\"Useful for searching for specific facts.\"\n",
        "    )\n",
        ")\n",
        "\n",
        "summary_tool = QueryEngineTool(\n",
        "    summary_index.as_query_engine(response_mode=\"tree_summarize\"),\n",
        "    metadata=ToolMetadata(\n",
        "        name=\"summary\",\n",
        "        description=\"Useful for summarizing an entire document.\"\n",
        "    )\n",
        ")"
      ],
      "metadata": {
        "id": "fWZeIoZgNyVc"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import nest_asyncio\n",
        "nest_asyncio.apply()"
      ],
      "metadata": {
        "id": "s4aihq2Gagsl"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from llama_index.query_engine import SubQuestionQueryEngine\n",
        "\n",
        "query_engine = SubQuestionQueryEngine.from_defaults(\n",
        "    [vector_tool, summary_tool],\n",
        "    service_context=service_context,\n",
        "    verbose=True,\n",
        ")\n",
        "\n",
        "response = query_engine.query(\"What was mentioned about Meta? How Does it differ from how OpenAI is talked about?\")\n",
        "\n",
        "display_response(response)"
      ],
      "metadata": {
        "id": "iTnorfujaT8H"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## SQL Query Engine\n",
        "\n",
        "Here, we download and use a sample SQLite database with 11 tables, with various info about music, playlists, and customers. We will limit to a select few tables for this test."
      ],
      "metadata": {
        "id": "Lkkx29-lN0jJ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import locale\n",
        "locale.getpreferredencoding = lambda: \"UTF-8\""
      ],
      "metadata": {
        "id": "KYdtjgKlO8th"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "!curl https://www.sqlitetutorial.net/wp-content/uploads/2018/03/chinook.zip -O /content/chinook.zip\n",
        "!unzip /content/chinook.zip"
      ],
      "metadata": {
        "id": "7IrjMhxlN1eI"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from sqlalchemy import create_engine, MetaData, Table, Column, String, Integer, select, column\n",
        "\n",
        "engine = create_engine(\"sqlite:////content/chinook.db\")"
      ],
      "metadata": {
        "id": "__yAhMDtbsDb"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from llama_index import SQLDatabase\n",
        "\n",
        "sql_database = SQLDatabase(engine)"
      ],
      "metadata": {
        "id": "IUV_QU0cb7S4"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from llama_index.indices.struct_store import NLSQLTableQueryEngine\n",
        "\n",
        "query_engine = NLSQLTableQueryEngine(\n",
        "    sql_database=sql_database,\n",
        "    tables=[\"albums\", \"tracks\", \"artists\"],\n",
        "    service_context=service_context\n",
        ")"
      ],
      "metadata": {
        "id": "WQZbzQuRcHd5"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "response = query_engine.query(\"What are some albums? Limit to 5.\")\n",
        "\n",
        "display_response(response)"
      ],
      "metadata": {
        "id": "W-_rUwszcQCD"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "response = query_engine.query(\"What are some artists? Limit it to 5.\")\n",
        "\n",
        "display_response(response)"
      ],
      "metadata": {
        "id": "300SGKKzcb4b"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "This last query should be a more complex join"
      ],
      "metadata": {
        "id": "smjsFMQSeQMJ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "response = query_engine.query(\"What are some tracks from the artist AC/DC? Limit it to 3\")\n",
        "\n",
        "display_response(response)"
      ],
      "metadata": {
        "id": "EeXqw3sacmie"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "print(response.metadata['sql_query'])"
      ],
      "metadata": {
        "id": "M_6qV0y7eS_t"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Programs\n",
        "\n",
        "Depending the LLM, you will have to test with either `OpenAIPydanticProgram` or `LLMTextCompletionProgram`"
      ],
      "metadata": {
        "id": "UGDmqtwBO0wM"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from typing import List\n",
        "from pydantic import BaseModel\n",
        "\n",
        "from llama_index.program import OpenAIPydanticProgram, LLMTextCompletionProgram\n",
        "\n",
        "class Song(BaseModel):\n",
        "    \"\"\"Data model for a song.\"\"\"\n",
        "\n",
        "    title: str\n",
        "    length_seconds: int\n",
        "\n",
        "\n",
        "class Album(BaseModel):\n",
        "    \"\"\"Data model for an album.\"\"\"\n",
        "\n",
        "    name: str\n",
        "    artist: str\n",
        "    songs: List[Song]"
      ],
      "metadata": {
        "id": "UJA5VfB6O0kp"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from llama_index.output_parsers import PydanticOutputParser\n",
        "\n",
        "prompt_template_str = \"\"\"\\\n",
        "Generate an example album, with an artist and a list of songs. \\\n",
        "Using the movie {movie_name} as inspiration.\\\n",
        "\"\"\"\n",
        "program = LLMTextCompletionProgram.from_defaults(\n",
        "    output_parser=PydanticOutputParser(Album),\n",
        "    prompt_template_str=prompt_template_str,\n",
        "    llm=llm,\n",
        "    verbose=True,\n",
        ")"
      ],
      "metadata": {
        "id": "sYinxYiufjLe"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "output = program(movie_name=\"The Shining\")"
      ],
      "metadata": {
        "id": "vDGnFiMcgAli"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "print(output)"
      ],
      "metadata": {
        "id": "JvoAxKsvA_b8"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Data Agent\n",
        "\n",
        "Similar to programs, OpenAI LLMs will use `OpenAIAgent`, while other LLMs will use `ReActAgent`."
      ],
      "metadata": {
        "id": "Pzh9E0jRN_mD"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from llama_index.agent import OpenAIAgent, ReActAgent\n",
        "\n",
        "agent = ReActAgent.from_tools(\n",
        "    [vector_tool, summary_tool],\n",
        "    llm=llm,\n",
        "    verbose=True\n",
        ")"
      ],
      "metadata": {
        "id": "BXdFusIKN_aD"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Some inputs are hallucinated, causing issues with responses. Likely a better system prompt or tool descriptions could help."
      ],
      "metadata": {
        "id": "Ud3SE4fVB3NZ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "response = agent.chat(\"Hello!\")\n",
        "print(response)"
      ],
      "metadata": {
        "id": "O53iJJ5Yg3Wd"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "response = agent.chat(\"What was mentioned about Meta? How Does it differ from how OpenAI is talked about?\")\n",
        "print(response)"
      ],
      "metadata": {
        "id": "0HMgXZ4Zg_ls"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "p5t7TjwhBBCt"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}