a0c8464e58
Build Package / build (ubuntu-latest) (push) Failing after 1s
CodeQL / Analyze (python) (push) Failing after 1s
Core Typecheck / core-typecheck (push) Failing after 1s
Linting / lint (push) Failing after 1s
llama-dev tests / test-llama-dev (push) Failing after 1s
Publish Sub-Package to PyPI if Needed / publish_subpackage_if_needed (push) Has been skipped
Sync Docs to Developer Hub / sync-docs (push) Failing after 0s
Build Package / build (windows-latest) (push) Has been cancelled
400 lines
9.6 KiB
Plaintext
400 lines
9.6 KiB
Plaintext
{
|
|
"cells": [
|
|
{
|
|
"attachments": {},
|
|
"cell_type": "markdown",
|
|
"id": "d6509c3a",
|
|
"metadata": {},
|
|
"source": [
|
|
"<a href=\"https://colab.research.google.com/github/run-llama/llama_index/blob/main/docs/examples/observability/WandbCallbackHandler.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
|
|
]
|
|
},
|
|
{
|
|
"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
|
|
}
|