Assistant\n",
+ "\"\"\"\n",
+ ")\n",
+ "system = \"You are a helpful, respectful and honest assistant. Always answer as helpfully as possible, while being safe. Please ensure that your responses are positive in nature.\"\n",
+ "prompt = 'What is the fastest land animal?'\n",
+ "prompt = NEMOTRON_PROMPT_TEMPLATE.format(prompt=prompt, system=system)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "c0e7ec48-3865-4811-a175-6a6142240da2",
+ "metadata": {},
+ "source": [
+ "#### Step 2: Create the Triton Client"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "5b8476aa-6364-408a-840a-a8afb83ea08e",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "triton_url = \"llm:8000\"\n",
+ "client = HttpTritonClient(triton_url)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "e9c89aca-1eae-4afc-b520-ed14ebef2606",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "pload = {\n",
+ " 'prompt':[[prompt]],\n",
+ " 'tokens':64,\n",
+ " 'temperature':1.0,\n",
+ " 'top_k':1,\n",
+ " 'top_p':0,\n",
+ " 'beam_width':1,\n",
+ " 'repetition_penalty':1.0,\n",
+ " 'length_penalty':1.0\n",
+ "}"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "0c32f8d8-f383-4faf-997a-1c8412412a5a",
+ "metadata": {},
+ "source": [
+ "#### Step 3: Load the Model and Generate Response"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "b02689bd-36d6-41fe-8636-9deeb52bc68f",
+ "metadata": {
+ "scrolled": true
+ },
+ "outputs": [],
+ "source": [
+ "model_name = \"ensemble\"\n",
+ "client.load_model(model_name)\n",
+ "val = client.request(model_name, **pload)\n",
+ "print(val)"
+ ]
+ }
+ ],
+ "metadata": {
+ "kernelspec": {
+ "display_name": "Python 3 (ipykernel)",
+ "language": "python",
+ "name": "python3"
+ },
+ "language_info": {
+ "codemirror_mode": {
+ "name": "ipython",
+ "version": 3
+ },
+ "file_extension": ".py",
+ "mimetype": "text/x-python",
+ "name": "python",
+ "nbconvert_exporter": "python",
+ "pygments_lexer": "ipython3",
+ "version": "3.10.12"
+ }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}
diff --git a/notebooks/01-llm-streaming-client.ipynb b/notebooks/01-llm-streaming-client.ipynb
index 7168e1cc..13c5e9c1 100644
--- a/notebooks/01-llm-streaming-client.ipynb
+++ b/notebooks/01-llm-streaming-client.ipynb
@@ -9,9 +9,9 @@
"This notebook demonstrates how to stream responses from the LLM. \n",
"\n",
"### Triton Inference Server\n",
- "The LLM has been deployed to [NVIDIA Triton Inference Server](https://developer.nvidia.com/triton-inference-server) and leverages NVIDIA TensorRT-LLM (TRT-LLM), so it's optimized for low latency and high throughput inference. \n",
+ "The LLM has been deployed to [NVIDIA Triton Inference Server](https://developer.nvidia.com/triton-inference-server) and leverages NVIDIA TensorRT-LLM (TRT-LLM), so it's optimized for low latency and high throughput inference.\n",
"\n",
- "The **Triton client** is used to communicate with the inference server hosting the LLM. \n",
+ "The **Triton client** is used to communicate with the inference server hosting the LLM and is available in [Langchain](https://github.com/langchain-ai/langchain/tree/master/libs/partners/nvidia-trt). \n",
"\n",
"### Streaming LLM Responses\n",
"TRT-LLM on its own can provide drastic improvements to LLM response latency, but streaming can take the user-experience to the next level. Instead of waiting for an entire response to be returned from the LLM, chunks of it can be processed as soon as they are available. This helps reduce the perceived latency by the user. "
@@ -40,7 +40,7 @@
},
{
"cell_type": "code",
- "execution_count": null,
+ "execution_count": 34,
"id": "42a2f2cb",
"metadata": {},
"outputs": [],
@@ -81,20 +81,31 @@
},
{
"cell_type": "code",
- "execution_count": null,
+ "execution_count": 35,
"id": "5670011e-f52b-4c16-be4d-8a782b622541",
"metadata": {},
"outputs": [],
"source": [
- "from triton_trt_llm import GrpcTritonClient\n",
+ "from langchain_nvidia_trt.llms import TritonTensorRTLLM\n",
"\n",
"triton_url = \"llm:8001\"\n",
- "client = GrpcTritonClient(triton_url)"
+ "pload = {\n",
+ " 'tokens':300,\n",
+ " 'server_url': triton_url,\n",
+ " 'model_name': \"ensemble\",\n",
+ " 'temperature':1.0,\n",
+ " 'top_k':1,\n",
+ " 'top_p':0,\n",
+ " 'beam_width':1,\n",
+ " 'repetition_penalty':1.0,\n",
+ " 'length_penalty':1.0\n",
+ "}\n",
+ "client = TritonTensorRTLLM(**pload)"
]
},
{
"cell_type": "markdown",
- "id": "eec1fda2-e974-4c7d-b656-0b81e452cefb",
+ "id": "03676629-33d8-46d8-b5fc-557d526609b4",
"metadata": {},
"source": [
"Additional inputs to the LLM can be modified:\n",
@@ -106,25 +117,6 @@
"- length_penalty: 1 means no penalty for length of generation"
]
},
- {
- "cell_type": "code",
- "execution_count": null,
- "id": "6907e68d-8ab9-4f94-b6a0-825446261bfc",
- "metadata": {},
- "outputs": [],
- "source": [
- "pload = {\n",
- " 'prompt':[[prompt]],\n",
- " 'tokens':300,\n",
- " 'temperature':1.0,\n",
- " 'top_k':1,\n",
- " 'top_p':0,\n",
- " 'beam_width':1,\n",
- " 'repetition_penalty':1.0,\n",
- " 'length_penalty':1.0\n",
- "}"
- ]
- },
{
"cell_type": "markdown",
"id": "c526b20b-258a-4eb7-87e6-5430d57e32ea",
@@ -136,18 +128,17 @@
{
"cell_type": "code",
"execution_count": null,
- "id": "d389f26d-e68d-4466-a7cb-89e62683af0b",
+ "id": "274e4164-7460-4471-8625-90562237cf11",
"metadata": {},
"outputs": [],
"source": [
- "model_name = \"ensemble\"\n",
- "client.load_model(model_name)\n",
"import time\n",
"import random\n",
"\n",
"start_time = time.time()\n",
"tokens_generated = 0\n",
- "for val in client.request_streaming(model_name, request_id = str(random.getrandbits(64)), **pload):\n",
+ "\n",
+ "for val in client.stream(prompt):\n",
" tokens_generated += 1\n",
" print(val, end=\"\", flush=True)\n",
"\n",
diff --git a/notebooks/02_langchain_simple.ipynb b/notebooks/02_langchain_simple.ipynb
index e4415b22..8b5cc548 100644
--- a/notebooks/02_langchain_simple.ipynb
+++ b/notebooks/02_langchain_simple.ipynb
@@ -34,73 +34,47 @@
"id": "9dd510db",
"metadata": {},
"source": [
- "### Step 1: Integrate TensorRT-LLM to LangChain [*(Model I/O)*](https://python.langchain.com/docs/modules/model_io/)\n",
- "\n",
- "#### Custom TRT-LLM Langchain integration.\n",
- "Langchain allows you to [create custom wrappers for your LLM](https://python.langchain.com/docs/modules/model_io/models/llms/custom_llm) in case you want to use your own LLM or a different wrapper than the one that is supported in LangChain. Since we are using a custom Llama2 model hosted on Triton with TRT-LLM, we have written a custom wrapper for our LLM. \n",
- "\n",
- "Below is a snippet of the custom wrapper. Take a look at ```triton_trt_llm.py``` under `integrations/langchain/llms` of the project root for the full implementation.\n",
- "```\n",
- "class TensorRTLLM(LLM):\n",
- " server_url: str = Field(None, alias=\"server_url\")\n",
- "\n",
- " # some of the optional arguments\n",
- " model_name: str = \"ensemble\"\n",
- " temperature: Optional[float] = 1.0\n",
- " top_p: Optional[float] = 0\n",
- "\n",
- " @property\n",
- " def _llm_type(self) -> str:\n",
- " return \"triton_tensorrt\"\n",
- "\n",
- " def _call(\n",
- " self,\n",
- " prompt: str,\n",
- " run_manager: Optional[CallbackManagerForLLMRun] = None,\n",
- " **kwargs,\n",
- " ) -> str:\n",
- " \"\"\"\n",
- " Args:\n",
- " prompt: The prompt to pass into the model.\n",
- " stop: A list of strings to stop generation when encountered\n",
- "\n",
- " Returns:\n",
- " The string generated by the model\n",
- " \"\"\"\n",
- "\n",
- "```\n",
- "\n",
- "A ```_call``` method that takes in a string, some optional stop words, and returns a string. Take a look at ```triton_trt_llm.py``` for the code of LangChain wrapper for a Llama2 model deployed on Triton with TRT-LLM.\n",
- "\n",
- "``llm = TensorRTLLM(server_url =\"triton_ip:8001\", model_name=\"ensemble\", callbacks=callbacks, tokens=500)``\n",
- "\n",
- "\n",
- " \n",
- "WARNING! Be sure to replace server_url with the address and port of the Triton server hosting the LLM. If you are running this as part of the generative AI Workflow, you don't have to replace the Triton url.\n",
- "\n",
- "
"
+ "### Step 1: Integrate TensorRT-LLM to LangChain [*(Connector)*](https://docs.llamaindex.ai/en/stable/examples/llm/nvidia_tensorrt.html)"
]
},
{
"cell_type": "code",
"execution_count": null,
+ "id": "f7c4cb93-95c7-4665-b091-7719d996acb8",
"metadata": {},
"outputs": [],
- "source": []
+ "source": [
+ "from langchain_nvidia_trt.llms import TritonTensorRTLLM\n",
+ "\n",
+ "# Connect to the TRT-LLM Llama-2 model running on the Triton server at the url below\n",
+ "# Replace \"llm\" with the url of the system where llama2 is hosted\n",
+ "triton_url = \"llm:8001\"\n",
+ "pload = {\n",
+ " 'tokens':500,\n",
+ " 'server_url': triton_url,\n",
+ " 'model_name': \"ensemble\"\n",
+ "}\n",
+ "llm = TritonTensorRTLLM(**pload)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "552c7fc8",
+ "metadata": {},
+ "source": [
+ "#### Note: Follow this step for nemotron models\n",
+ "1. In case you have deployed a trt-llm optimized nemotron model following steps [here](../RetrievalAugmentedGeneration/README.md#6-qa-chatbot----nemotron-model), execute the cell below by uncommenting the lines. Here we use a custom wrapper for talking with the model server."
+ ]
},
{
"cell_type": "code",
"execution_count": null,
- "id": "f7c4cb93-95c7-4665-b091-7719d996acb8",
+ "id": "4f125aea",
"metadata": {},
"outputs": [],
"source": [
- "from triton_trt_llm import TensorRTLLM\n",
- "from langchain.callbacks import streaming_stdout\n",
- "\n",
- "callbacks = [streaming_stdout.StreamingStdOutCallbackHandler()]\n",
- "# Connect to the TRT-LLM Llama-2 model running on the Triton server at the url below\n",
- "llm = TensorRTLLM(server_url =\"llm:8001\", model_name=\"ensemble\", callbacks=callbacks, tokens=500)"
+ "# from triton_trt_llm import TensorRTLLM\n",
+ "# llm = TensorRTLLM(server_url =\"llm:8000\", model_name=\"ensemble\", tokens=500, streaming=False)"
]
},
{
@@ -198,12 +172,12 @@
"import time\n",
"from langchain.text_splitter import SentenceTransformersTokenTextSplitter\n",
"TEXT_SPLITTER_MODEL = \"intfloat/e5-large-v2\"\n",
- "TEXT_SPLITTER_CHUNCK_SIZE = 510\n",
+ "TEXT_SPLITTER_TOKENS_PER_CHUNK = 510\n",
"TEXT_SPLITTER_CHUNCK_OVERLAP = 200\n",
"\n",
"text_splitter = SentenceTransformersTokenTextSplitter(\n",
" model_name=TEXT_SPLITTER_MODEL,\n",
- " chunk_size=TEXT_SPLITTER_CHUNCK_SIZE,\n",
+ " tokens_per_chunk=TEXT_SPLITTER_TOKENS_PER_CHUNK,\n",
" chunk_overlap=TEXT_SPLITTER_CHUNCK_OVERLAP,\n",
")\n",
"start_time = time.time()\n",
@@ -320,24 +294,11 @@
"metadata": {},
"source": [
"### Step 6: Compose a streamed answer using a Chain\n",
- "We have already integrated the Llama2 TRT LLM into LangChain with a custom wrapper, loaded and transformed documents, and generated and stored document embeddings in a vector database. To finish the pipeline, we need to add a few more LangChain components and combine all the components together with a [chain](https://python.langchain.com/docs/modules/chains/).\n",
+ "We have already integrated the Llama2 TRT LLM with the help of LangChain connector, loaded and transformed documents, and generated and stored document embeddings in a vector database. To finish the pipeline, we need to add a few more LangChain components and combine all the components together with a [chain](https://python.langchain.com/docs/modules/chains/).\n",
"\n",
- "A [LangChain chain](https://python.langchain.com/docs/modules/chains/) combines components together. In this case, we use a [RetrievalQA chain](https://js.langchain.com/docs/modules/chains/popular/vector_db_qa/), which is a chain type for question-answering against a vector index. It combines a *Retriever* and a *question answering (QA) chain*.\n",
- "\n",
- "We pass it 3 of our LangChain components:\n",
- "- Our instance of the LLM (from step 1).\n",
- "- A [retriever](https://python.langchain.com/docs/modules/data_connection/retrievers/), which is an interface that returns documents given an unstructured query. In this case, we use our vector store as the retriever.\n",
- "- Our prompt template constructed from the prompt format for Llama2 (from step 2)\n",
- "\n",
- "```\n",
- "qa_chain = RetrievalQA.from_chain_type(\n",
- " llm,\n",
- " retriever=vectorstore.as_retriever(),\n",
- " chain_type_kwargs={\"prompt\": LLAMA_PROMPT}\n",
- ")\n",
- "```\n",
+ "A [LangChain chain](https://python.langchain.com/docs/modules/chains/) combines components together. In this case, we use [Langchain Expression Language](https://python.langchain.com/docs/expression_language/why) to build a chain.\n",
"\n",
- "Lastly, we pass a user query to the chain and stream the result. "
+ "We formulate the prompt placeholders (context and question) and pipe it to our trt-llm connector as shown below and finally stream the result."
]
},
{
@@ -347,16 +308,17 @@
"metadata": {},
"outputs": [],
"source": [
- "from langchain.chains import RetrievalQA\n",
+ "from langchain_core.runnables import RunnablePassthrough\n",
"import time\n",
"\n",
- "qa_chain = RetrievalQA.from_chain_type(\n",
- " llm,\n",
- " retriever=vectorstore.as_retriever(),\n",
- " chain_type_kwargs={\"prompt\": LLAMA_PROMPT}\n",
+ "chain = (\n",
+ " {\"context\": vectorstore.as_retriever(), \"question\": RunnablePassthrough()}\n",
+ " | LLAMA_PROMPT\n",
+ " | llm\n",
")\n",
"start_time = time.time()\n",
- "result = qa_chain({\"query\": question})\n",
+ "for token in chain.stream(question):\n",
+ " print(token, end=\"\", flush=True)\n",
"print(f\"\\n--- {time.time() - start_time} seconds ---\")"
]
}
diff --git a/notebooks/03_llama_index_simple.ipynb b/notebooks/03_llama_index_simple.ipynb
index 3c067663..4657b4d7 100644
--- a/notebooks/03_llama_index_simple.ipynb
+++ b/notebooks/03_llama_index_simple.ipynb
@@ -38,7 +38,7 @@
"source": [
"### Step 1: Integrate TensorRT-LLM to LangChain *and* LlamaIndex\n",
"#### Customized LangChain LLM in LlamaIndex\n",
- "As noted in the previous notebook, Langchain allows you to [create custom wrappers for your LLM](https://python.langchain.com/docs/modules/model_io/models/llms/custom_llm) in case you want to use your own LLM or a different wrapper than the one that is supported in LangChain. Since we are using a custom Llama2 model hosted on Triton with TRT-LLM, we have written a custom wrapper for our LLM. \n",
+ "Langchain allows you to create custom wrappers for your LLM in case you want to use your own LLM or a different wrapper than the one that is supported in LangChain. Since we are using LlamaIndex, we have written a custom langchain wrapper compatible with LlamaIndex.\n",
"\n",
"We can easily take a custom LLM that has been wrapped for LangChain and plug it into [LlamaIndex as an LLM](https://docs.llamaindex.ai/en/stable/understanding/using_llms/using_llms.html#using-llms)! We use the [LlamaIndex LangChainLLM library](https://gpt-index.readthedocs.io/en/latest/api_reference/llms/langchain.html) so the LangChain LLM can be used in LlamaIndex. \n",
"\n",
@@ -206,12 +206,12 @@
"\n",
"\n",
"TEXT_SPLITTER_MODEL = \"intfloat/e5-large-v2\"\n",
- "TEXT_SPLITTER_CHUNCK_SIZE = 510\n",
+ "TEXT_SPLITTER_TOKENS_PER_CHUNK = 510\n",
"TEXT_SPLITTER_CHUNCK_OVERLAP = 200\n",
"\n",
"text_splitter = SentenceTransformersTokenTextSplitter(\n",
" model_name=TEXT_SPLITTER_MODEL,\n",
- " chunk_size=TEXT_SPLITTER_CHUNCK_SIZE,\n",
+ " tokens_per_chunk=TEXT_SPLITTER_TOKENS_PER_CHUNK,\n",
" chunk_overlap=TEXT_SPLITTER_CHUNCK_OVERLAP,\n",
")\n",
"\n",
diff --git a/notebooks/04_llamaindex_hier_node_parser.ipynb b/notebooks/04_llamaindex_hier_node_parser.ipynb
index 5f36ec3d..640188f8 100644
--- a/notebooks/04_llamaindex_hier_node_parser.ipynb
+++ b/notebooks/04_llamaindex_hier_node_parser.ipynb
@@ -316,13 +316,14 @@
"source": [
"from typing import Callable, Optional\n",
"\n",
- "from llama_index.utils import globals_helper\n",
+ "from llama_index.utils import globals_helper, get_tokenizer\n",
"from llama_index.schema import MetadataMode\n",
"\n",
"class LimitRetrievedNodesLength:\n",
"\n",
" def __init__(self, limit: int = 2500, tokenizer: Optional[Callable] = None):\n",
- " self._tokenizer = tokenizer or globals_helper.tokenizer\n",
+ " self._tokenizer = tokenizer or get_tokenizer()\n",
+ "\n",
" self.limit = limit\n",
"\n",
" def postprocess_nodes(self, nodes, query_bundle):\n",
diff --git a/notebooks/06_AI_playground.ipynb b/notebooks/06_AI_playground.ipynb
deleted file mode 100644
index 94c401b5..00000000
--- a/notebooks/06_AI_playground.ipynb
+++ /dev/null
@@ -1,276 +0,0 @@
-{
- "cells": [
- {
- "cell_type": "markdown",
- "id": "65d4e77f-cc3d-402d-987b-a9c361ce99e1",
- "metadata": {},
- "source": [
- "# Notebook 6: RAG with NVIDIA AI Playground and Langchain\n",
- "\n",
- "**NVIDIA AI Playground** on NGC allows developers to experience state of the art LLMs accelerated on NVIDIA DGX Cloud with NVIDIA TensorRT nd Triton Inference Server. Developers get **free credits for 10K requests** to any of the available models. Sign up process is easy. Follow the instructions here \n",
- "\n",
- "This notebook demonstrates how to use LangChain and NVIDIA AI Playground to build a chatbot that references a custom knowledge-base. \n",
- "\n",
- "Suppose you have some text documents (PDF, blog, Notion pages, etc.) and want to ask questions related to the contents of those documents. LLMs, given their proficiency in understanding text, are a great tool for this. \n",
- "\n",
- "### [LangChain](https://python.langchain.com/docs/get_started/introduction)\n",
- "[**LangChain**](https://python.langchain.com/docs/get_started/introduction) provides a simple framework for connecting LLMs to your own data sources. Since LLMs are both only trained up to a fixed point in time and do not contain knowledge that is proprietary to an Enterprise, they can't answer questions about new or proprietary knowledge. LangChain solves this problem."
- ]
- },
- {
- "cell_type": "markdown",
- "id": "23110dc7-df11-4413-b8d3-9db2f3e1187a",
- "metadata": {},
- "source": [
- ""
- ]
- },
- {
- "cell_type": "markdown",
- "id": "e6806b29-cbc1-4dc8-b378-5c6333eb7229",
- "metadata": {},
- "source": [
- "### Step 1: Load Documents [*(Retrieval)*](https://python.langchain.com/docs/modules/data_connection/)\n",
- "LangChain provides a variety of [document loaders](https://python.langchain.com/docs/integrations/document_loaders) that load various types of documents (HTML, PDF, code) from many different sources and locations (private s3 buckets, public websites).\n",
- "\n",
- "Document loaders load data from a source as **Documents**. A **Document** is a piece of text (the page_content) and associated metadata. Document loaders provide a ``load`` method for loading data as documents from a configured source. \n",
- "\n",
- "In this example, we use a LangChain [`UnstructuredFileLoader`](https://python.langchain.com/docs/integrations/document_loaders/unstructured_file) to load a research paper about Llama2 from Meta.\n",
- "\n",
- "[Here](https://python.langchain.com/docs/integrations/document_loaders) are some of the other document loaders available from LangChain."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "id": "2c09ede8-58cc-44d5-b46c-b1060dd578c5",
- "metadata": {},
- "outputs": [],
- "source": [
- "!wget -O \"llama2_paper.pdf\" -nc --user-agent=\"Mozilla\" https://arxiv.org/pdf/2307.09288.pdf"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "id": "43793ae9-f25f-4e9d-a7a5-13b64f02b115",
- "metadata": {},
- "outputs": [],
- "source": [
- "from langchain.document_loaders import UnstructuredFileLoader\n",
- "loader = UnstructuredFileLoader(\"llama2_paper.pdf\")\n",
- "data = loader.load()"
- ]
- },
- {
- "cell_type": "markdown",
- "id": "21c8568d-c20e-4d79-a828-c2bcd2b6cd20",
- "metadata": {},
- "source": [
- "### Step 2: Transform Documents [*(Retrieval)*](https://python.langchain.com/docs/modules/data_connection/)\n",
- "Once documents have been loaded, they are often transformed. One method of transformation is known as **chunking**, which breaks down large pieces of text, for example, a long document, into smaller segments. This technique is valuable because it helps [optimize the relevance of the content returned from the vector database](https://www.pinecone.io/learn/chunking-strategies/). \n",
- "\n",
- "LangChain provides a [variety of document transformers](https://python.langchain.com/docs/integrations/document_transformers/), such as text splitters. In this example, we use a [``SentenceTransformersTokenTextSplitter``](https://api.python.langchain.com/en/latest/text_splitter/langchain.text_splitter.SentenceTransformersTokenTextSplitter.html#langchain.text_splitter.SentenceTransformersTokenTextSplitter). The ``SentenceTransformersTokenTextSplitter`` is a specialized text splitter for use with the sentence-transformer models. The default behaviour is to split the text into chunks that fit the token window of the sentence transformer model that you would like to use. This sentence transformer model is used to generate the embeddings from documents. \n",
- "\n",
- "There are some nuanced complexities to text splitting since semantically related text, in theory, should be kept together. "
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "id": "5fea998e-c776-4763-b811-b42b6372fcb9",
- "metadata": {},
- "outputs": [],
- "source": [
- "import time\n",
- "from langchain.text_splitter import SentenceTransformersTokenTextSplitter\n",
- "TEXT_SPLITTER_MODEL = \"intfloat/e5-large-v2\"\n",
- "TEXT_SPLITTER_CHUNCK_SIZE = 510\n",
- "TEXT_SPLITTER_CHUNCK_OVERLAP = 200\n",
- "\n",
- "text_splitter = SentenceTransformersTokenTextSplitter(\n",
- " model_name=TEXT_SPLITTER_MODEL,\n",
- " chunk_size=TEXT_SPLITTER_CHUNCK_SIZE,\n",
- " chunk_overlap=TEXT_SPLITTER_CHUNCK_OVERLAP,\n",
- ")\n",
- "start_time = time.time()\n",
- "documents = text_splitter.split_documents(data)\n",
- "print(f\"--- {time.time() - start_time} seconds ---\")"
- ]
- },
- {
- "cell_type": "markdown",
- "id": "c72a07ff-119a-468c-ab68-f8f6e9f7c6fc",
- "metadata": {},
- "source": [
- "### Step 3: Generate Embeddings and Store Embeddings in the Vector Store [*(Retrieval)*](https://python.langchain.com/docs/modules/data_connection/)\n",
- "\n",
- "#### a) Generate Embeddings\n",
- "[Embeddings](https://python.langchain.com/docs/modules/data_connection/text_embedding/) for documents are created by vectorizing the document text; this vectorization captures the semantic meaning of the text. This allows you to quickly and efficiently find other pieces of text that are similar. The embedding model used below is [intfloat/e5-large-v2](https://huggingface.co/intfloat/e5-large-v2).\n",
- "\n",
- "LangChain provides a wide variety of [embedding models](https://python.langchain.com/docs/integrations/text_embedding) from many providers and makes it simple to swap out the models. \n",
- "\n",
- "When a user sends in their query, the query is also embedded using the same embedding model that was used to embed the documents. As explained earlier, this allows to find similar (relevant) documents to the user's query. \n",
- "\n",
- "#### b) Store Document Embeddings in the Vector Store\n",
- "Once the document embeddings are generated, they are stored in a vector store so that at query time we can:\n",
- "1) Embed the user query and\n",
- "2) Retrieve the embedding vectors that are most similar to the embedding query.\n",
- "\n",
- "A vector store takes care of storing the embedded data and performing a vector search.\n",
- "\n",
- "LangChain provides support for a [great selection of vector stores](https://python.langchain.com/docs/integrations/vectorstores/). \n",
- "\n",
- "\n",
- " \n",
- "⚠️ For this workflow, [Milvus](https://milvus.io/) vector database is running as a microservice. \n",
- "\n",
- "
"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "id": "3ee3be33-5f78-4937-9e07-a7cc93d5d7df",
- "metadata": {},
- "outputs": [],
- "source": [
- "from langchain.embeddings import HuggingFaceEmbeddings\n",
- "from langchain.vectorstores import Milvus\n",
- "import torch\n",
- "import time\n",
- "\n",
- "#Running the model on CPU as we want to conserve gpu memory.\n",
- "#In the production deployment (API server shown as part of the 5th notebook we run the model on GPU)\n",
- "model_name = \"intfloat/e5-large-v2\"\n",
- "model_kwargs = {\"device\": \"cuda:0\"}\n",
- "encode_kwargs = {\"normalize_embeddings\": False}\n",
- "hf_embeddings = HuggingFaceEmbeddings(\n",
- " model_name=model_name,\n",
- " model_kwargs=model_kwargs,\n",
- " encode_kwargs=encode_kwargs,\n",
- ")\n",
- "start_time = time.time()\n",
- "vectorstore = Milvus.from_documents(documents=documents, embedding=hf_embeddings, connection_args={\"host\": \"milvus\", \"port\": \"19530\"})\n",
- "print(f\"--- {time.time() - start_time} seconds ---\")"
- ]
- },
- {
- "cell_type": "markdown",
- "id": "bc71d0c9-aeab-4a26-959f-5493ca1f2f02",
- "metadata": {},
- "source": [
- "### Step 4: Sign up to NVIDIA AI Playground \n",
- "\n",
- "**NVIDIA AI Playground** on NGC allows developers to experience state of the art LLMs accelerated on NVIDIA DGX Cloud with NVIDIA TensorRT nd Triton Inference Server. Developers get **free credits for 10K requests** to any of the available models. Sign up process is easy. Follow the instructions here and replace the API key below. For this example we will be using the llama2 13B model"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "id": "8cddceb0-1f58-48c4-aba4-cc1962ed8806",
- "metadata": {},
- "outputs": [],
- "source": [
- "import os\n",
- "from nv_aiplay import GeneralLLM\n",
- "from langchain.callbacks import streaming_stdout\n",
- "\n",
- "os.environ['NVAPI_KEY'] = \"REPLACE_WITH_API_KEY\"\n",
- "\n",
- "callbacks = [streaming_stdout.StreamingStdOutCallbackHandler()]\n",
- "\n",
- "llm = GeneralLLM(\n",
- " temperature=0.2,\n",
- " max_tokens=300,\n",
- " streaming=True,\n",
- " callbacks = callbacks\n",
- ")"
- ]
- },
- {
- "cell_type": "markdown",
- "id": "8f25e8af-9aff-460a-a214-87a3e285469d",
- "metadata": {},
- "source": [
- "### Step 5: Ask a question without context\n",
- "\n",
- "Send request to the llm without any context from the vector DB. The answer is generic and irrelvant."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "id": "3fced303-d825-4797-9e4a-9a8c6e5f82aa",
- "metadata": {},
- "outputs": [],
- "source": [
- "question = \"Can you talk about the safety features of llama2 chat?\"\n",
- "answer = llm(question)"
- ]
- },
- {
- "cell_type": "markdown",
- "id": "225fc935-4bc9-4957-a476-6d23d147e664",
- "metadata": {},
- "source": [
- "### Step 6: Compose a streamed answer using a Chain\n",
- "We have already integrated the Llama2 TRT LLM into LangChain with a custom wrapper, loaded and transformed documents, and generated and stored document embeddings in a vector database. To finish the pipeline, we need to add a few more LangChain components and combine all the components together with a [chain](https://python.langchain.com/docs/modules/chains/).\n",
- "\n",
- "A [LangChain chain](https://python.langchain.com/docs/modules/chains/) combines components together. In this case, we use a [RetrievalQA chain](https://js.langchain.com/docs/modules/chains/popular/vector_db_qa/), which is a chain type for question-answering against a vector index. It combines a *Retriever* and a *question answering (QA) chain*.\n",
- "\n",
- "We pass it 3 of our LangChain components:\n",
- "- Our instance of the LLM (from step 1).\n",
- "- A [retriever](https://python.langchain.com/docs/modules/data_connection/retrievers/), which is an interface that returns documents given an unstructured query. In this case, we use our vector store as the retriever.\n",
- "- Our prompt template constructed from the prompt format for Llama2 (from step 2)\n",
- "\n",
- "```\n",
- "qa_chain = RetrievalQA.from_chain_type(\n",
- " llm,\n",
- " retriever=vectorstore.as_retriever(),\n",
- " chain_type_kwargs={\"prompt\": LLAMA_PROMPT}\n",
- ")\n",
- "```\n",
- "\n",
- "Lastly, we pass a user query to the chain and stream the result. "
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "id": "65b7bac7-fda6-4584-b469-43f9eba75468",
- "metadata": {},
- "outputs": [],
- "source": [
- "from langchain.chains import RetrievalQA\n",
- "\n",
- "qa_chain = RetrievalQA.from_chain_type(\n",
- " llm,\n",
- " retriever=vectorstore.as_retriever()\n",
- ")\n",
- "result = qa_chain({\"query\": question})"
- ]
- }
- ],
- "metadata": {
- "kernelspec": {
- "display_name": "Python 3 (ipykernel)",
- "language": "python",
- "name": "python3"
- },
- "language_info": {
- "codemirror_mode": {
- "name": "ipython",
- "version": 3
- },
- "file_extension": ".py",
- "mimetype": "text/x-python",
- "name": "python",
- "nbconvert_exporter": "python",
- "pygments_lexer": "ipython3",
- "version": "3.10.13"
- }
- },
- "nbformat": 4,
- "nbformat_minor": 5
-}
diff --git a/notebooks/imgs/chrome_flags_fix_media_device_access_error.png b/notebooks/imgs/chrome_flags_fix_media_device_access_error.png
new file mode 100755
index 00000000..a7cc6f51
Binary files /dev/null and b/notebooks/imgs/chrome_flags_fix_media_device_access_error.png differ
diff --git a/notebooks/imgs/grace_answer.png b/notebooks/imgs/grace_answer.png
new file mode 100644
index 00000000..76afc8c0
Binary files /dev/null and b/notebooks/imgs/grace_answer.png differ
diff --git a/notebooks/imgs/grace_answer_with_riva.png b/notebooks/imgs/grace_answer_with_riva.png
new file mode 100755
index 00000000..458555bf
Binary files /dev/null and b/notebooks/imgs/grace_answer_with_riva.png differ
diff --git a/notebooks/imgs/grace_noanswer.png b/notebooks/imgs/grace_noanswer.png
new file mode 100644
index 00000000..f957762c
Binary files /dev/null and b/notebooks/imgs/grace_noanswer.png differ
diff --git a/notebooks/imgs/grace_noanswer_with_riva.png b/notebooks/imgs/grace_noanswer_with_riva.png
new file mode 100755
index 00000000..5a1e2548
Binary files /dev/null and b/notebooks/imgs/grace_noanswer_with_riva.png differ
diff --git a/notebooks/imgs/media_device_access_error.png b/notebooks/imgs/media_device_access_error.png
new file mode 100755
index 00000000..d7427477
Binary files /dev/null and b/notebooks/imgs/media_device_access_error.png differ
diff --git a/notebooks/requirements.txt b/notebooks/requirements.txt
index 6e68824f..b430e057 100644
--- a/notebooks/requirements.txt
+++ b/notebooks/requirements.txt
@@ -1,13 +1,16 @@
fastapi==0.104.1
uvicorn[standard]==0.24.0
python-multipart==0.0.6
-langchain==0.0.330
-tritonclient[all]==2.39.0
+langchain==0.0.352
unstructured[all-docs]==0.11.2
sentence-transformers==2.2.2
-llama-index==0.9.13
+llama-index==0.9.22
dataclass-wizard==0.22.2
opencv-python==4.8.0.74
llama-hub==0.0.43
pymilvus==2.3.1
-jupyterlab==4.0.8
\ No newline at end of file
+jupyterlab==4.0.8
+langchain-nvidia-trt==0.0.1rc0
+langchain-core==0.1.3
+langchain-nvidia-ai-endpoints==0.0.1
+atlassian-python-api==3.41.4
\ No newline at end of file
diff --git a/tools/__init__.py b/tools/__init__.py
new file mode 100644
index 00000000..e42268fe
--- /dev/null
+++ b/tools/__init__.py
@@ -0,0 +1,14 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
\ No newline at end of file
diff --git a/evaluation/01_synthetic_data_generation.ipynb b/tools/evaluation/01_synthetic_data_generation.ipynb
similarity index 96%
rename from evaluation/01_synthetic_data_generation.ipynb
rename to tools/evaluation/01_synthetic_data_generation.ipynb
index abbc4eb8..262e8e79 100644
--- a/evaluation/01_synthetic_data_generation.ipynb
+++ b/tools/evaluation/01_synthetic_data_generation.ipynb
@@ -207,7 +207,7 @@
"source": [
"#### a) AI Playground LLM generator\n",
"\n",
- "**NVIDIA AI Playground** on NGC allows developers to experience state of the art LLMs accelerated on NVIDIA DGX Cloud with NVIDIA TensorRT nd Triton Inference Server. Developers get **free credits for 10K requests** to any of the available models. Sign up process is easy. follow the steps here. \n",
+ "**NVIDIA AI Playground** on NGC allows developers to experience state of the art LLMs accelerated on NVIDIA DGX Cloud with NVIDIA TensorRT nd Triton Inference Server. Developers get **free credits for 10K requests** to any of the available models. Sign up process is easy. follow the steps here. \n",
"\n",
"We are going to use theAI playground'ss `llama2-70B `LLM to generate the Question-Answer pairs."
]
@@ -219,8 +219,8 @@
"metadata": {},
"outputs": [],
"source": [
- "# import the relevant libraries\n",
- "from nv_aiplay import GeneralLLM"
+ "# import the relevant libraries from langchain\n",
+ "from langchain_nvidia_ai_endpoints import ChatNVIDIA"
]
},
{
@@ -239,9 +239,9 @@
"outputs": [],
"source": [
"import os\n",
- "os.environ['NVAPI_KEY'] = \"nvapi-*\"\n",
+ "os.environ['NVIDIA_API_KEY'] = \"nvapi-*\"\n",
"\n",
- "llm = GeneralLLM(\n",
+ "llm = ChatNVIDIA(\n",
" model=\"llama2_70b\",\n",
" temperature=0.2,\n",
" max_tokens=300\n",
@@ -256,7 +256,7 @@
"outputs": [],
"source": [
"# check the output\n",
- "answer = llm(context)"
+ "answer = llm.invoke(context)"
]
},
{
@@ -317,7 +317,7 @@
"source": [
"# Synthetic Data Post-processing \n",
"\n",
- "So far, the generated JSON file structure embeds `gt_context`, `document` and the `question`, `answer` pair.\n",
+ "So far, the generated JSON file structure embeds `gt_context`, `document`, the `question` and `gt_answer` pair.\n",
"\n",
"In order to evaluate Retrieval Augmented Generation (RAG) systems, we need to add the RAG results fields (To be populated in the next notebook):\n",
" - `contexts`: Retrieved documents by the retriever \n",
diff --git a/evaluation/02_filling_RAG_outputs_for_Evaluation.ipynb b/tools/evaluation/02_filling_RAG_outputs_for_Evaluation.ipynb
similarity index 88%
rename from evaluation/02_filling_RAG_outputs_for_Evaluation.ipynb
rename to tools/evaluation/02_filling_RAG_outputs_for_Evaluation.ipynb
index 8e05d6cd..b82497b8 100644
--- a/evaluation/02_filling_RAG_outputs_for_Evaluation.ipynb
+++ b/tools/evaluation/02_filling_RAG_outputs_for_Evaluation.ipynb
@@ -28,7 +28,7 @@
"#### Define the LLM\n",
"Here we are using a local llm on triton and the address and gRPC port that the Triton is available on. \n",
"\n",
- "***If you are using AI Playground (no local GPU) replace, the code in the cell below with the following: ***\n",
+ "***If you are using AI Playground (no local GPU) replace, the code in the cell two cells below with the following: ***\n",
"\n",
"```\n",
"import os\n",
@@ -43,6 +43,17 @@
"```"
]
},
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "a18dfc7b",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "%%capture\n",
+ "!test -d dataset || unzip dataset.zip"
+ ]
+ },
{
"cell_type": "code",
"execution_count": null,
@@ -125,12 +136,12 @@
"\n",
"# setup the text splitter\n",
"TEXT_SPLITTER_MODEL = \"intfloat/e5-large-v2\"\n",
- "TEXT_SPLITTER_CHUNCK_SIZE = 510\n",
+ "TEXT_SPLITTER_TOKENS_PER_CHUNK = 510\n",
"TEXT_SPLITTER_CHUNCK_OVERLAP = 200\n",
"\n",
"text_splitter = SentenceTransformersTokenTextSplitter(\n",
" model_name=TEXT_SPLITTER_MODEL,\n",
- " chunk_size=TEXT_SPLITTER_CHUNCK_SIZE,\n",
+ " tokens_per_chunk=TEXT_SPLITTER_TOKENS_PER_CHUNK,\n",
" chunk_overlap=TEXT_SPLITTER_CHUNCK_OVERLAP,\n",
")\n",
"\n",
@@ -234,6 +245,60 @@
"set_global_service_context(service_context)"
]
},
+ {
+ "cell_type": "markdown",
+ "id": "44e10c13",
+ "metadata": {},
+ "source": [
+ "Ingest the dataset using the /uploadDocument endpoint in the chain-server."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "acdc51db",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "import os\n",
+ "import requests\n",
+ "import mimetypes\n",
+ "\n",
+ "def upload_document(file_path, url):\n",
+ " headers = {\n",
+ " 'accept': 'application/json'\n",
+ " }\n",
+ " mime_type, _ = mimetypes.guess_type(file_path)\n",
+ " files = {\n",
+ " 'file': (file_path, open(file_path, 'rb'), mime_type)\n",
+ " }\n",
+ " response = requests.post(url, headers=headers, files=files)\n",
+ "\n",
+ " return response.text\n",
+ "\n",
+ "def upload_pdf_files(folder_path, upload_url):\n",
+ " for files in os.listdir(folder_path):\n",
+ " _, ext = os.path.splitext(files)\n",
+ " # Ingest only pdf files\n",
+ " if ext.lower() == \".pdf\":\n",
+ " file_path = os.path.join(folder_path, files)\n",
+ " print(upload_document(file_path, upload_url))"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "823c89f9",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "import time\n",
+ "\n",
+ "start_time = time.time()\n",
+ "upload_pdf_files(\"dataset\", \"http://query:8081/uploadDocument\")\n",
+ "print(f\"--- {time.time() - start_time} seconds ---\")"
+ ]
+ },
{
"attachments": {},
"cell_type": "markdown",
@@ -259,7 +324,13 @@
"from llama_index.vector_stores import MilvusVectorStore\n",
"\n",
"# store\n",
- "vector_store = MilvusVectorStore(uri=\"http://milvus:19530\", dim=1024, overwrite=False)\n",
+ "vector_store = MilvusVectorStore(uri=\"http://milvus:19530\",\n",
+ " dim=1024,\n",
+ " collection_name=\"document_store_ivfflat\",\n",
+ " index_config={\"index_type\": \"IVF_FLAT\", \"nlist\": 64},\n",
+ " search_config={\"nprobe\": 16},\n",
+ " overwrite=False\n",
+ ")\n",
"storage_context = StorageContext.from_defaults(vector_store=vector_store)\n",
"index = VectorStoreIndex.from_vector_store(vector_store)"
]
diff --git a/evaluation/03_eval_ragas.ipynb b/tools/evaluation/03_eval_ragas.ipynb
similarity index 100%
rename from evaluation/03_eval_ragas.ipynb
rename to tools/evaluation/03_eval_ragas.ipynb
diff --git a/evaluation/04_Human_Like_RAG_Evaluation-AIP.ipynb b/tools/evaluation/04_Human_Like_RAG_Evaluation-AIP.ipynb
similarity index 97%
rename from evaluation/04_Human_Like_RAG_Evaluation-AIP.ipynb
rename to tools/evaluation/04_Human_Like_RAG_Evaluation-AIP.ipynb
index c28e30eb..fb310d08 100644
--- a/evaluation/04_Human_Like_RAG_Evaluation-AIP.ipynb
+++ b/tools/evaluation/04_Human_Like_RAG_Evaluation-AIP.ipynb
@@ -92,8 +92,9 @@
"invoke_url = \"https://api.nvcf.nvidia.com/v2/nvcf/pexec/functions/0e349b44-440a-44e1-93e9-abe8dcb27158\" #Llama 2 70B\n",
"fetch_url_format = \"https://api.nvcf.nvidia.com/v2/nvcf/pexec/status/\"\n",
"\n",
+ "# do not remove Bearer from Authorization, replace with api key\n",
"headers = {\n",
- " \"Authorization\": \"Bearer REPLACE_THIS_WITH_API_KEY\",\n",
+ " \"Authorization\": \"Bearer \",\n",
" \"Accept\": \"application/json\",\n",
"}\n"
]
@@ -221,7 +222,7 @@
" response_body = response.json()\n",
" llama_judge_responses.append(response_body['choices'][0]['message']['content'])\n",
" except Exception as e:\n",
- " print(\"pass\")\n",
+ " print(\"Exception:\", e)\n",
" llama_judge_responses.append(None)\n"
]
},
@@ -256,11 +257,11 @@
"\n",
" # Extract and print the rating and explanation\n",
" llama_ratings.append(int(rating_match.group(1)) if rating_match else None)\n",
- " llama_explanations.append(explanation_match.group(1) if explanation_match else None)\n",
+ " llama_explanations.append(explanation_match.group(1) if explanation_match else response)\n",
" except Exception as e:\n",
- " print(\"pass\")\n",
+ " print(\"Exception\", e)\n",
" llama_ratings.append(None)\n",
- " llama_explanations.append(None)\n"
+ " llama_explanations.append(response)\n"
]
},
{
diff --git a/evaluation/Dockerfile.eval b/tools/evaluation/Dockerfile.eval
similarity index 78%
rename from evaluation/Dockerfile.eval
rename to tools/evaluation/Dockerfile.eval
index e76734f3..d0fc923a 100644
--- a/evaluation/Dockerfile.eval
+++ b/tools/evaluation/Dockerfile.eval
@@ -4,21 +4,21 @@ FROM python:3.10-slim
WORKDIR /app
#COPY notebooks
-COPY ./evaluation/*.ipynb .
+COPY ./tools/evaluation/*.ipynb .
RUN mkdir -p /app/imgs
COPY ./notebooks/dataset.zip .
-COPY ./evaluation/imgs/* imgs/
+COPY ./tools/evaluation/imgs/* imgs/
COPY ./integrations/langchain/llms/triton_trt_llm.py .
COPY ./integrations/langchain/llms/nv_aiplay.py .
-COPY ./evaluation/requirements.txt .
+COPY ./tools/evaluation/requirements.txt .
-COPY ./evaluation/qa_generation.json .
+COPY ./tools/evaluation/qa_generation.json .
# Run pip dependencies
RUN pip3 install -r requirements.txt
diff --git a/evaluation/imgs/ragas.png b/tools/evaluation/imgs/ragas.png
similarity index 100%
rename from evaluation/imgs/ragas.png
rename to tools/evaluation/imgs/ragas.png
diff --git a/evaluation/imgs/synthetic_data_pipeline.png b/tools/evaluation/imgs/synthetic_data_pipeline.png
similarity index 100%
rename from evaluation/imgs/synthetic_data_pipeline.png
rename to tools/evaluation/imgs/synthetic_data_pipeline.png
diff --git a/evaluation/qa_generation.json b/tools/evaluation/qa_generation.json
similarity index 100%
rename from evaluation/qa_generation.json
rename to tools/evaluation/qa_generation.json
diff --git a/evaluation/requirements.txt b/tools/evaluation/requirements.txt
similarity index 54%
rename from evaluation/requirements.txt
rename to tools/evaluation/requirements.txt
index 001a80ca..7c471e45 100644
--- a/evaluation/requirements.txt
+++ b/tools/evaluation/requirements.txt
@@ -1,5 +1,4 @@
-langchain==0.0.348
-tritonclient[all]==2.39.0
+langchain==0.0.352
unstructured[all-docs]==0.11.2
sentence-transformers==2.2.2
llama-index==0.9.13
@@ -8,3 +7,7 @@ pymilvus==2.3.1
jupyterlab==4.0.8
ragas==0.0.21
seaborn==0.13.0
+langchain-core==0.1.3
+langchain-nvidia-ai-endpoints==0.0.1
+langchain-nvidia-trt==0.0.1rc0
+atlassian-python-api==3.41.4
diff --git a/tools/observability/__init__.py b/tools/observability/__init__.py
new file mode 100644
index 00000000..e42268fe
--- /dev/null
+++ b/tools/observability/__init__.py
@@ -0,0 +1,14 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
\ No newline at end of file
diff --git a/tools/observability/llamaindex/__init__.py b/tools/observability/llamaindex/__init__.py
new file mode 100644
index 00000000..e42268fe
--- /dev/null
+++ b/tools/observability/llamaindex/__init__.py
@@ -0,0 +1,14 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
\ No newline at end of file
diff --git a/tools/observability/llamaindex/opentelemetry_callback.py b/tools/observability/llamaindex/opentelemetry_callback.py
new file mode 100644
index 00000000..48d89e85
--- /dev/null
+++ b/tools/observability/llamaindex/opentelemetry_callback.py
@@ -0,0 +1,198 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from opentelemetry.trace import Tracer, get_tracer, set_span_in_context, Status, StatusCode
+from opentelemetry.trace.span import Span
+from opentelemetry.context import Context, get_current, attach, detach
+from typing import Any, Dict, List, Optional, Callable
+from llama_index.callbacks.base_handler import BaseCallbackHandler
+from llama_index.callbacks.base import CallbackManager
+from llama_index.callbacks.schema import CBEventType, EventPayload, BASE_TRACE_EVENT
+from llama_index.callbacks.token_counting import get_llm_token_counts, TokenCountingEvent
+from llama_index.utilities.token_counting import TokenCounter
+from llama_index.utils import get_tokenizer
+from dataclasses import dataclass
+from contextvars import ContextVar
+import threading
+
+global_root_trace = ContextVar("trace", default=None)
+
+@dataclass
+class SpanWithContext:
+ """Object for tracking a span, its context, and its context token"""
+ span: Span
+ context: Context
+ token: object
+
+ def __init__(self, span: Span, context: Context, token: object, thread_identity):
+ self.span = span
+ self.context = context
+ self.token = token
+ self.thread_identity = thread_identity
+
+class OpenTelemetryCallbackHandler(BaseCallbackHandler):
+ """Callback handler for creating OpenTelemetry traces from llamaindex traces and events."""
+
+ def __init__(
+ self,
+ tracer: Optional[Tracer] = get_tracer(__name__),
+ tokenizer: Optional[Callable[[str], List]] = None,
+ ) -> None:
+ """Initializes the OpenTelemetryCallbackHandler.
+
+ Args:
+ tracer: Optional[Tracer]: A OpenTelemetry tracer used to create OpenTelemetry spans
+ """
+ super().__init__(event_starts_to_ignore=[], event_ends_to_ignore=[])
+ self._tracer = tracer
+ self._event_map: Dict[str, SpanWithContext] = {}
+ self.tokenizer = tokenizer or get_tokenizer()
+ self._token_counter = TokenCounter(tokenizer=self.tokenizer)
+
+ def start_trace(self, trace_id: Optional[str] = None) -> None:
+ trace_name = "llamaindex.trace"
+ if trace_id is not None:
+ trace_name = "llamaindex.trace." + trace_id
+ span = self._tracer.start_span(trace_name)
+ ctx = set_span_in_context(span)
+ token = attach(ctx)
+ global_root_trace.set(SpanWithContext(span=span, context=ctx, token=token, thread_identity=threading.get_ident()))
+
+ def end_trace(
+ self,
+ trace_id: Optional[str] = None,
+ trace_map: Optional[Dict[str, List[str]]] = None,
+ ) -> None:
+ root_trace = global_root_trace.get()
+ if root_trace is not None:
+ if root_trace.thread_identity == threading.get_ident():
+ detach(root_trace.token)
+ root_trace.span.end()
+
+ def on_event_start(
+ self,
+ event_type: CBEventType,
+ payload: Optional[Dict[str, Any]] = None,
+ event_id: str = "",
+ parent_id: str = "",
+ **kwargs: Any,
+ ) -> str:
+ parent_ctx = None
+ # Case where the parent of this event is another event
+ if parent_id in self._event_map:
+ parent_ctx = self._event_map[parent_id].context
+ # Case where the parent of this event is the root trace, and the root trace exists
+ elif parent_id is BASE_TRACE_EVENT and global_root_trace.get() is not None:
+ parent_ctx = global_root_trace.get().context
+ # Case where the parent of this event is the root trace, but the trace does not exist
+ else:
+ return
+
+ span_prefix = "llamaindex.event."
+ span = self._tracer.start_span(span_prefix + event_type.value, context=parent_ctx)
+ ctx = set_span_in_context(span)
+ token = attach(ctx)
+ self._event_map[event_id] = SpanWithContext(span=span, context=ctx, token=token, thread_identity=threading.get_ident())
+
+ span.set_attribute("event_id", event_id)
+ if payload is not None:
+ if event_type is CBEventType.QUERY:
+ span.set_attribute("query.text", payload[EventPayload.QUERY_STR])
+ elif event_type is CBEventType.RETRIEVE:
+ pass
+ elif event_type is CBEventType.EMBEDDING:
+ span.set_attribute("embedding.model", payload[EventPayload.SERIALIZED]['model_name'])
+ span.set_attribute("embedding.batch_size", payload[EventPayload.SERIALIZED]['embed_batch_size'])
+ span.set_attribute("embedding.class_name", payload[EventPayload.SERIALIZED]['class_name'])
+ elif event_type is CBEventType.SYNTHESIZE:
+ span.set_attribute("synthesize.query_text", payload[EventPayload.QUERY_STR])
+ elif event_type is CBEventType.CHUNKING:
+ for i, chunk in enumerate(payload[EventPayload.CHUNKS]):
+ span.set_attribute(f"chunk.{i}", chunk)
+ elif event_type is CBEventType.TEMPLATING:
+ if payload[EventPayload.QUERY_WRAPPER_PROMPT]:
+ span.set_attribute("query_wrapper_prompt", payload[EventPayload.QUERY_WRAPPER_PROMPT])
+ if payload[EventPayload.SYSTEM_PROMPT]:
+ span.set_attribute("system_prompt", payload[EventPayload.SYSTEM_PROMPT])
+ if payload[EventPayload.TEMPLATE]:
+ span.set_attribute("template", payload[EventPayload.TEMPLATE])
+ if payload[EventPayload.TEMPLATE_VARS]:
+ for key, var in payload[EventPayload.TEMPLATE_VARS].items():
+ span.set_attribute(f"template_variables.{key}", var)
+ elif event_type is CBEventType.LLM:
+ span.set_attribute("llm.class_name", payload[EventPayload.SERIALIZED]['class_name'])
+ span.set_attribute("llm.formatted_prompt", payload[EventPayload.PROMPT])
+ span.set_attribute("llm.additional_kwargs", str(payload[EventPayload.ADDITIONAL_KWARGS]))
+ elif event_type is CBEventType.NODE_PARSING:
+ span.set_attribute("node_parsing.num_documents", len(payload[EventPayload.DOCUMENTS]))
+ elif event_type is CBEventType.EXCEPTION:
+ span.set_status(Status(StatusCode.ERROR))
+ span.record_exception(payload[EventPayload.EXCEPTION])
+ return event_id
+
+ def on_event_end(
+ self,
+ event_type: CBEventType,
+ payload: Optional[Dict[str, Any]] = None,
+ event_id: str = "",
+ **kwargs: Any,
+ ) -> None:
+ if event_id in self._event_map:
+ span = self._event_map[event_id].span
+ span.set_attribute("event_id", event_id)
+ if payload is not None:
+ if event_type is CBEventType.QUERY:
+ pass
+ elif event_type is CBEventType.RETRIEVE:
+ for i, node_with_score in enumerate(payload[EventPayload.NODES]):
+ node = node_with_score.node
+ score = node_with_score.score
+ span.set_attribute(f"query.node.{i}.id", node.hash)
+ span.set_attribute(f"query.node.{i}.score", score)
+ span.set_attribute(f"query.node.{i}.text", node.text)
+ elif event_type is CBEventType.EMBEDDING:
+ texts = payload[EventPayload.CHUNKS]
+ vectors = payload[EventPayload.EMBEDDINGS]
+ total_chunk_tokens = 0
+ for text, vector in zip(texts, vectors) :
+ span.set_attribute(f"embedding_text_{texts.index(text)}", text)
+ span.set_attribute(f"embedding_vector_{vectors.index(vector)}", vector)
+ total_chunk_tokens +=self._token_counter.get_string_tokens(text)
+ span.set_attribute(f"embedding_token_usage", total_chunk_tokens)
+ elif event_type is CBEventType.SYNTHESIZE:
+ pass
+ elif event_type is CBEventType.CHUNKING:
+ pass
+ elif event_type is CBEventType.TEMPLATING:
+ pass
+ elif event_type is CBEventType.LLM:
+ span.set_attribute("response.text", str(
+ payload.get(EventPayload.RESPONSE, "")
+ ) or str(payload.get(EventPayload.COMPLETION, ""))
+ )
+ token_counts = get_llm_token_counts(self._token_counter, payload, event_id)
+ span.set_attribute("llm_prompt.token_usage", token_counts.prompt_token_count)
+ span.set_attribute("llm_completion.token_usage", token_counts.completion_token_count)
+ span.set_attribute("total_tokens_used", token_counts.total_token_count)
+ elif event_type is CBEventType.NODE_PARSING:
+ span.set_attribute("node_parsing.num_nodes", len(payload[EventPayload.NODES]))
+ elif event_type is CBEventType.EXCEPTION:
+ span.set_status(Status(StatusCode.ERROR))
+ span.record_exception(payload[EventPayload.EXCEPTION])
+ if self._event_map[event_id].thread_identity == threading.get_ident():
+ detach(self._event_map[event_id].token)
+ self._event_map.pop(event_id, None)
+ span.end()
+
\ No newline at end of file