{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# DuckDB\n",
    "\n",
    ">[DuckDB](https://duckdb.org/docs/api/python/overview) is a fast in-process analytical database. DuckDB is under an MIT license.\n",
    "\n",
    "In this notebook we are going to show how to use DuckDB as a Vector store to be used in LlamaIndex.\n",
    "\n",
    "Install DuckDB with:\n",
    "\n",
    "```sh\n",
    "pip install duckdb\n",
    "```\n",
    "\n",
    "Make sure to use the latest DuckDB version (>= 0.10.0).\n",
    "\n",
    "You can run DuckDB in different modes depending on persistence:\n",
    "- `in-memory` is the default mode, where the database is created in memory, you can force this to be use by setting `database_name = \":memory:\"` when initializing the vector store.\n",
    "- `persistence` is set by using a name for a database and setting a persistence directory `database_name = \"my_vector_store.duckdb\"` where the database is persisted in the default `persist_dir` or to the one you set it to.\n",
    "\n",
    "With the vector store created, you can:\n",
    "- `.add` \n",
    "- `.get` \n",
    "- `.update`\n",
    "- `.upsert`\n",
    "- `.delete`\n",
    "- `.peek`\n",
    "- `.query` to run a search. \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Basic example\n",
    "\n",
    "In this basic example, we take the Paul Graham essay, split it into chunks, embed it using an open-source embedding model, load it into `DuckDBVectorStore`, and then query it.\n",
    "\n",
    "For the embedding model we will use OpenAI. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you're opening this Notebook on colab, you will probably need to install LlamaIndex 🦙."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install llama-index"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Creating a DuckDB Index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install duckdb\n",
    "!pip install llama-index-vector-stores-duckdb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core import VectorStoreIndex, SimpleDirectoryReader\n",
    "from llama_index.vector_stores.duckdb import DuckDBVectorStore\n",
    "from llama_index.core import StorageContext\n",
    "\n",
    "from IPython.display import Markdown, display"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Setup OpenAI API\n",
    "import os\n",
    "import openai\n",
    "\n",
    "openai.api_key = os.environ[\"OPENAI_API_KEY\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Download and prepare the sample dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--2024-02-16 19:38:34--  https://raw.githubusercontent.com/run-llama/llama_index/main/docs/docs/examples/data/paul_graham/paul_graham_essay.txt\n",
      "Resolving raw.githubusercontent.com (raw.githubusercontent.com)... 185.199.110.133, 185.199.111.133, 185.199.108.133, ...\n",
      "Connecting to raw.githubusercontent.com (raw.githubusercontent.com)|185.199.110.133|:443... connected.\n",
      "HTTP request sent, awaiting response... 200 OK\n",
      "Length: 75042 (73K) [text/plain]\n",
      "Saving to: ‘data/paul_graham/paul_graham_essay.txt’\n",
      "\n",
      "data/paul_graham/pa 100%[===================>]  73.28K  --.-KB/s    in 0.06s   \n",
      "\n",
      "2024-02-16 19:38:34 (1.24 MB/s) - ‘data/paul_graham/paul_graham_essay.txt’ saved [75042/75042]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "!mkdir -p 'data/paul_graham/'\n",
    "!wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/docs/examples/data/paul_graham/paul_graham_essay.txt' -O 'data/paul_graham/paul_graham_essay.txt'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "documents = SimpleDirectoryReader(\"data/paul_graham/\").load_data()\n",
    "\n",
    "vector_store = DuckDBVectorStore()\n",
    "storage_context = StorageContext.from_defaults(vector_store=vector_store)\n",
    "\n",
    "index = VectorStoreIndex.from_documents(\n",
    "    documents, storage_context=storage_context\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "The author mentions that before college, they worked on two main things outside of school: writing and programming. They wrote short stories and also tried writing programs on an IBM 1401 computer. They later got a microcomputer and started programming more extensively."
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "query_engine = index.as_query_engine()\n",
    "response = query_engine.query(\"What did the author do growing up?\")\n",
    "display(Markdown(f\"{response}\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Persisting to disk example\n",
    "\n",
    "Extending the previous example, if you want to save to disk, simply initialize the DuckDBVectorStore by specifying a database name and persist directory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Save to disk\n",
    "documents = SimpleDirectoryReader(\"data/paul_graham/\").load_data()\n",
    "\n",
    "vector_store = DuckDBVectorStore(\"pg.duckdb\", persist_dir=\"./persist/\")\n",
    "storage_context = StorageContext.from_defaults(vector_store=vector_store)\n",
    "\n",
    "index = VectorStoreIndex.from_documents(\n",
    "    documents, storage_context=storage_context\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load from disk\n",
    "vector_store = DuckDBVectorStore.from_local(\"./persist/pg.duckdb\")\n",
    "index = VectorStoreIndex.from_vector_store(vector_store)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "The author mentions that before college, they worked on two main things outside of school: writing and programming. They wrote short stories and also tried writing programs on an IBM 1401 computer. They later got a microcomputer and started programming more extensively."
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Query Data\n",
    "query_engine = index.as_query_engine()\n",


{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "714eb664",
   "metadata": {},
   "source": [
    ""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "307804a3-c02b-4a57-ac0d-172c30ddc851",
   "metadata": {},
   "source": [
    "# Pinecone Vector Store"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "36be66bf",
   "metadata": {},
   "source": [
    "If you're opening this Notebook on colab, you will probably need to install LlamaIndex 🦙."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9ddff1e4",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install llama-index-vector-stores-pinecone"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6807106d",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install llama-index>=0.9.31 pinecone-client>=3.0.0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d48af8e1",
   "metadata": {},
   "outputs": [],
   "source": [
    "import logging\n",
    "import sys\n",
    "import os\n",
    "\n",
    "logging.basicConfig(stream=sys.stdout, level=logging.INFO)\n",
    "logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f7010b1d-d1bb-4f08-9309-a328bb4ea396",
   "metadata": {},
   "source": [
    "#### Creating a Pinecone Index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0ce3143d-198c-4dd2-8e5a-c5cdf94f017a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pinecone import Pinecone, ServerlessSpec"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4ad14111-0bbb-4c62-906d-6d6253e0cdee",
   "metadata": {},
   "outputs": [],
   "source": [
    "os.environ[\n",
    "    \"PINECONE_API_KEY\"\n",
    "] = \"<Your Pinecone API key, from app.pinecone.io>\"\n",
    "os.environ[\"OPENAI_API_KEY\"] = \"sk-...\"\n",
    "\n",
    "api_key = os.environ[\"PINECONE_API_KEY\"]\n",
    "\n",
    "pc = Pinecone(api_key=api_key)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "233a080f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# delete if needed\n",
    "# pc.delete_index(\"quickstart\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c2c90087-bdd9-4ca4-b06b-2af883559f88",
   "metadata": {},
   "outputs": [],
   "source": [
    "# dimensions are for text-embedding-ada-002\n",
    "\n",
    "pc.create_index(\n",
    "    name=\"quickstart\",\n",
    "    dimension=1536,\n",
    "    metric=\"euclidean\",\n",
    "    spec=ServerlessSpec(cloud=\"aws\", region=\"us-west-2\"),\n",
    ")\n",
    "\n",
    "# If you need to create a PodBased Pinecone index, you could alternatively do this:\n",
    "#\n",
    "# from pinecone import Pinecone, PodSpec\n",
    "#\n",
    "# pc = Pinecone(api_key='xxx')\n",
    "#\n",
    "# pc.create_index(\n",
    "# \t name='my-index',\n",
    "# \t dimension=1536,\n",
    "# \t metric='cosine',\n",
    "# \t spec=PodSpec(\n",
    "# \t\t environment='us-east1-gcp',\n",
    "# \t\t pod_type='p1.x1',\n",
    "# \t\t pods=1\n",
    "# \t )\n",
    "# )\n",
    "#"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "667f3cb3-ce18-48d5-b9aa-bfc1a1f0f0f6",
   "metadata": {},
   "outputs": [],
   "source": [
    "pinecone_index = pc.Index(\"quickstart\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8ee4473a-094f-4d0a-a825-e1213db07240",
   "metadata": {},
   "source": [
    "#### Load documents, build the PineconeVectorStore and VectorStoreIndex"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0a2bcc07",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core import VectorStoreIndex, SimpleDirectoryReader\n",
    "from llama_index.vector_stores.pinecone import PineconeVectorStore\n",
    "from IPython.display import Markdown, display"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7d782f76",
   "metadata": {},
   "source": [
    "Download Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5104674e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Will not apply HSTS. The HSTS database must be a regular and non-world-writable file.\n",
      "ERROR: could not open HSTS store at '/home/loganm/.wget-hsts'. HSTS will be disabled.\n",
      "--2024-01-16 11:56:25--  https://raw.githubusercontent.com/run-llama/llama_index/main/docs/docs/examples/data/paul_graham/paul_graham_essay.txt\n",
      "Resolving raw.githubusercontent.com (raw.githubusercontent.com)... 185.199.108.133, 185.199.111.133, 185.199.110.133, ...\n",
      "Connecting to raw.githubusercontent.com (raw.githubusercontent.com)|185.199.108.133|:443... connected.\n",
      "HTTP request sent, awaiting response... 200 OK\n",
      "Length: 75042 (73K) [text/plain]\n",
      "Saving to: ‘data/paul_graham/paul_graham_essay.txt’\n",
      "\n",
      "data/paul_graham/pa 100%[===================>]  73.28K  --.-KB/s    in 0.04s   \n",
      "\n",
      "2024-01-16 11:56:25 (1.79 MB/s) - ‘data/paul_graham/paul_graham_essay.txt’ saved [75042/75042]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "!mkdir -p 'data/paul_graham/'\n",
    "!wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/docs/examples/data/paul_graham/paul_graham_essay.txt' -O 'data/paul_graham/paul_graham_essay.txt'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "68cbd239-880e-41a3-98d8-dbb3fab55431",
   "metadata": {},
   "outputs": [],
   "source": [
    "# load documents\n",
    "documents = SimpleDirectoryReader(\"./data/paul_graham\").load_data()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ba1558b3",
   "metadata": {},


{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "f5860e4e-3775-41b9-8329-46a6d9568056",
   "metadata": {},
   "source": [
    "## Couchbase Vector Store\n",
    "[Couchbase](https://couchbase.com/) is an award-winning distributed NoSQL cloud database that delivers unmatched versatility, performance, scalability, and financial value for all of your cloud, mobile, AI, and edge computing applications. Couchbase embraces AI with coding assistance for developers and vector search for their applications.\n",
    "\n",
    "Vector Search is a part of the [Full Text Search Service](https://docs.couchbase.com/server/current/learn/services-and-indexes/services/search-service.html) (Search Service) in Couchbase.\n",
    "\n",
    "This tutorial explains how to use Vector Search in Couchbase. You can work with both [Couchbase Capella](https://www.couchbase.com/products/capella/) and your self-managed Couchbase Server."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0fb82a99-965c-4a04-80d0-2baa91f5dcf0",
   "metadata": {},
   "source": [
    "### Installation\n",
    "If you're opening this Notebook on colab, you will probably need to install LlamaIndex 🦙."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6e0e1c57-f30c-4dd2-b1c0-91f80df7012a",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install llama-index-vector-stores-couchbase"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "629d75ac-c7c2-444f-9a3d-adbdc6533160",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install llama-index"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3321020f-849a-4848-9747-b154ecd0a15e",
   "metadata": {},
   "source": [
    "### Creating Couchbase Connection\n",
    "We create a connection to the Couchbase cluster initially and then pass the cluster object to the Vector Store.\n",
    "\n",
    "Here, we are connecting using the username and password. You can also connect using any other supported way to your cluster.\n",
    "\n",
    "For more information on connecting to the Couchbase cluster, please check the [Python SDK documentation](https://docs.couchbase.com/python-sdk/current/hello-world/start-using-sdk.html#connect).\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "97805d01-3bd5-4933-82e1-6876c5101e13",
   "metadata": {},
   "outputs": [],
   "source": [
    "COUCHBASE_CONNECTION_STRING = (\n",
    "    \"couchbase://localhost\"  # or \"couchbases://localhost\" if using TLS\n",
    ")\n",
    "DB_USERNAME = \"Administrator\"\n",
    "DB_PASSWORD = \"P@ssword1!\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "22874017-6010-4a2f-899e-08904a3108ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "from datetime import timedelta\n",
    "\n",
    "from couchbase.auth import PasswordAuthenticator\n",
    "from couchbase.cluster import Cluster\n",
    "from couchbase.options import ClusterOptions\n",
    "\n",
    "auth = PasswordAuthenticator(DB_USERNAME, DB_PASSWORD)\n",
    "options = ClusterOptions(auth)\n",
    "cluster = Cluster(COUCHBASE_CONNECTION_STRING, options)\n",
    "\n",
    "# Wait until the cluster is ready for use.\n",
    "cluster.wait_until_ready(timedelta(seconds=5))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5368c1c4-3f31-40e2-afee-5b6567445ff0",
   "metadata": {},
   "source": [
    "### Creating the Search Index\n",
    "Currently, the Search index needs to be created from the Couchbase Capella or Server UI or using the REST interface.\n",
    "\n",
    "Let us define a Search index with the name `vector-index` on the `testing` bucket\n",
    "\n",
    "For this example, let us use the Import Index feature on the Search Service on the UI.\n",
    "\n",
    "We are defining an index on the testing bucket’s `_default` scope on the `_default` collection with the vector field set to `embedding` with 1536 dimensions and the text field set to text. We are also indexing and storing all the fields under metadata in the document as a dynamic mapping to account for varying document structures. The similarity metric is set to `dot_product`.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2b944c44-a891-455a-b6d9-9c5c44fdd397",
   "metadata": {},
   "source": [
    "#### How to Import an Index to the Full Text Search service?\n",
    "\n",
    "- [Couchbase Server](https://docs.couchbase.com/server/current/search/import-search-index.html)\n",
    "    - Click on Search -> Add Index -> Import\n",
    "    - Copy the following Index definition in the Import screen\n",
    "    - Click on Create Index to create the index.\n",
    "\n",
    "\n",
    "- [Couchbase Capella](https://docs.couchbase.com/cloud/search/import-search-index.html)\n",
    "    - Copy the index definition to a new file `index.json`\n",
    "    - Import the file in Capella using the instructions in the documentation.\n",
    "    - Click on Create Index to create the index."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6984c13c-149e-461d-9859-be50bce17bab",
   "metadata": {},
   "source": [
    "#### Index Definition\n",
    "```\n",
    "{\n",
    " \"name\": \"vector-index\",\n",
    " \"type\": \"fulltext-index\",\n",
    " \"params\": {\n",
    "  \"doc_config\": {\n",
    "   \"docid_prefix_delim\": \"\",\n",
    "   \"docid_regexp\": \"\",\n",
    "   \"mode\": \"type_field\",\n",
    "   \"type_field\": \"type\"\n",
    "  },\n",
    "  \"mapping\": {\n",
    "   \"default_analyzer\": \"standard\",\n",
    "   \"default_datetime_parser\": \"dateTimeOptional\",\n",
    "   \"default_field\": \"_all\",\n",
    "   \"default_mapping\": {\n",
    "    \"dynamic\": true,\n",
    "    \"enabled\": true,\n",
    "    \"properties\": {\n",
    "     \"metadata\": {\n",
    "      \"dynamic\": true,\n",
    "      \"enabled\": true\n",
    "     },\n",
    "     \"embedding\": {\n",
    "      \"enabled\": true,\n",
    "      \"dynamic\": false,\n",
    "      \"fields\": [\n",
    "       {\n",
    "        \"dims\": 1536,\n",
    "        \"index\": true,\n",
    "        \"name\": \"embedding\",\n",
    "        \"similarity\": \"dot_product\",\n",
    "        \"type\": \"vector\",\n",
    "        \"vector_index_optimized_for\": \"recall\"\n",
    "       }\n",
    "      ]\n",
    "     },\n",
    "     \"text\": {\n",
    "      \"enabled\": true,\n",
    "      \"dynamic\": false,\n",
    "      \"fields\": [\n",
    "       {\n",
    "        \"index\": true,\n",
    "        \"name\": \"text\",\n",
    "        \"store\": true,\n",
    "        \"type\": \"text\"\n",
    "       }\n",
    "      ]\n",
    "     }\n",
    "    }\n",
    "   },\n",
    "   \"default_type\": \"_default\",\n",
    "   \"docvalues_dynamic\": false,\n",
    "   \"index_dynamic\": true,\n",
    "   \"store_dynamic\": true,\n",
    "   \"type_field\": \"_type\"\n",
    "  },\n",
    "  \"store\": {\n",


{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Hybrid Search with Qdrant BM42\n",
    "\n",
    "Qdrant recently released a new lightweight approach to sparse embeddings, [BM42](https://qdrant.tech/articles/bm42/).\n",
    "\n",
    "In this notebook, we walk through how to use BM42 with llama-index, for effecient hybrid search."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup\n",
    "\n",
    "First, we need a few packages\n",
    "- `llama-index`\n",
    "- `llama-index-vector-stores-qdrant`\n",
    "- `fastembed` or `fastembed-gpu`\n",
    "\n",
    "`llama-index` will automatically run fastembed models on GPU if the provided libraries are installed. Check out their [full installation guide](https://qdrant.github.io/fastembed/examples/FastEmbed_GPU/)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install llama-index llama-index-vector-stores-qdrant fastembed"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## (Optional) Test the fastembed package\n",
    "\n",
    "To confirm the installation worked (and also to confirm GPU usage, if used), we can run the following code.\n",
    "\n",
    "This will first download (and cache) the model locally, and then embed it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "238632fe2708433c9cbc8fffd4ccbc05",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Fetching 6 files:   0%|          | 0/6 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[613153351, 74040069] [0.3703993395381275, 0.3338314745830077]\n"
     ]
    }
   ],
   "source": [
    "from fastembed import SparseTextEmbedding\n",
    "\n",
    "model = SparseTextEmbedding(\n",
    "    model_name=\"Qdrant/bm42-all-minilm-l6-v2-attentions\",\n",
    "    # if using fastembed-gpu with cuda+onnx installed\n",
    "    # providers=[\"CudaExecutionProvider\"],\n",
    ")\n",
    "\n",
    "embeddings = model.embed([\"hello world\", \"goodbye world\"])\n",
    "\n",
    "indices, values = zip(\n",
    "    *[\n",
    "        (embedding.indices.tolist(), embedding.values.tolist())\n",
    "        for embedding in embeddings\n",
    "    ]\n",
    ")\n",
    "\n",
    "print(indices[0], values[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Construct our Hybrid Index\n",
    "\n",
    "In llama-index, we can construct a hybrid index in just a few lines of code.\n",
    "\n",
    "If you've tried hybrid in the past with splade, you will notice that this is much faster!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Loading Data\n",
    "\n",
    "Here, we use `llama-parse` to read in the Llama2 paper! Using the JSON result mode, we can get detailed data about each page, including layout and images. For now, we will use the page numbers and text.\n",
    "\n",
    "You can get a free api key for `llama-parse` by visiting [https://cloud.llamaindex.ai](https://cloud.llamaindex.ai)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir -p 'data/'\n",
    "!wget --user-agent \"Mozilla\" \"https://arxiv.org/pdf/2307.09288.pdf\" -O \"data/llama2.pdf\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import nest_asyncio\n",
    "\n",
    "nest_asyncio.apply()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Started parsing the file under job_id cac11eca-4058-4a89-a94a-5603dea3d851\n"
     ]
    }
   ],
   "source": [
    "from llama_parse import LlamaParse\n",
    "from llama_index.core import Document\n",
    "\n",
    "parser = LlamaParse(result_type=\"text\", api_key=\"llx-...\")\n",
    "\n",
    "# get per-page results, along with detailed layout info and metadata\n",
    "json_data = parser.get_json_result(\"data/llama2.pdf\")\n",
    "\n",
    "documents = []\n",
    "for document_json in json_data:\n",
    "    for page in document_json[\"pages\"]:\n",
    "        documents.append(\n",
    "            Document(text=page[\"text\"], metadata={\"page_number\": page[\"page\"]})\n",
    "        )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Construct the Index /w Qdrant\n",
    "\n",
    "With our nodes, we can construct our index with Qdrant and BM42!\n",
    "\n",
    "In this case, Qdrant is being hosted in a docker container.\n",
    "\n",
    "You can pull the latest:\n",
    "\n",
    "```\n",
    "docker pull qdrant/qdrant\n",
    "```\n",
    "\n",
    "And then to launch:\n",
    "\n",
    "```\n",
    "docker run -p 6333:6333 -p 6334:6334 \\\n",
    "    -v $(pwd)/qdrant_storage:/qdrant/storage:z \\\n",
    "    qdrant/qdrant\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Both client and aclient are provided. If using `:memory:` mode, the data between clients is not synced.\n"
     ]
    }
   ],
   "source": [
    "import qdrant_client\n",
    "from llama_index.vector_stores.qdrant import QdrantVectorStore\n",
    "\n",
    "client = qdrant_client.QdrantClient(\"http://localhost:6333\")\n",
    "aclient = qdrant_client.AsyncQdrantClient(\"http://localhost:6333\")\n",
    "\n",
    "# delete collection if it exists\n",
    "if client.collection_exists(\"llama2_bm42\"):\n",
    "    client.delete_collection(\"llama2_bm42\")\n",
    "\n",
    "vector_store = QdrantVectorStore(\n",
    "    collection_name=\"llama2_bm42\",\n",
    "    client=client,\n",
    "    aclient=aclient,\n",
    "    fastembed_sparse_model=\"Qdrant/bm42-all-minilm-l6-v2-attentions\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core import VectorStoreIndex, StorageContext\n",
    "from llama_index.embeddings.openai import OpenAIEmbedding\n",
    "\n",
    "storage_context = StorageContext.from_defaults(vector_store=vector_store)\n",
    "index = VectorStoreIndex.from_documents(\n",
    "    documents,\n",
    "    # our dense embedding model\n",


"    embed_model=OpenAIEmbedding(\n",
    "        model_name=\"text-embedding-3-small\", api_key=\"sk-proj-...\"\n",
    "    ),\n",
    "    storage_context=storage_context,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As we can see, both the dense and sparse embeddings were generated super quickly!\n",
    "\n",
    "Even though the sparse model is running locally on CPU, its very small and fast."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Test out the Index\n",
    "\n",
    "Using the powers of sparse embeddings, we can query for some very specific facts, and get the correct data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.llms.openai import OpenAI\n",
    "\n",
    "chat_engine = index.as_chat_engine(\n",
    "    chat_mode=\"condense_plus_context\",\n",
    "    llm=OpenAI(model=\"gpt-4o\", api_key=\"sk-proj-...\"),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The training hardware for Llama 2 included Meta’s Research Super Cluster (RSC) and internal production clusters. Both clusters utilized NVIDIA A100 GPUs. There were two key differences between these clusters:\n",
      "\n",
      "1. **Interconnect Type**:\n",
      "   - RSC used NVIDIA Quantum InfiniBand.\n",
      "   - The internal production cluster used a RoCE (RDMA over Converged Ethernet) solution based on commodity Ethernet switches.\n",
      "\n",
      "2. **Per-GPU Power Consumption Cap**:\n",
      "   - RSC had a power consumption cap of 400W per GPU.\n",
      "   - The internal production cluster had a power consumption cap of 350W per GPU.\n",
      "\n",
      "This setup allowed for a comparison of the suitability of these different types of interconnects for large-scale training.\n"
     ]
    }
   ],
   "source": [
    "response = chat_engine.chat(\"What training hardware was used for Llama2?\")\n",
    "print(str(response))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The main idea of Llama 2 is to provide an updated and improved version of the original Llama model, designed to be more efficient, scalable, and safe for various applications, including research and commercial use. Here are the key aspects of Llama 2:\n",
      "\n",
      "1. **Enhanced Pretraining**: Llama 2 is trained on a new mix of publicly available data, with a 40% increase in the size of the pretraining corpus compared to Llama 1. This aims to improve the model's performance and knowledge base.\n",
      "\n",
      "2. **Improved Architecture**: The model incorporates several architectural enhancements, such as increased context length and grouped-query attention (GQA), to improve inference scalability and overall performance.\n",
      "\n",
      "3. **Safety and Responsiveness**: Llama 2-Chat, a fine-tuned version of Llama 2, is optimized for dialogue use cases. It undergoes supervised fine-tuning and iterative refinement using Reinforcement Learning with Human Feedback (RLHF) to ensure safer and more helpful interactions.\n",
      "\n",
      "4. **Open Release**: Meta is releasing Llama 2 models with 7B, 13B, and 70B parameters to the general public for research and commercial use, promoting transparency and collaboration in the AI community.\n",
      "\n",
      "5. **Responsible Use**: The release includes guidelines and code examples to facilitate the safe deployment of Llama 2 and Llama 2-Chat, emphasizing the importance of safety testing and tuning tailored to specific applications.\n",
      "\n",
      "Overall, Llama 2 aims to be a more robust, scalable, and safer large language model that can be widely used and further developed by the AI community.\n"
     ]
    }
   ],
   "source": [
    "response = chat_engine.chat(\"What is the main idea of Llama2?\")\n",
    "print(str(response))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Llama 2 was evaluated and compared against several other models, both open-source and closed-source, across a variety of benchmarks. Here are the key comparisons:\n",
      "\n",
      "### Open-Source Models:\n",
      "1. **Llama 1**: Llama 2 models were compared to their predecessors, Llama 1 models. For example, Llama 2 70B showed improvements of approximately 5 points on MMLU and 8 points on BBH compared to Llama 1 65B.\n",
      "2. **MPT Models**: Llama 2 7B and 30B models outperformed MPT models of corresponding sizes in all categories except code benchmarks.\n",
      "3. **Falcon Models**: Llama 2 7B and 34B models outperformed Falcon 7B and 40B models across all benchmark categories.\n",
      "\n",
      "### Closed-Source Models:\n",
      "1. **GPT-3.5**: Llama 2 70B was compared to GPT-3.5, showing close performance on MMLU and GSM8K but a significant gap on coding benchmarks.\n",
      "2. **PaLM (540B)**: Llama 2 70B performed on par or better than PaLM (540B) on almost all benchmarks.\n",
      "3. **GPT-4 and PaLM-2-L**: There remains a large performance gap between Llama 2 70B and these more advanced models.\n",
      "\n",
      "### Benchmarks:\n",
      "Llama 2 was evaluated on a variety of benchmarks, including:\n",
      "1. **MMLU (Massive Multitask Language Understanding)**: Evaluated in a 5-shot setting.\n",
      "2. **BBH (Big Bench Hard)**: Evaluated in a 3-shot setting.\n",
      "3. **AGI Eval**: Evaluated in 3-5 shot settings, focusing on English tasks.\n",
      "4. **GSM8K**: For math problem-solving.\n",
      "5. **Human-Eval and MBPP**: For code generation.\n",
      "6. **NaturalQuestions and TriviaQA**: For world knowledge.\n",
      "7. **SQUAD and QUAC**: For reading comprehension.\n",
      "8. **BoolQ, PIQA, SIQA, Hella-Swag, ARC-e, ARC-c, NQ, TQA**: Various other benchmarks for different aspects of language understanding and reasoning.\n",
      "\n",
      "These evaluations demonstrate that Llama 2 models generally outperform their predecessors and other open-source models, while also being competitive with some of the leading closed-source models.\n"
     ]
    }
   ],
   "source": [
    "response = chat_engine.chat(\"What was Llama2 evaluated and compared against?\")\n",
    "print(str(response))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Loading from existing store\n",
    "\n",
    "With your vector index created, we can easily connect back to it!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import qdrant_client\n",
    "from llama_index.core import VectorStoreIndex\n",
    "from llama_index.embeddings.openai import OpenAIEmbedding\n",
    "from llama_index.vector_stores.qdrant import QdrantVectorStore\n",
    "\n",
    "client = qdrant_client.QdrantClient(\"http://localhost:6333\")\n",
    "aclient = qdrant_client.AsyncQdrantClient(\"http://localhost:6333\")\n",
    "\n",
    "# delete collection if it exists\n",
    "if client.collection_exists(\"llama2_bm42\"):\n",
    "    client.delete_collection(\"llama2_bm42\")\n",
    "\n",
    "vector_store = QdrantVectorStore(\n",


"chroma_client = chromadb.EphemeralClient()\n",
    "chroma_collection = chroma_client.create_collection(\"quickstart\")\n",
    "\n",
    "# define embedding function\n",
    "embed_model = HuggingFaceEmbedding(model_name=\"BAAI/bge-base-en-v1.5\")\n",
    "\n",
    "# load documents\n",
    "documents = SimpleDirectoryReader(\"./data/paul_graham/\").load_data()\n",
    "\n",
    "# set up ChromaVectorStore and load in data\n",
    "vector_store = ChromaVectorStore(chroma_collection=chroma_collection)\n",
    "storage_context = StorageContext.from_defaults(vector_store=vector_store)\n",
    "index = VectorStoreIndex.from_documents(\n",
    "    documents, storage_context=storage_context, embed_model=embed_model\n",
    ")\n",
    "\n",
    "# Query Data\n",
    "query_engine = index.as_query_engine()\n",
    "response = query_engine.query(\"What did the author do growing up?\")\n",
    "display(Markdown(f\"{response}\"))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "349de571",
   "metadata": {},
   "source": [
    "## Basic Example (including saving to disk)\n",
    "\n",
    "Extending the previous example, if you want to save to disk, simply initialize the Chroma client and pass the directory where you want the data to be saved to. \n",
    "\n",
    "`Caution`: Chroma makes a best-effort to automatically save data to disk, however multiple in-memory clients can stomp each other's work. As a best practice, only have one client per path running at any given time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9c3a56a5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "The author worked on writing and programming growing up. They wrote short stories and tried writing programs on an IBM 1401 computer. Later, they got a microcomputer and started programming games and a word processor."
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# save to disk\n",
    "\n",
    "db = chromadb.PersistentClient(path=\"./chroma_db\")\n",
    "chroma_collection = db.get_or_create_collection(\"quickstart\")\n",
    "vector_store = ChromaVectorStore(chroma_collection=chroma_collection)\n",
    "storage_context = StorageContext.from_defaults(vector_store=vector_store)\n",
    "\n",
    "index = VectorStoreIndex.from_documents(\n",
    "    documents, storage_context=storage_context, embed_model=embed_model\n",
    ")\n",
    "\n",
    "# load from disk\n",
    "db2 = chromadb.PersistentClient(path=\"./chroma_db\")\n",
    "chroma_collection = db2.get_or_create_collection(\"quickstart\")\n",
    "vector_store = ChromaVectorStore(chroma_collection=chroma_collection)\n",
    "index = VectorStoreIndex.from_vector_store(\n",
    "    vector_store,\n",
    "    embed_model=embed_model,\n",
    ")\n",
    "\n",
    "# Query Data from the persisted index\n",
    "query_engine = index.as_query_engine()\n",
    "response = query_engine.query(\"What did the author do growing up?\")\n",
    "display(Markdown(f\"{response}\"))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "d596e475",
   "metadata": {},
   "source": [
    "## Basic Example (using the Docker Container)\n",
    "\n",
    "You can also run the Chroma Server in a Docker container separately, create a Client to connect to it, and then pass that to LlamaIndex. \n",
    "\n",
    "Here is how to clone, build, and run the Docker Image:\n",
    "```\n",
    "git clone git@github.com:chroma-core/chroma.git\n",
    "docker-compose up -d --build\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d6c9bd64",
   "metadata": {},
   "outputs": [],
   "source": [
    "# create the chroma client and add our data\n",
    "import chromadb\n",
    "\n",
    "remote_db = chromadb.HttpClient()\n",
    "chroma_collection = remote_db.get_or_create_collection(\"quickstart\")\n",
    "vector_store = ChromaVectorStore(chroma_collection=chroma_collection)\n",
    "storage_context = StorageContext.from_defaults(vector_store=vector_store)\n",
    "\n",
    "index = VectorStoreIndex.from_documents(\n",
    "    documents, storage_context=storage_context, embed_model=embed_model\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "88e10c26",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "\n",
       "Growing up, the author wrote short stories, programmed on an IBM 1401, and wrote programs on a TRS-80 microcomputer. He also took painting classes at Harvard and worked as a de facto studio assistant for a painter. He also tried to start a company to put art galleries online, and wrote software to build online stores."
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Query Data from the Chroma Docker index\n",
    "query_engine = index.as_query_engine()\n",
    "response = query_engine.query(\"What did the author do growing up?\")\n",
    "display(Markdown(f\"{response}\"))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "0a0e79f7",
   "metadata": {},
   "source": [
    "## Update and Delete\n",
    "\n",
    "While building toward a real application, you want to go beyond adding data, and also update and delete data. \n",
    "\n",
    "Chroma has users provide `ids` to simplify the bookkeeping here. `ids` can be the name of the file, or a combined has like `filename_paragraphNumber`, etc.\n",
    "\n",
    "Here is a basic example showing how to do various operations:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d9411826",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [


{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "1496f9de",
   "metadata": {},
   "source": [
    ""
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "0b692c73",
   "metadata": {},
   "source": [
    "# Milvus Vector Store\n",
    "\n",
    "This guide demonstrates how to build a Retrieval-Augmented Generation (RAG) system using LlamaIndex and Milvus.\n",
    "\n",
    "The RAG system combines a retrieval system with a generative model to generate new text based on a given prompt. The system first retrieves relevant documents from a corpus using a vector similarity search engine like Milvus, and then uses a generative model to generate new text based on the retrieved documents.\n",
    "\n",
    "[Milvus](https://milvus.io/) is the world's most advanced open-source vector database, built to power embedding similarity search and AI applications.\n",
    "\n",
    "In this notebook we are going to show a quick demo of using the MilvusVectorStore."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "f81e2c81",
   "metadata": {},
   "source": [
    "## Before you begin\n",
    "\n",
    "### Install dependencies"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0d0e46d8",
   "metadata": {},
   "source": [
    "If you're opening this Notebook on colab, you will probably need to install LlamaIndex 🦙."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3e0c18ca",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install llama-index-vector-stores-milvus"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6b80700a",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install llama-index"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eab0d1a3",
   "metadata": {},
   "source": [
    "This notebook will use [Milvus Lite](https://milvus.io/docs/milvus_lite.md) requiring a higher version of pymilvus:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7661b098",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install pymilvus>=2.4.2"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "70cc8c56",
   "metadata": {},
   "source": [
    "> If you are using Google Colab, to enable dependencies just installed, you may need to **restart the runtime** (click on the \"Runtime\" menu at the top of the screen, and select \"Restart session\" from the dropdown menu)."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "f9b97a89",
   "metadata": {},
   "source": [
    "### Setup OpenAI\n",
    "\n",
    "Lets first begin by adding the openai api key. This will allow us to access chatgpt."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0c9f4d21-145a-401e-95ff-ccb259e8ef84",
   "metadata": {},
   "outputs": [],
   "source": [
    "import openai\n",
    "\n",
    "openai.api_key = \"sk-***********\""
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "a3d4e638",
   "metadata": {},
   "source": [
    "### Prepare data\n",
    "\n",
    "You can download sample data with the following commands:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2a2e24d1",
   "metadata": {},
   "outputs": [],
   "source": [
    "! mkdir -p 'data/'\n",
    "! wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/docs/examples/data/paul_graham/paul_graham_essay.txt' -O 'data/paul_graham_essay.txt'\n",
    "! wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/docs/examples/data/10k/uber_2021.pdf' -O 'data/uber_2021.pdf'"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "59ff935d",
   "metadata": {},
   "source": [
    "## Getting Started\n",
    "\n",
    "### Generate our data\n",
    "As a first example, lets generate a document from the file `paul_graham_essay.txt`. It is a single essay from Paul Graham titled `What I Worked On`. To generate the documents we will use the SimpleDirectoryReader."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "68cbd239-880e-41a3-98d8-dbb3fab55431",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Document ID: 95f25e4d-f270-4650-87ce-006d69d82033\n"
     ]
    }
   ],
   "source": [
    "from llama_index.core import SimpleDirectoryReader\n",
    "\n",
    "# load documents\n",
    "documents = SimpleDirectoryReader(\n",
    "    input_files=[\"./data/paul_graham_essay.txt\"]\n",
    ").load_data()\n",
    "\n",
    "print(\"Document ID:\", documents[0].doc_id)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "dd270925",
   "metadata": {},
   "source": [
    "### Create an index across the data\n",
    "\n",
    "Now that we have a document, we can can create an index and insert the document.\n",
    "\n",
    "> Please note that **Milvus Lite** requires `pymilvus>=2.4.2`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ba1558b3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create an index over the documents\n",
    "from llama_index.core import VectorStoreIndex, StorageContext\n",
    "from llama_index.vector_stores.milvus import MilvusVectorStore\n",
    "\n",
    "\n",
    "vector_store = MilvusVectorStore(\n",
    "    uri=\"./milvus_demo.db\", dim=1536, overwrite=True\n",
    ")\n",
    "storage_context = StorageContext.from_defaults(vector_store=vector_store)\n",
    "index = VectorStoreIndex.from_documents(\n",
    "    documents, storage_context=storage_context\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a75a5773",
   "metadata": {},
   "source": [
    "> For the parameters of `MilvusVectorStore`:\n",
    "> - Setting the `uri` as a local file, e.g.`./milvus.db`, is the most convenient method, as it automatically utilizes [Milvus Lite](https://milvus.io/docs/milvus_lite.md) to store all data in this file.\n",
    "> - If you have large scale of data, you can set up a more performant Milvus server on [docker or kubernetes](https://milvus.io/docs/quickstart.md). In this setup, please use the server uri, e.g.`http://localhost:19530`, as your `uri`.\n",


{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    ""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Azure AI Search"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Basic Example\n",
    "\n",
    "In this notebook, we take a Paul Graham essay, split it into chunks, embed it using an Azure OpenAI embedding model, load it into an Azure AI Search index, and then query it."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you're opening this Notebook on colab, you will probably need to install LlamaIndex 🦙."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install llama-index\n",
    "!pip install wget\n",
    "%pip install llama-index-vector-stores-azureaisearch\n",
    "%pip install azure-search-documents==11.5.1\n",
    "%llama-index-embeddings-azure-openai\n",
    "%llama-index-llms-azure-openai"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import logging\n",
    "import sys\n",
    "from azure.core.credentials import AzureKeyCredential\n",
    "from azure.search.documents import SearchClient\n",
    "from azure.search.documents.indexes import SearchIndexClient\n",
    "from IPython.display import Markdown, display\n",
    "from llama_index.core import (\n",
    "    SimpleDirectoryReader,\n",
    "    StorageContext,\n",
    "    VectorStoreIndex,\n",
    ")\n",
    "from llama_index.core.settings import Settings\n",
    "from llama_index.llms.azure_openai import AzureOpenAI\n",
    "from llama_index.embeddings.azure_openai import AzureOpenAIEmbedding\n",
    "from llama_index.vector_stores.azureaisearch import AzureAISearchVectorStore\n",
    "from llama_index.vector_stores.azureaisearch import (\n",
    "    IndexManagement,\n",
    "    MetadataIndexFieldType,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup Azure OpenAI"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "aoai_api_key = \"YOUR_AZURE_OPENAI_API_KEY\"\n",
    "aoai_endpoint = \"YOUR_AZURE_OPENAI_ENDPOINT\"\n",
    "aoai_api_version = \"2024-02-01\"\n",
    "\n",
    "llm = AzureOpenAI(\n",
    "    model=\"YOUR_AZURE_OPENAI_COMPLETION_MODEL_NAME\",\n",
    "    deployment_name=\"YOUR_AZURE_OPENAI_COMPLETION_DEPLOYMENT_NAME\",\n",
    "    api_key=aoai_api_key,\n",
    "    azure_endpoint=aoai_endpoint,\n",
    "    api_version=aoai_api_version,\n",
    ")\n",
    "\n",
    "# You need to deploy your own embedding model as well as your own chat completion model\n",
    "embed_model = AzureOpenAIEmbedding(\n",
    "    model=\"YOUR_AZURE_OPENAI_EMBEDDING_MODEL_NAME\",\n",
    "    deployment_name=\"YOUR_AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME\",\n",
    "    api_key=aoai_api_key,\n",
    "    azure_endpoint=aoai_endpoint,\n",
    "    api_version=aoai_api_version,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup Azure AI Search"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "search_service_api_key = \"YOUR-AZURE-SEARCH-SERVICE-ADMIN-KEY\"\n",
    "search_service_endpoint = \"YOUR-AZURE-SEARCH-SERVICE-ENDPOINT\"\n",
    "search_service_api_version = \"2024-07-01\"\n",
    "credential = AzureKeyCredential(search_service_api_key)\n",
    "\n",
    "\n",
    "# Index name to use\n",
    "index_name = \"llamaindex-vector-demo\"\n",
    "\n",
    "# Use index client to demonstrate creating an index\n",
    "index_client = SearchIndexClient(\n",
    "    endpoint=search_service_endpoint,\n",
    "    credential=credential,\n",
    ")\n",
    "\n",
    "# Use search client to demonstration using existing index\n",
    "search_client = SearchClient(\n",
    "    endpoint=search_service_endpoint,\n",
    "    index_name=index_name,\n",
    "    credential=credential,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create Index (if it does not exist)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Demonstrates creating a vector index named \"llamaindex-vector-demo\" if one doesn't exist. The index has the following fields:\n",
    "| Field Name | OData Type                |  \n",
    "|------------|---------------------------|  \n",
    "| id         | `Edm.String`              |  \n",
    "| chunk      | `Edm.String`              |  \n",
    "| embedding  | `Collection(Edm.Single)`  |  \n",
    "| metadata   | `Edm.String`              |  \n",
    "| doc_id     | `Edm.String`              |  \n",
    "| author     | `Edm.String`              |  \n",
    "| theme      | `Edm.String`              |  \n",
    "| director   | `Edm.String`              |  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "metadata_fields = {\n",
    "    \"author\": \"author\",\n",
    "    \"theme\": (\"topic\", MetadataIndexFieldType.STRING),\n",
    "    \"director\": \"director\",\n",
    "}\n",
    "\n",
    "vector_store = AzureAISearchVectorStore(\n",
    "    search_or_index_client=index_client,\n",
    "    filterable_metadata_field_keys=metadata_fields,\n",
    "    index_name=index_name,\n",
    "    index_management=IndexManagement.CREATE_IF_NOT_EXISTS,\n",
    "    id_field_key=\"id\",\n",
    "    chunk_field_key=\"chunk\",\n",
    "    embedding_field_key=\"embedding\",\n",
    "    embedding_dimensionality=1536,\n",
    "    metadata_string_field_key=\"metadata\",\n",
    "    doc_id_field_key=\"doc_id\",\n",
    "    language_analyzer=\"en.lucene\",\n",
    "    vector_algorithm_type=\"exhaustiveKnn\",\n",
    "    # compression_type=\"binary\" # Option to use \"scalar\" or \"binary\". NOTE: compression is only supported for HNSW\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir -p 'data/paul_graham/'\n",
    "!wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/docs/examples/data/paul_graham/paul_graham_essay.txt' -O 'data/paul_graham/paul_graham_essay.txt'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Loading documents\n",
    "Load the documents stored in the `data/paul_graham/` using the SimpleDirectoryReader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load documents\n",
    "documents = SimpleDirectoryReader(\"../data/paul_graham/\").load_data()\n",
    "storage_context = StorageContext.from_defaults(vector_store=vector_store)\n",
    "\n",
    "Settings.llm = llm\n",


{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0d6aba5a-abde-4e41-8a54-ef7e0de9e8a3",
   "metadata": {},
   "source": [
    "# Query Pipeline over Pandas DataFrames\n",
    "\n",
    "This is a simple example that builds a query pipeline that can perform structured operations over a Pandas DataFrame to satisfy a user query, using LLMs to infer the set of operations.\n",
    "\n",
    "This can be treated as the \"from-scratch\" version of our `PandasQueryEngine`.\n",
    "\n",
    "WARNING: This tool provides the LLM access to the `eval` function.\n",
    "Arbitrary code execution is possible on the machine running this tool.\n",
    "This tool is not recommended to be used in a production setting, and would\n",
    "require heavy sandboxing or virtual machines."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "beea049c",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install llama-index-llms-openai llama-index-experimental"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e16e7f46-6a3f-4390-b149-51d49a255d54",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.query_pipeline import (\n",
    "    QueryPipeline as QP,\n",
    "    Link,\n",
    "    InputComponent,\n",
    ")\n",
    "from llama_index.experimental.query_engine.pandas import (\n",
    "    PandasInstructionParser,\n",
    ")\n",
    "from llama_index.llms.openai import OpenAI\n",
    "from llama_index.core import PromptTemplate"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "844f1db7-8eaa-4fb2-91f4-a363db66bd5b",
   "metadata": {},
   "source": [
    "## Download Data\n",
    "\n",
    "Here we load the Titanic CSV dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "31c3d6f4-ec01-4992-8dbf-696b6feeda5f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--2024-01-13 18:39:07--  https://raw.githubusercontent.com/jerryjliu/llama_index/main/docs/docs/examples/data/csv/titanic_train.csv\n",
      "Resolving raw.githubusercontent.com (raw.githubusercontent.com)... 2606:50c0:8003::154, 2606:50c0:8001::154, 2606:50c0:8002::154, ...\n",
      "Connecting to raw.githubusercontent.com (raw.githubusercontent.com)|2606:50c0:8003::154|:443... connected.\n",
      "HTTP request sent, awaiting response... 200 OK\n",
      "Length: 57726 (56K) [text/plain]\n",
      "Saving to: ‘titanic_train.csv’\n",
      "\n",
      "titanic_train.csv   100%[===================>]  56.37K  --.-KB/s    in 0.007s  \n",
      "\n",
      "2024-01-13 18:39:07 (7.93 MB/s) - ‘titanic_train.csv’ saved [57726/57726]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "!wget 'https://raw.githubusercontent.com/jerryjliu/llama_index/main/docs/docs/examples/data/csv/titanic_train.csv' -O 'titanic_train.csv'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b2f94d45-607c-4f04-a05c-b66ecedfe233",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "\n",
    "df = pd.read_csv(\"./titanic_train.csv\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b42ccfe1-fd13-4356-b5b1-c9eabac5f391",
   "metadata": {},
   "source": [
    "## Define Modules\n",
    "\n",
    "Here we define the set of modules: \n",
    "1. Pandas prompt to infer pandas instructions from user query\n",
    "2. Pandas output parser to execute pandas instructions on dataframe, get back dataframe\n",
    "3. Response synthesis prompt to synthesize a final response given the dataframe\n",
    "4. LLM\n",
    "\n",
    "The pandas output parser specifically is designed to safely execute Python code. It includes a lot of safety checks that may be annoying to write from scratch. This includes only importing from a set of approved modules (e.g. no modules that would alter the file system like `os`), and also making sure that no private/dunder methods are being called."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d47ac723-e1fc-47c2-9736-19deb50d87ae",
   "metadata": {},
   "outputs": [],
   "source": [
    "instruction_str = (\n",
    "    \"1. Convert the query to executable Python code using Pandas.\\n\"\n",
    "    \"2. The final line of code should be a Python expression that can be called with the `eval()` function.\\n\"\n",
    "    \"3. The code should represent a solution to the query.\\n\"\n",
    "    \"4. PRINT ONLY THE EXPRESSION.\\n\"\n",
    "    \"5. Do not quote the expression.\\n\"\n",
    ")\n",
    "\n",
    "pandas_prompt_str = (\n",
    "    \"You are working with a pandas dataframe in Python.\\n\"\n",
    "    \"The name of the dataframe is `df`.\\n\"\n",
    "    \"This is the result of `print(df.head())`:\\n\"\n",
    "    \"{df_str}\\n\\n\"\n",
    "    \"Follow these instructions:\\n\"\n",
    "    \"{instruction_str}\\n\"\n",
    "    \"Query: {query_str}\\n\\n\"\n",
    "    \"Expression:\"\n",
    ")\n",
    "response_synthesis_prompt_str = (\n",
    "    \"Given an input question, synthesize a response from the query results.\\n\"\n",
    "    \"Query: {query_str}\\n\\n\"\n",
    "    \"Pandas Instructions (optional):\\n{pandas_instructions}\\n\\n\"\n",
    "    \"Pandas Output: {pandas_output}\\n\\n\"\n",
    "    \"Response: \"\n",
    ")\n",
    "\n",
    "pandas_prompt = PromptTemplate(pandas_prompt_str).partial_format(\n",
    "    instruction_str=instruction_str, df_str=df.head(5)\n",
    ")\n",
    "pandas_output_parser = PandasInstructionParser(df)\n",
    "response_synthesis_prompt = PromptTemplate(response_synthesis_prompt_str)\n",
    "llm = OpenAI(model=\"gpt-3.5-turbo\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1b7907f6-2926-4c83-9541-d089d0b0af88",
   "metadata": {},
   "source": [
    "## Build Query Pipeline\n",
    "\n",
    "Looks like this:\n",
    "input query_str -> pandas_prompt -> llm1 -> pandas_output_parser -> response_synthesis_prompt -> llm2\n",
    "\n",
    "Additional connections to response_synthesis_prompt: llm1 -> pandas_instructions, and pandas_output_parser -> pandas_output."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3544a6a8-3866-4806-8108-d908bd5d17e6",
   "metadata": {},
   "outputs": [],
   "source": [
    "qp = QP(\n",
    "    modules={\n",
    "        \"input\": InputComponent(),\n",
    "        \"pandas_prompt\": pandas_prompt,\n",
    "        \"llm1\": llm,\n",
    "        \"pandas_output_parser\": pandas_output_parser,\n",
    "        \"response_synthesis_prompt\": response_synthesis_prompt,\n",


{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Query Pipeline Chat Engine\n",
    "\n",
    "By combining a query pipeline with a memory buffer, we can design our own custom chat engine loop."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install llama-index-core\n",
    "%pip install llama-index-llms-openai\n",
    "%pip install llama-index-embeddings-openai\n",
    "%pip install llama-index-postprocessor-colbert-rerank\n",
    "%pip install llama-index-readers-web"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ[\"OPENAI_API_KEY\"] = \"sk-...\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Index Construction\n",
    "\n",
    "As a test, we will index Anthropic's latest documentation about tool/function calling."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.readers.web import BeautifulSoupWebReader\n",
    "\n",
    "reader = BeautifulSoupWebReader()\n",
    "\n",
    "documents = reader.load_data(\n",
    "    [\"https://docs.anthropic.com/claude/docs/tool-use\"]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you inspected the document text, you'd notice that there are way too many blank lines, lets clean that up a bit."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lines = documents[0].text.split(\"\\n\")\n",
    "\n",
    "# remove sections with more than two empty lines in a row\n",
    "fixed_lines = [lines[0]]\n",
    "for idx in range(1, len(lines)):\n",
    "    if lines[idx].strip() == \"\" and lines[idx - 1].strip() == \"\":\n",
    "        continue\n",
    "    fixed_lines.append(lines[idx])\n",
    "\n",
    "documents[0].text = \"\\n\".join(fixed_lines)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, we can create our index using OpenAI embeddings."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core import VectorStoreIndex\n",
    "from llama_index.embeddings.openai import OpenAIEmbedding\n",
    "\n",
    "index = VectorStoreIndex.from_documents(\n",
    "    documents,\n",
    "    embed_model=OpenAIEmbedding(\n",
    "        model=\"text-embedding-3-large\", embed_batch_size=256\n",
    "    ),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Query Pipeline Contruction\n",
    "\n",
    "As a demonstration, lets make a robust query pipeline with HyDE for retrieval and Colbert for reranking."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.query_pipeline import (\n",
    "    QueryPipeline,\n",
    "    InputComponent,\n",
    "    ArgPackComponent,\n",
    ")\n",
    "from llama_index.core.prompts import PromptTemplate\n",
    "from llama_index.llms.openai import OpenAI\n",
    "from llama_index.postprocessor.colbert_rerank import ColbertRerank\n",
    "\n",
    "# First, we create an input component to capture the user query\n",
    "input_component = InputComponent()\n",
    "\n",
    "# Next, we use the LLM to rewrite a user query\n",
    "rewrite = (\n",
    "    \"Please write a query to a semantic search engine using the current conversation.\\n\"\n",
    "    \"\\n\"\n",
    "    \"\\n\"\n",
    "    \"{chat_history_str}\"\n",
    "    \"\\n\"\n",
    "    \"\\n\"\n",
    "    \"Latest message: {query_str}\\n\"\n",
    "    'Query:\"\"\"\\n'\n",
    ")\n",
    "rewrite_template = PromptTemplate(rewrite)\n",
    "llm = OpenAI(\n",
    "    model=\"gpt-4-turbo-preview\",\n",
    "    temperature=0.2,\n",
    ")\n",
    "\n",
    "# we will retrieve two times, so we need to pack the retrieved nodes into a single list\n",
    "argpack_component = ArgPackComponent()\n",
    "\n",
    "# using that, we will retrieve...\n",
    "retriever = index.as_retriever(similarity_top_k=6)\n",
    "\n",
    "# then postprocess/rerank with Colbert\n",
    "reranker = ColbertRerank(top_n=3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For generating a response using chat history + retrieved nodes, lets create a custom component."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# then lastly, we need to create a response using the nodes AND chat history\n",
    "from typing import Any, Dict, List, Optional\n",
    "from llama_index.core.bridge.pydantic import Field\n",
    "from llama_index.core.llms import ChatMessage\n",
    "from llama_index.core.query_pipeline import CustomQueryComponent\n",
    "from llama_index.core.schema import NodeWithScore\n",
    "\n",
    "DEFAULT_CONTEXT_PROMPT = (\n",
    "    \"Here is some context that may be relevant:\\n\"\n",
    "    \"-----\\n\"\n",
    "    \"{node_context}\\n\"\n",
    "    \"-----\\n\"\n",
    "    \"Please write a response to the following question, using the above context:\\n\"\n",
    "    \"{query_str}\\n\"\n",
    ")\n",
    "\n",
    "\n",
    "class ResponseWithChatHistory(CustomQueryComponent):\n",
    "    llm: OpenAI = Field(..., description=\"OpenAI LLM\")\n",
    "    system_prompt: Optional[str] = Field(\n",
    "        default=None, description=\"System prompt to use for the LLM\"\n",
    "    )\n",
    "    context_prompt: str = Field(\n",
    "        default=DEFAULT_CONTEXT_PROMPT,\n",
    "        description=\"Context prompt to use for the LLM\",\n",
    "    )\n",
    "\n",
    "    def _validate_component_inputs(\n",
    "        self, input: Dict[str, Any]\n",
    "    ) -> Dict[str, Any]:\n",
    "        \"\"\"Validate component inputs during run_component.\"\"\"\n",
    "        # NOTE: this is OPTIONAL but we show you where to do validation as an example\n",
    "        return input\n",
    "\n",
    "    @property\n",
    "    def _input_keys(self) -> set:\n",
    "        \"\"\"Input keys dict.\"\"\"\n",
    "        # NOTE: These are required inputs. If you have optional inputs please override\n",
    "        # `optional_input_keys_dict`\n",
    "        return {\"chat_history\", \"nodes\", \"query_str\"}\n",
    "\n",
    "    @property\n",
    "    def _output_keys(self) -> set:\n",
    "        return {\"response\"}\n",
    "\n",
    "    def _prepare_context(\n",
    "        self,\n",
    "        chat_history: List[ChatMessage],\n",
    "        nodes: List[NodeWithScore],\n",
    "        query_str: str,\n",
    "    ) -> List[ChatMessage]:\n",
    "        node_context = \"\"\n",
    "        for idx, node in enumerate(nodes):\n",
    "            node_text = node.get_content(metadata_mode=\"llm\")\n",
    "            node_context += f\"Context Chunk {idx}:\\n{node_text}\\n\\n\"\n",
    "\n",
    "        formatted_context = self.context_prompt.format(\n",


{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "cd032bcb-fefb-48ec-94da-08d49ac26120",
   "metadata": {},
   "source": [
    "# Query Pipeline with Async/Parallel Execution\n",
    "\n",
    "Here we showcase our query pipeline with async + parallel execution.\n",
    "\n",
    "We do this by setting up a RAG pipeline that does the following:\n",
    "1. Send query to multiple RAG query engines.\n",
    "2. Combine results.\n",
    "\n",
    "In the process we'll also show some nice abstractions for joining results (e.g. our `ArgPackComponent()`)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3531eedc-4f65-457e-8844-55fcc1773154",
   "metadata": {},
   "source": [
    "## Load Data\n",
    "\n",
    "Load in the Paul Graham essay as an example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "462a4f6f",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install llama-index-llms-openai"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2a441905-9007-44d6-b71a-6fc3e5023e49",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--2024-01-10 12:31:00--  https://raw.githubusercontent.com/run-llama/llama_index/main/docs/docs/examples/data/paul_graham/paul_graham_essay.txt\n",
      "Resolving raw.githubusercontent.com (raw.githubusercontent.com)... 185.199.111.133, 185.199.110.133, 185.199.108.133, ...\n",
      "Connecting to raw.githubusercontent.com (raw.githubusercontent.com)|185.199.111.133|:443... connected.\n",
      "HTTP request sent, awaiting response... 200 OK\n",
      "Length: 75042 (73K) [text/plain]\n",
      "Saving to: ‘pg_essay.txt’\n",
      "\n",
      "pg_essay.txt        100%[===================>]  73.28K  --.-KB/s    in 0.01s   \n",
      "\n",
      "2024-01-10 12:31:00 (6.32 MB/s) - ‘pg_essay.txt’ saved [75042/75042]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "!wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/docs/examples/data/paul_graham/paul_graham_essay.txt' -O 'data/paul_graham/paul_graham_essay.txt' -O pg_essay.txt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3533149c-4312-4444-9b45-52afe21731ed",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core import SimpleDirectoryReader\n",
    "\n",
    "reader = SimpleDirectoryReader(input_files=[\"pg_essay.txt\"])\n",
    "documents = reader.load_data()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6c1d5ff8-ae04-4ea3-bbe0-2c097af71efd",
   "metadata": {},
   "source": [
    "## Setup Query Pipeline\n",
    "\n",
    "We setup a parallel query pipeline that executes multiple chunk sizes at once, and combines the results."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "63caf998-0a88-4c50-b6a4-2a0c412bde5b",
   "metadata": {},
   "source": [
    "### Define Modules\n",
    "\n",
    "This includes:\n",
    "- LLM\n",
    "- Chunk Sizes\n",
    "- Query Engines"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "01fcbdb2-6747-4e65-b1ce-5d40febccb81",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.query_pipeline import (\n",
    "    QueryPipeline,\n",
    "    InputComponent,\n",
    "    ArgPackComponent,\n",
    ")\n",
    "from typing import Dict, Any, List, Optional\n",
    "from llama_index.core.llama_pack import BaseLlamaPack\n",
    "from llama_index.core.llms import LLM\n",
    "from llama_index.llms.openai import OpenAI\n",
    "from llama_index.core import Document, VectorStoreIndex\n",
    "from llama_index.core.response_synthesizers import TreeSummarize\n",
    "from llama_index.core.schema import NodeWithScore, TextNode\n",
    "from llama_index.core.node_parser import SentenceSplitter\n",
    "\n",
    "\n",
    "llm = OpenAI(model=\"gpt-3.5-turbo\")\n",
    "chunk_sizes = [128, 256, 512, 1024]\n",
    "query_engines = {}\n",
    "for chunk_size in chunk_sizes:\n",
    "    splitter = SentenceSplitter(chunk_size=chunk_size, chunk_overlap=0)\n",
    "    nodes = splitter.get_nodes_from_documents(documents)\n",
    "    vector_index = VectorStoreIndex(nodes)\n",
    "    query_engines[str(chunk_size)] = vector_index.as_query_engine(llm=llm)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7a87a439-88e6-4130-b28f-45268330d3e4",
   "metadata": {},
   "source": [
    "### Construct Query Pipeline\n",
    "\n",
    "Connect input to multiple query engines, and join the results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ff95be2e-517f-4632-a7b8-a2e0dec11d73",
   "metadata": {},
   "outputs": [],
   "source": [
    "# construct query pipeline\n",
    "p = QueryPipeline(verbose=True)\n",
    "module_dict = {\n",
    "    **query_engines,\n",
    "    \"input\": InputComponent(),\n",
    "    \"summarizer\": TreeSummarize(),\n",
    "    \"join\": ArgPackComponent(\n",
    "        convert_fn=lambda x: NodeWithScore(node=TextNode(text=str(x)))\n",
    "    ),\n",
    "}\n",
    "p.add_modules(module_dict)\n",
    "# add links from input to query engine (id'ed by chunk_size)\n",
    "for chunk_size in chunk_sizes:\n",
    "    p.add_link(\"input\", str(chunk_size))\n",
    "    p.add_link(str(chunk_size), \"join\", dest_key=str(chunk_size))\n",
    "p.add_link(\"join\", \"summarizer\", dest_key=\"nodes\")\n",
    "p.add_link(\"input\", \"summarizer\", dest_key=\"query_str\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bda05274-09c5-4b56-b2ba-57f445346e73",
   "metadata": {},
   "source": [
    "## Try out Queries\n",
    "\n",
    "Let's compare the async performance vs. synchronous performance!\n",
    "\n",
    "In our experiments we get a 2x speedup."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d3e161ce-ef10-446f-acfb-f6d3a1d291bb",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1;3;38;2;155;135;227m> Running modules and inputs in parallel: \n",
      "Module key: input. Input: \n",
      "input: What did the author do during his time in YC?\n",
      "\n",
      "\n",
      "\u001b[0m\u001b[1;3;38;2;155;135;227m> Running modules and inputs in parallel: \n",
      "Module key: 128. Input: \n",


},
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "de0e3aa5",
   "metadata": {},
   "outputs": [],
   "source": [
    "db = chromadb.PersistentClient(path=\"./chroma_db\")\n",
    "chroma_collection = db.get_or_create_collection(\"dense_vectors\")\n",
    "vector_store = ChromaVectorStore(chroma_collection=chroma_collection)\n",
    "\n",
    "docstore = SimpleDocumentStore.from_persist_path(\"./docstore.json\")\n",
    "\n",
    "storage_context = StorageContext.from_defaults(\n",
    "    docstore=docstore, vector_store=vector_store\n",
    ")\n",
    "\n",
    "index = VectorStoreIndex(nodes=[], storage_context=storage_context)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "llama-index-caVs7DDe-py3.11",
   "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"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}


"Doc 8 (node score, doc similarity, full similarity): (0.5672766061523489, 0.24827341793941335, 0.4077750120458811)\n",
      "Doc 9 (node score, doc similarity, full similarity): (0.5671131641337652, 0.24827341793941335, 0.4076932910365893)\n",
      "The LLM interface is a unified interface provided by LlamaIndex for defining Large Language Models (LLMs) from different sources such as OpenAI, Hugging Face, or LangChain. This interface eliminates the need to write the boilerplate code for defining the LLM interface yourself. The LLM interface supports text completion and chat endpoints, as well as streaming and non-streaming endpoints. It also supports both synchronous and asynchronous endpoints.\n",
      "\n",
      "LLMs are a core component of LlamaIndex and can be used as standalone modules or plugged into other core LlamaIndex modules such as indices, retrievers, and query engines. They are primarily used during the response synthesis step, which occurs after retrieval. Depending on the type of index being used, LLMs may also be used during index construction, insertion, and query traversal.\n",
      "\n",
      "To use LLMs, you can import the necessary modules and instantiate the LLM object. You can then use the LLM object to generate responses or complete text prompts. LlamaIndex provides examples and code snippets to help you get started with using LLMs.\n",
      "\n",
      "It's important to note that tokenization plays a crucial role in LLMs. LlamaIndex uses a global tokenizer by default, but if you change the LLM, you may need to update the tokenizer to ensure accurate token counts, chunking, and prompting. LlamaIndex provides instructions on how to set a global tokenizer using libraries like tiktoken or Hugging Face's AutoTokenizer.\n",
      "\n",
      "Overall, LLMs are powerful tools for building LlamaIndex applications and can be customized within the LlamaIndex abstractions. While LLMs from paid APIs like OpenAI and Anthropic are generally considered more reliable, local open-source models are gaining popularity due to their customizability and transparency. LlamaIndex offers integrations with various LLMs and provides documentation on their compatibility and performance. Contributions to improve the setup and performance of existing LLMs or to add new LLMs are welcome.\n"
     ]
    }
   ],
   "source": [
    "response = query_engine.query(query_str)\n",
    "print(str(response))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The LLM interface is a unified interface provided by LlamaIndex for defining Large Language Model (LLM) modules. It allows users to easily integrate LLMs from different providers such as OpenAI, Hugging Face, or LangChain into their applications without having to write the boilerplate code for defining the LLM interface themselves.\n",
      "\n",
      "LLMs are a core component of LlamaIndex and can be used as standalone modules or plugged into other core LlamaIndex modules such as indices, retrievers, and query engines. They are primarily used during the response synthesis step, which occurs after retrieval. Depending on the type of index being used, LLMs may also be used during index construction, insertion, and query traversal.\n",
      "\n",
      "The LLM interface supports various functionalities, including text completion and chat endpoints. It also provides support for streaming and non-streaming endpoints, as well as synchronous and asynchronous endpoints.\n",
      "\n",
      "To use LLMs, you can import the necessary modules and make use of the provided functions. For example, you can use the OpenAI module to interact with the gpt-3.5-turbo LLM by calling the `OpenAI()` function. You can then use the `complete()` function to generate completions based on a given prompt.\n",
      "\n",
      "It's important to note that LlamaIndex uses a global tokenizer called cl100k from tiktoken by default for all token counting. If you change the LLM being used, you may need to update the tokenizer to ensure accurate token counts, chunking, and prompting.\n",
      "\n",
      "Overall, LLMs and the LLM interface provided by LlamaIndex are essential for building LLM applications and integrating them into the LlamaIndex ecosystem.\n"
     ]
    }
   ],
   "source": [
    "base_response = base_query_engine.query(query_str)\n",
    "print(str(base_response))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "llama_index_v2",
   "language": "python",
   "name": "llama_index_v2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}


{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "6340e329",
   "metadata": {},
   "source": [
    ""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8d09c269-526b-4858-979d-d77285c25260",
   "metadata": {},
   "source": [
    "# Azure OpenAI"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d981ca8e",
   "metadata": {},
   "source": [
    "If you're opening this Notebook on colab, you will probably need to install LlamaIndex 🦙."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24e9114d",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install llama-index-llms-azure-openai"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "02fcb2e6",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install llama-index"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5386685e-fa9e-46b4-8bd9-2cdec2d9903e",
   "metadata": {},
   "source": [
    "## Prerequisites"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5a3491e8-11e8-4548-b6b9-d08246c7ef9b",
   "metadata": {},
   "source": [
    "1. Setup an Azure subscription - you can create one for free [here](https://azure.microsoft.com/en-us/free/cognitive-services/)\n",
    "2. Apply for access to Azure OpenAI Service [here](https://customervoice.microsoft.com/Pages/ResponsePage.aspx?id=v4j5cvGGr0GRqy180BHbR7en2Ais5pxKtso_Pz4b1_xUOFA5Qk1UWDRBMjg0WFhPMkIzTzhKQ1dWNyQlQCN0PWcu) \n",
    "3. Create a resource in the Azure portal [here](https://portal.azure.com/?microsoft_azure_marketplace_ItemHideKey=microsoft_openai_tip#create/Microsoft.CognitiveServicesOpenAI)\n",
    "4. Deploy a model in Azure OpenAI Studio [here](https://oai.azure.com/)\n",
    "\n",
    "\n",
    "You can find more details in [this guide.](https://learn.microsoft.com/en-us/azure/cognitive-services/openai/how-to/create-resource?pivots=web-portal)\n",
    "\n",
    "Note down the **\"model name\"** and **\"deployment name\"**, you'll need it when connecting to your LLM."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14050fdc-890d-464c-89ff-9f444672de1d",
   "metadata": {},
   "source": [
    "## Environment Setup"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3635cdc5-d27c-48c1-9041-ccb91d239956",
   "metadata": {},
   "source": [
    "### Find your setup information - API base, API key, deployment name (i.e. engine), etc"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f7b24c58-b8fe-4687-bfa7-a659c2a86d93",
   "metadata": {},
   "source": [
    "To find the setup information necessary, do the following setups:\n",
    "1. Go to the Azure OpenAI Studio [here](https://oai.azure.com/)\n",
    "2. Go to the chat or completions playground (depending on which LLM you are setting up)\n",
    "3. Click \"view code\" (shown in image below)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aa6cd770-8b22-425d-b5cc-af2613a71c10",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from IPython.display import Image\n",
    "\n",
    "Image(filename=\"./azure_playground.png\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "271c6132-dc8f-4c5e-a440-471c81bb0302",
   "metadata": {},
   "source": [
    "4. Note down the `api_type`, `api_base`, `api_version`, `engine` (this should be the same as the \"deployment name\" from before), and the `key`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "434775cb-2ca0-408c-b6f9-a08b3f0c63b9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from IPython.display import Image\n",
    "\n",
    "Image(filename=\"./azure_env.png\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7a5bf13f-4d64-446c-bb0e-fdf3e03d22d4",
   "metadata": {},
   "source": [
    "### Configure environment variables"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ae5dd044-6c03-4854-b643-84e79a1c05f2",
   "metadata": {},
   "source": [
    "Using Azure deployment of OpenAI models is very similar to normal OpenAI. \n",
    "You just need to configure a couple more environment variables.\n",
    "\n",
    "- `OPENAI_API_VERSION`: set this to `2023-07-01-preview`\n",
    "    This may change in the future.\n",
    "- `AZURE_OPENAI_ENDPOINT`: your endpoint should look like the following\n",
    "    https://YOUR_RESOURCE_NAME.openai.azure.com/\n",
    "- `OPENAI_API_KEY`: your API key"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "82233222-5a91-473b-b968-10bf8b7105e7",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ[\"OPENAI_API_KEY\"] = \"<your-api-key>\"\n",
    "os.environ[\n",
    "    \"AZURE_OPENAI_ENDPOINT\"\n",
    "] = \"https://<your-resource-name>.openai.azure.com/\"\n",
    "os.environ[\"OPENAI_API_VERSION\"] = \"2023-07-01-preview\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a593031b-c872-4360-8775-dff4844ccead",
   "metadata": {},
   "source": [
    "## Use your LLM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fd389e2c-a3d5-4b47-acbe-b22b3da17670",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.llms.azure_openai import AzureOpenAI"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ae049806-09b3-46fe-b589-2ae2f33beda9",
   "metadata": {},
   "source": [
    "Unlike normal `OpenAI`, you need to pass a `engine` argument in addition to `model`. The `engine` is the name of your model deployment you selected in Azure OpenAI Studio. See previous section on \"find your setup information\" for more details."
   ]
  },
  {


{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "3ac9adb4",
   "metadata": {},
   "source": [
    ""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "368686b4-f487-4dd4-aeff-37823976529d",
   "metadata": {},
   "source": [
    "# LlamaCPP \n",
    "\n",
    "In this short notebook, we show how to use the [llama-cpp-python](https://github.com/abetlen/llama-cpp-python) library with LlamaIndex.\n",
    "\n",
    "In this notebook, we use the [`llama-2-chat-13b-ggml`](https://huggingface.co/TheBloke/Llama-2-13B-chat-GGML) model, along with the proper prompt formatting. \n",
    "\n",
    "Note that if you're using a version of `llama-cpp-python` after version `0.1.79`, the model format has changed from `ggmlv3` to `gguf`. Old model files like the used in this notebook can be converted using scripts in the [`llama.cpp`](https://github.com/ggerganov/llama.cpp) repo. Alternatively, you can download the GGUF version of the model above from [huggingface](https://huggingface.co/TheBloke/Llama-2-13B-chat-GGUF).\n",
    "\n",
    "By default, if model_path and model_url are blank, the `LlamaCPP` module will load llama2-chat-13B in either format depending on your version.\n",
    "\n",
    "## Installation\n",
    "\n",
    "To get the best performance out of `LlamaCPP`, it is recomended to install the package so that it is compilied with GPU support. A full guide for installing this way is [here](https://github.com/abetlen/llama-cpp-python#installation-with-openblas--cublas--clblast--metal).\n",
    "\n",
    "Full MACOS instructions are also [here](https://llama-cpp-python.readthedocs.io/en/latest/install/macos/).\n",
    "\n",
    "In general:\n",
    "- Use `CuBLAS` if you have CUDA and an NVidia GPU\n",
    "- Use `METAL` if you are running on an M1/M2 MacBook\n",
    "- Use `CLBLAST` if you are running on an AMD/Intel GPU"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aff273be",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install llama-index-embeddings-huggingface\n",
    "%pip install llama-index-llms-llama-cpp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "40a33749",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core import SimpleDirectoryReader, VectorStoreIndex\n",
    "from llama_index.llms.llama_cpp import LlamaCPP\n",
    "from llama_index.llms.llama_cpp.llama_utils import (\n",
    "    messages_to_prompt,\n",
    "    completion_to_prompt,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e7927630-0044-41fb-a8a6-8dc3d2adb608",
   "metadata": {},
   "source": [
    "## Setup LLM\n",
    "\n",
    "The LlamaCPP llm is highly configurable. Depending on the model being used, you'll want to pass in `messages_to_prompt` and `completion_to_prompt` functions to help format the model inputs.\n",
    "\n",
    "Since the default model is llama2-chat, we use the util functions found in [`llama_index.llms.llama_utils`](https://github.com/jerryjliu/llama_index/blob/main/llama_index/llms/llama_utils.py).\n",
    "\n",
    "For any kwargs that need to be passed in during initialization, set them in `model_kwargs`. A full list of available model kwargs is available in the [LlamaCPP docs](https://llama-cpp-python.readthedocs.io/en/latest/api-reference/#llama_cpp.llama.Llama.__init__).\n",
    "\n",
    "For any kwargs that need to be passed in during inference, you can set them in `generate_kwargs`. See the full list of [generate kwargs here](https://llama-cpp-python.readthedocs.io/en/latest/api-reference/#llama_cpp.llama.Llama.__call__).\n",
    "\n",
    "In general, the defaults are a great starting point. The example below shows configuration with all defaults.\n",
    "\n",
    "As noted above, we're using the [`llama-2-chat-13b-ggml`](https://huggingface.co/TheBloke/Llama-2-13B-chat-GGML) model in this notebook which uses the `ggmlv3` model format. If you are running a version of `llama-cpp-python` greater than `0.1.79`, you can replace the `model_url` below with `\"https://huggingface.co/TheBloke/Llama-2-13B-chat-GGUF/resolve/main/llama-2-13b-chat.Q4_0.gguf\"`."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "59b27895",
   "metadata": {},
   "source": [
    "If you're opening this Notebook on colab, you will probably need to install LlamaIndex 🦙."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "439960c5",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install llama-index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2640c7a4",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_url = \"https://huggingface.co/TheBloke/Llama-2-13B-chat-GGML/resolve/main/llama-2-13b-chat.ggmlv3.q4_0.bin\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6fa0ec4f-03ff-4e28-957f-b4b99a0faa20",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "llama.cpp: loading model from /Users/rchan/Library/Caches/llama_index/models/llama-2-13b-chat.ggmlv3.q4_0.bin\n",
      "llama_model_load_internal: format     = ggjt v3 (latest)\n",
      "llama_model_load_internal: n_vocab    = 32000\n",
      "llama_model_load_internal: n_ctx      = 3900\n",
      "llama_model_load_internal: n_embd     = 5120\n",
      "llama_model_load_internal: n_mult     = 256\n",
      "llama_model_load_internal: n_head     = 40\n",
      "llama_model_load_internal: n_head_kv  = 40\n",
      "llama_model_load_internal: n_layer    = 40\n",
      "llama_model_load_internal: n_rot      = 128\n",
      "llama_model_load_internal: n_gqa      = 1\n",
      "llama_model_load_internal: rnorm_eps  = 5.0e-06\n",
      "llama_model_load_internal: n_ff       = 13824\n",
      "llama_model_load_internal: freq_base  = 10000.0\n",
      "llama_model_load_internal: freq_scale = 1\n",
      "llama_model_load_internal: ftype      = 2 (mostly Q4_0)\n",
      "llama_model_load_internal: model size = 13B\n",
      "llama_model_load_internal: ggml ctx size =    0.11 MB\n",
      "llama_model_load_internal: mem required  = 6983.72 MB (+ 3046.88 MB per state)\n",
      "llama_new_context_with_model: kv self size  = 3046.88 MB\n",
      "ggml_metal_init: allocating\n",


LlamaIndex supports embeddings from OpenAI, Azure, and Langchain. But if this isn't enough, you can also implement any embeddings model!

The example below uses Instructor Embeddings ([install/setup details here](https://huggingface.co/hkunlp/instructor-large)), and implements a custom embeddings class. Instructor embeddings work by providing text, as well as "instructions" on the domain of the text to embed. This is helpful when embedding text from a very specific and specialized topic.

If you're opening this Notebook on colab, you will probably need to install LlamaIndex 🦙.

# !pip install InstructorEmbedding torch transformers sentence-transformers

from InstructorEmbedding import INSTRUCTOR

from llama_index.core.bridge.pydantic import PrivateAttr

from llama_index.core.embeddings import BaseEmbedding

instructor_model_name: str = "hkunlp/instructor-large",

instruction: str = "Represent a document for semantic search:",

async def _aget_query_embedding(self, query: str) -> List[float]:

async def _aget_text_embedding(self, text: str) -> List[float]:

def _get_query_embedding(self, query: str) -> List[float]:

embeddings = self._model.encode([[self._instruction, query]])

def _get_text_embedding(self, text: str) -> List[float]:

embeddings = self._model.encode([[self._instruction, text]])

def _get_text_embeddings(self, texts: List[str]) -> List[List[float]]:

[[self._instruction, text] for text in texts]

from llama_index.core import SimpleDirectoryReader, VectorStoreIndex

stdout

load INSTRUCTOR_Transformer
max_seq_length  512

# if running for the first time, will download model weights first!

The author wrote short stories and also worked on programming, specifically on an IBM 1401 computer in 9th grade. They used an early version of Fortran and had to type programs on punch cards. Later on, they got a microcomputer, a TRS-80, and started programming more extensively, writing simple games and a word processor. They initially planned to study philosophy in college but eventually switched to AI.

response = index.as_query_engine().query("What did the author do growing up?")