{ "cells": [ { "attachments": {}, "cell_type": "markdown", "id": "d6509c3a", "metadata": {}, "source": [ "\"Open" ] }, { "cell_type": "markdown", "id": "c0d8b66c", "metadata": {}, "source": [ "# Wandb Callback Handler\n", "\n", "[Weights & Biases Prompts](https://docs.wandb.ai/guides/prompts) is a suite of LLMOps tools built for the development of LLM-powered applications.\n", "\n", "The `WandbCallbackHandler` is integrated with W&B Prompts to visualize and inspect the execution flow of your index construction, or querying over your index and more. You can use this handler to persist your created indices as W&B Artifacts allowing you to version control your indices.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "49c3527e", "metadata": {}, "outputs": [], "source": [ "%pip install llama-index-callbacks-wandb\n", "%pip install llama-index-llms-openai" ] }, { "cell_type": "code", "execution_count": null, "id": "612f35ad", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "OpenAI API key configured\n" ] } ], "source": [ "import os\n", "from getpass import getpass\n", "\n", "if os.getenv(\"OPENAI_API_KEY\") is None:\n", " os.environ[\"OPENAI_API_KEY\"] = getpass(\n", " \"Paste your OpenAI key from:\"\n", " \" https://platform.openai.com/account/api-keys\\n\"\n", " )\n", "assert os.getenv(\"OPENAI_API_KEY\", \"\").startswith(\n", " \"sk-\"\n", "), \"This doesn't look like a valid OpenAI API key\"\n", "print(\"OpenAI API key configured\")" ] }, { "cell_type": "code", "execution_count": null, "id": "78a29d9a", "metadata": {}, "outputs": [], "source": [ "from llama_index.core.callbacks import CallbackManager\n", "from llama_index.core.callbacks import LlamaDebugHandler\n", "from llama_index.callbacks.wandb import WandbCallbackHandler\n", "from llama_index.core import (\n", " VectorStoreIndex,\n", " SimpleDirectoryReader,\n", " SimpleKeywordTableIndex,\n", " StorageContext,\n", ")\n", "from llama_index.llms.openai import OpenAI" ] }, { "cell_type": "markdown", "id": "e6feb252", "metadata": {}, "source": [ "## Setup LLM" ] }, { "cell_type": "code", "execution_count": null, "id": "d22fee33", "metadata": {}, "outputs": [], "source": [ "from llama_index.core import Settings\n", "\n", "Settings.llm = OpenAI(model=\"gpt-4\", temperature=0)" ] }, { "cell_type": "markdown", "id": "8790f4c7", "metadata": {}, "source": [ "## W&B Callback Manager Setup" ] }, { "cell_type": "markdown", "id": "8a32b984-772e-4832-945e-cb6fc7be9e0b", "metadata": {}, "source": [ "**Option 1**: Set Global Evaluation Handler" ] }, { "cell_type": "code", "execution_count": null, "id": "2a3b9d22-cd67-4fb5-9785-254e58179a02", "metadata": {}, "outputs": [], "source": [ "import llama_index.core\n", "from llama_index.core import set_global_handler\n", "\n", "set_global_handler(\"wandb\", run_args={\"project\": \"llamaindex\"})\n", "wandb_callback = llama_index.core.global_handler" ] }, { "cell_type": "markdown", "id": "d1755516-f8ad-458e-b52f-f7665c023e43", "metadata": {}, "source": [ "**Option 2**: Manually Configure Callback Handler\n", "\n", "Also configure a debugger handler for extra notebook visibility." ] }, { "cell_type": "code", "execution_count": null, "id": "defa9155-daca-4a8f-8ca6-87d1ee98f084", "metadata": {}, "outputs": [], "source": [ "llama_debug = LlamaDebugHandler(print_trace_on_end=True)\n", "\n", "# wandb.init args\n", "run_args = dict(\n", " project=\"llamaindex\",\n", ")\n", "\n", "wandb_callback = WandbCallbackHandler(run_args=run_args)\n", "\n", "Settings.callback_manager = CallbackManager([llama_debug, wandb_callback])" ] }, { "cell_type": "markdown", "id": "c4cf969a", "metadata": {}, "source": [ "> After running the above cell, you will get the W&B run page URL. Here you will find a trace table with all the events tracked using [Weights and Biases' Prompts](https://docs.wandb.ai/guides/prompts) feature." ] }, { "cell_type": "markdown", "id": "a4a7c101", "metadata": {}, "source": [ "## 1. Indexing" ] }, { "attachments": {}, "cell_type": "markdown", "id": "e5d31f80", "metadata": {}, "source": [ "Download Data" ] }, { "cell_type": "code", "execution_count": null, "id": "1e7ad71e", "metadata": {}, "outputs": [], "source": [ "!mkdir -p 'data/paul_graham/'\n", "!wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/examples/data/paul_graham/paul_graham_essay.txt' -O 'data/paul_graham/paul_graham_essay.txt'" ] }, { "cell_type": "code", "execution_count": null, "id": "d1011596", "metadata": {}, "outputs": [], "source": [ "docs = SimpleDirectoryReader(\"./data/paul_graham/\").load_data()" ] }, { "cell_type": "code", "execution_count": null, "id": "d3d6975c", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "**********\n", "Trace: index_construction\n", " |_node_parsing -> 0.295179 seconds\n", " |_chunking -> 0.293976 seconds\n", " |_embedding -> 0.494492 seconds\n", " |_embedding -> 0.346162 seconds\n", "**********\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "\u001b[34m\u001b[1mwandb\u001b[0m: Logged trace tree to W&B.\n" ] } ], "source": [ "index = VectorStoreIndex.from_documents(docs)" ] }, { "cell_type": "markdown", "id": "0a948efc", "metadata": {}, "source": [ "### 1.1 Persist Index as W&B Artifacts" ] }, { "cell_type": "code", "execution_count": null, "id": "8ad58e67", "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "\u001b[34m\u001b[1mwandb\u001b[0m: Adding directory to artifact (/Users/loganmarkewich/llama_index/docs/examples/callbacks/wandb/run-20230801_152955-ds93prxa/files/storage)... Done. 0.0s\n" ] } ], "source": [ "wandb_callback.persist_index(index, index_name=\"simple_vector_store\")" ] }, { "cell_type": "markdown", "id": "7ed156a6", "metadata": {}, "source": [ "### 1.2 Download Index from W&B Artifacts" ] }, { "cell_type": "code", "execution_count": null, "id": "dc35f448", "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "\u001b[34m\u001b[1mwandb\u001b[0m: 3 of 3 files downloaded. \n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "**********\n", "Trace: index_construction\n", "**********\n" ] } ], "source": [ "from llama_index.core import load_index_from_storage\n", "\n", "storage_context = wandb_callback.load_storage_context(\n", " artifact_url=\"ayut/llamaindex/simple_vector_store:v0\"\n", ")\n", "\n", "# Load the index and initialize a query engine\n", "index = load_index_from_storage(\n", " storage_context,\n", ")" ] }, { "cell_type": "markdown", "id": "ae4de4a9", "metadata": {}, "source": [ "## 2. Query Over Index" ] }, { "cell_type": "code", "execution_count": null, "id": "42221465", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "**********\n", "Trace: query\n", " |_query -> 2.695958 seconds\n", " |_retrieve -> 0.806379 seconds\n", " |_embedding -> 0.802871 seconds\n", " |_synthesize -> 1.8893 seconds\n", " |_llm -> 1.842434 seconds\n", "**********\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "\u001b[34m\u001b[1mwandb\u001b[0m: Logged trace tree to W&B.\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "The text does not provide information on what the author did growing up.\n" ] } ], "source": [ "query_engine = index.as_query_engine()\n", "response = query_engine.query(\"What did the author do growing up?\")\n", "print(response, sep=\"\\n\")" ] }, { "cell_type": "markdown", "id": "c49ff101", "metadata": {}, "source": [ "## Close W&B Callback Handler\n", "\n", "When we are done tracking our events we can close the wandb run." ] }, { "cell_type": "code", "execution_count": null, "id": "28ef6a7b", "metadata": {}, "outputs": [], "source": [ "wandb_callback.finish()" ] } ], "metadata": { "kernelspec": { "display_name": "Python 3 (ipykernel)", "language": "python", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3" } }, "nbformat": 4, "nbformat_minor": 5 }