L3-Sentence_window_retrieval.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "21f9146a-c6f3-4a78-b3fb-0d262492e87c",
   "metadata": {},
   "source": [
    "# Lesson 3: Sentence Window Retrieval"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "541cadae-916c-42da-93ff-75d7f788ee8d",
   "metadata": {
    "height": 47,
    "ExecuteTime": {
     "end_time": "2023-12-07T08:19:39.114578Z",
     "start_time": "2023-12-07T08:19:38.756763Z"
    }
   },
   "outputs": [],
   "source": [
    "import warnings\n",
    "warnings.filterwarnings('ignore')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "78f48098-16d8-4209-b722-1ec6a0220c96",
   "metadata": {
    "height": 98,
    "ExecuteTime": {
     "end_time": "2023-12-07T08:19:48.750973Z",
     "start_time": "2023-12-07T08:19:39.106844Z"
    }
   },
   "outputs": [],
   "source": [
    "import utils\n",
    "import os\n",
    "from llama_index.llms import AzureOpenAI\n",
    "from llama_index.embeddings import AzureOpenAIEmbedding\n",
    "from dotenv import load_dotenv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "outputs": [],
   "source": [
    "load_dotenv()\n",
    "\n",
    "api_key = os.getenv(\"AZURE_OPENAI_API_KEY\")\n",
    "azure_endpoint = os.getenv(\"AZURE_OPENAI_ENDPOINT\")\n",
    "api_version = os.getenv(\"OPENAI_API_VERSION\")\n",
    "\n",
    "llm = AzureOpenAI(\n",
    "    model=\"gpt-4\",\n",
    "    deployment_name=\"gpt-4-32k\",\n",
    "    api_key=api_key,\n",
    "    azure_endpoint=azure_endpoint,\n",
    "    api_version=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=\"text-embedding-ada-002\",\n",
    "    deployment_name=\"embedding-ada-002\",\n",
    "    api_key=api_key,\n",
    "    azure_endpoint=azure_endpoint,\n",
    "    api_version=api_version,\n",
    ")"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2023-12-07T08:19:48.765263Z",
     "start_time": "2023-12-07T08:19:48.755808Z"
    }
   },
   "id": "8dfdcab28b2b8e67"
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "9ee55360-1fc4-41cc-bd49-82027797ea40",
   "metadata": {
    "height": 98,
    "tags": [],
    "ExecuteTime": {
     "end_time": "2023-12-07T08:19:49.628220Z",
     "start_time": "2023-12-07T08:19:48.765361Z"
    }
   },
   "outputs": [],
   "source": [
    "from llama_index import SimpleDirectoryReader\n",
    "\n",
    "documents = SimpleDirectoryReader(\n",
    "    input_files=[\"./eBook-How-to-Build-a-Career-in-AI.pdf\"]\n",
    ").load_data()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "a12ada81-5c1c-47c9-b7b4-ba621a80bbcd",
   "metadata": {
    "height": 81,
    "tags": [],
    "ExecuteTime": {
     "end_time": "2023-12-07T08:19:49.634826Z",
     "start_time": "2023-12-07T08:19:49.631233Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<class 'list'> \n",
      "\n",
      "41 \n",
      "\n",
      "<class 'llama_index.schema.Document'>\n",
      "Doc ID: bf212331-f913-4a3e-9356-53b467883945\n",
      "Text: PAGE 1Founder, DeepLearning.AICollected Insights from Andrew Ng\n",
      "How to  Build Your Career in AIA Simple Guide\n"
     ]
    }
   ],
   "source": [
    "print(type(documents), \"\\n\")\n",
    "print(len(documents), \"\\n\")\n",
    "print(type(documents[0]))\n",
    "print(documents[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "f2f662e4-3fb8-40c6-acc5-e8510348d113",
   "metadata": {
    "height": 64,
    "tags": [],
    "ExecuteTime": {
     "end_time": "2023-12-07T08:19:49.648651Z",
     "start_time": "2023-12-07T08:19:49.636501Z"
    }
   },
   "outputs": [],
   "source": [
    "from llama_index import Document\n",
    "\n",
    "document = Document(text=\"\\n\\n\".join([doc.text for doc in documents]))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "46ff01ea-b5b0-4e65-8565-b2444812bd84",
   "metadata": {},
   "source": [
    "## Window-sentence retrieval setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "6a194b93-f975-4d38-babe-218d3aae6117",
   "metadata": {
    "height": 149,
    "tags": [],
    "ExecuteTime": {
     "end_time": "2023-12-07T08:19:51.131789Z",
     "start_time": "2023-12-07T08:19:49.640627Z"
    }
   },
   "outputs": [],
   "source": [
    "from llama_index.node_parser import SentenceWindowNodeParser\n",
    "\n",
    "# create the sentence window node parser w/ default settings\n",
    "node_parser = SentenceWindowNodeParser.from_defaults(\n",
    "    window_size=3,\n",
    "    window_metadata_key=\"window\",\n",
    "    original_text_metadata_key=\"original_text\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "88973de6-be8f-4653-aca3-8d0e884d9470",
   "metadata": {
    "height": 64,
    "tags": [],
    "ExecuteTime": {
     "end_time": "2023-12-07T08:19:51.146542Z",
     "start_time": "2023-12-07T08:19:51.133347Z"
    }
   },
   "outputs": [],
   "source": [
    "text = \"hello. how are you? I am fine!  \"\n",
    "\n",
    "nodes = node_parser.get_nodes_from_documents([Document(text=text)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "860e40a8-6f50-4345-b314-c4d9349681c6",
   "metadata": {
    "height": 30,
    "tags": [],
    "ExecuteTime": {
     "end_time": "2023-12-07T08:19:51.152645Z",
     "start_time": "2023-12-07T08:19:51.144069Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['hello. ', 'how are you? ', 'I am fine!  ']\n"
     ]
    }
   ],
   "source": [
    "print([x.text for x in nodes])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "b4334e02-c423-4555-8446-4794eadccd0a",
   "metadata": {
    "height": 30,
    "tags": [],
    "ExecuteTime": {
     "end_time": "2023-12-07T08:19:51.206559Z",
     "start_time": "2023-12-07T08:19:51.150998Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "hello.  how are you?  I am fine!  \n"
     ]
    }
   ],
   "source": [
    "print(nodes[1].metadata[\"window\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "03fea59c-ed0f-4cc8-87b4-871798edb094",
   "metadata": {
    "height": 64,
    "tags": [],
    "ExecuteTime": {
     "end_time": "2023-12-07T08:19:51.210111Z",
     "start_time": "2023-12-07T08:19:51.159840Z"
    }
   },
   "outputs": [],
   "source": [
    "text = \"hello. foo bar. cat dog. mouse\"\n",
    "\n",
    "nodes = node_parser.get_nodes_from_documents([Document(text=text)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "c329808d-248f-422e-bce0-1ac1ecba79a5",
   "metadata": {
    "height": 30,
    "tags": [],
    "ExecuteTime": {
     "end_time": "2023-12-07T08:19:51.229508Z",
     "start_time": "2023-12-07T08:19:51.170473Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['hello. ', 'foo bar. ', 'cat dog. ', 'mouse']\n"
     ]
    }
   ],
   "source": [
    "print([x.text for x in nodes])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "cfd1eccb-823d-4f6f-8d6c-a9064dc76922",
   "metadata": {
    "height": 30,
    "tags": [],
    "ExecuteTime": {
     "end_time": "2023-12-07T08:19:51.242561Z",
     "start_time": "2023-12-07T08:19:51.177448Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "hello.  foo bar.  cat dog. \n"
     ]
    }
   ],
   "source": [
    "print(nodes[0].metadata[\"window\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fea97dda-8457-4f32-a2c7-26ae92eaf0b4",
   "metadata": {},
   "source": [
    "### Building the index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "5102d65b-3584-4a25-88be-7d5dbc70c678",
   "metadata": {
    "height": 149,
    "tags": [],
    "ExecuteTime": {
     "end_time": "2023-12-07T08:19:56.118776Z",
     "start_time": "2023-12-07T08:19:51.188142Z"
    }
   },
   "outputs": [],
   "source": [
    "from llama_index import ServiceContext\n",
    "\n",
    "sentence_context = ServiceContext.from_defaults(\n",
    "    llm=llm,\n",
    "    embed_model=\"local:BAAI/bge-small-en-v1.5\",\n",
    "    # embed_model=\"local:BAAI/bge-large-en-v1.5\"\n",
    "    node_parser=node_parser,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "ba4e7371-7a8e-451c-a5e7-e9e3d7371614",
   "metadata": {
    "height": 336,
    "tags": [],
    "ExecuteTime": {
     "end_time": "2023-12-07T08:19:57.013416Z",
     "start_time": "2023-12-07T08:19:56.119169Z"
    }
   },
   "outputs": [],
   "source": [
    "# This block of code is optional to check\n",
    "# if an index file exist, then it will load it\n",
    "# if not, it will rebuild it\n",
    "\n",
    "import os\n",
    "from llama_index import VectorStoreIndex, StorageContext, load_index_from_storage\n",
    "from llama_index import load_index_from_storage\n",
    "\n",
    "if not os.path.exists(\"./sentence_index\"):\n",
    "    sentence_index = VectorStoreIndex.from_documents(\n",
    "        [document], service_context=sentence_context\n",
    "    )\n",
    "\n",
    "    sentence_index.storage_context.persist(persist_dir=\"./sentence_index\")\n",
    "else:\n",
    "    sentence_index = load_index_from_storage(\n",
    "        StorageContext.from_defaults(persist_dir=\"./sentence_index\"),\n",
    "        service_context=sentence_context\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6b934879-e2c3-44f8-b98c-0300dc6389a9",
   "metadata": {},
   "source": [
    "### Building the postprocessor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "9f22f435-a90a-447c-9e63-8aebec1e968d",
   "metadata": {
    "height": 98,
    "tags": [],
    "ExecuteTime": {
     "end_time": "2023-12-07T08:19:57.019464Z",
     "start_time": "2023-12-07T08:19:57.015682Z"
    }
   },
   "outputs": [],
   "source": [
    "from llama_index.indices.postprocessor import MetadataReplacementPostProcessor\n",
    "\n",
    "postproc = MetadataReplacementPostProcessor(\n",
    "    target_metadata_key=\"window\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "77fad243-7ea8-4a81-b85e-91c3e638d932",
   "metadata": {
    "height": 98,
    "tags": [],
    "ExecuteTime": {
     "end_time": "2023-12-07T08:19:57.034546Z",
     "start_time": "2023-12-07T08:19:57.020414Z"
    }
   },
   "outputs": [],
   "source": [
    "from llama_index.schema import NodeWithScore\n",
    "from copy import deepcopy\n",
    "\n",
    "scored_nodes = [NodeWithScore(node=x, score=1.0) for x in nodes]\n",
    "nodes_old = [deepcopy(n) for n in nodes]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "946c4219-1d6e-4717-8c8f-40db659ea517",
   "metadata": {
    "height": 30,
    "tags": [],
    "ExecuteTime": {
     "end_time": "2023-12-07T08:19:57.049203Z",
     "start_time": "2023-12-07T08:19:57.031076Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": "'foo bar. '"
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "nodes_old[1].text"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "89900cb4-5d03-43f6-9d1d-de1350bfd299",
   "metadata": {
    "height": 30,
    "tags": [],
    "ExecuteTime": {
     "end_time": "2023-12-07T08:19:57.086081Z",
     "start_time": "2023-12-07T08:19:57.037341Z"
    }
   },
   "outputs": [],
   "source": [
    "replaced_nodes = postproc.postprocess_nodes(scored_nodes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "72f98170-e368-461b-9939-5da80c4940e4",
   "metadata": {
    "height": 30,
    "tags": [],
    "ExecuteTime": {
     "end_time": "2023-12-07T08:19:57.088087Z",
     "start_time": "2023-12-07T08:19:57.044557Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "hello.  foo bar.  cat dog.  mouse\n"
     ]
    }
   ],
   "source": [
    "print(replaced_nodes[1].text)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "72aa4a50-5c0b-4f23-9eae-ec67cb701e29",
   "metadata": {},
   "source": [
    "### Adding a reranker"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "57608e01-ee96-4619-aab1-81d7fce1cbd1",
   "metadata": {
    "height": 132,
    "tags": [],
    "ExecuteTime": {
     "end_time": "2023-12-07T08:20:06.323438Z",
     "start_time": "2023-12-07T08:19:57.050545Z"
    }
   },
   "outputs": [],
   "source": [
    "from llama_index.indices.postprocessor import SentenceTransformerRerank\n",
    "\n",
    "# BAAI/bge-reranker-base\n",
    "# link: https://huggingface.co/BAAI/bge-reranker-base\n",
    "rerank = SentenceTransformerRerank(\n",
    "    top_n=2, model=\"BAAI/bge-reranker-base\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "4d89da07-6b0c-477f-bd0a-2f466c9b159b",
   "metadata": {
    "height": 166,
    "tags": [],
    "ExecuteTime": {
     "end_time": "2023-12-07T08:20:06.336243Z",
     "start_time": "2023-12-07T08:20:06.324211Z"
    }
   },
   "outputs": [],
   "source": [
    "from llama_index import QueryBundle\n",
    "from llama_index.schema import TextNode, NodeWithScore\n",
    "\n",
    "query = QueryBundle(\"I want a dog.\")\n",
    "\n",
    "scored_nodes = [\n",
    "    NodeWithScore(node=TextNode(text=\"This is a cat\"), score=0.6),\n",
    "    NodeWithScore(node=TextNode(text=\"This is a dog\"), score=0.4),\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "1adc6cb4-e741-480a-ab13-06e9194b13b2",
   "metadata": {
    "height": 64,
    "tags": [],
    "ExecuteTime": {
     "end_time": "2023-12-07T08:20:06.853671Z",
     "start_time": "2023-12-07T08:20:06.332428Z"
    }
   },
   "outputs": [],
   "source": [
    "reranked_nodes = rerank.postprocess_nodes(\n",
    "    scored_nodes, query_bundle=query\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "3be841f6-0ab2-4c60-b52f-d27f81fcb1bb",
   "metadata": {
    "height": 30,
    "tags": [],
    "ExecuteTime": {
     "end_time": "2023-12-07T08:20:06.864397Z",
     "start_time": "2023-12-07T08:20:06.854420Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[('This is a dog', 0.91827416), ('This is a cat', 0.0014040753)]\n"
     ]
    }
   ],
   "source": [
    "print([(x.text, x.score) for x in reranked_nodes])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "74d51921-4b0b-4d89-963f-9a9e0ea439d9",
   "metadata": {},
   "source": [
    "### Runing the query engine"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "47bfe5a5-5ce1-4baa-ba61-3b39d16a2337",
   "metadata": {
    "height": 64,
    "tags": [],
    "ExecuteTime": {
     "end_time": "2023-12-07T08:20:06.870082Z",
     "start_time": "2023-12-07T08:20:06.860168Z"
    }
   },
   "outputs": [],
   "source": [
    "sentence_window_engine = sentence_index.as_query_engine(\n",
    "    similarity_top_k=6, node_postprocessors=[postproc, rerank]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "c4f166b9-99f9-4507-af59-3347eaf596c6",
   "metadata": {
    "height": 64,
    "tags": [],
    "ExecuteTime": {
     "end_time": "2023-12-07T08:20:13.751037Z",
     "start_time": "2023-12-07T08:20:06.864239Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...\n",
      "To disable this warning, you can either:\n",
      "\t- Avoid using `tokenizers` before the fork if possible\n",
      "\t- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)\n"
     ]
    }
   ],
   "source": [
    "window_response = sentence_window_engine.query(\n",
    "    \"What are the keys to building a career in AI?\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "ada6304b-a86c-4f6c-a2e9-8ac5b1091a58",
   "metadata": {
    "height": 64,
    "ExecuteTime": {
     "end_time": "2023-12-07T08:20:14.243172Z",
     "start_time": "2023-12-07T08:20:13.754041Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": "<IPython.core.display.Markdown object>",
      "text/markdown": "**`Final Response:`** The keys to building a career in AI include learning foundational technical skills, working on projects, and finding a job. It's also important to be part of a supportive community."
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from llama_index.response.notebook_utils import display_response\n",
    "\n",
    "display_response(window_response)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cbcbe614-befe-4bf7-9813-2c91899650e9",
   "metadata": {},
   "source": [
    "## Putting it all Together"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "313e8b1a-e9e7-4141-a4d2-72fb31a7e057",
   "metadata": {
    "height": 914,
    "ExecuteTime": {
     "end_time": "2023-12-07T08:20:14.260369Z",
     "start_time": "2023-12-07T08:20:14.242271Z"
    }
   },
   "outputs": [],
   "source": [
    "import os\n",
    "from llama_index import ServiceContext, VectorStoreIndex, StorageContext\n",
    "from llama_index.node_parser import SentenceWindowNodeParser\n",
    "from llama_index.indices.postprocessor import MetadataReplacementPostProcessor\n",
    "from llama_index.indices.postprocessor import SentenceTransformerRerank\n",
    "from llama_index import load_index_from_storage\n",
    "\n",
    "\n",
    "def build_sentence_window_index(\n",
    "    documents,\n",
    "    llm,\n",
    "    embed_model=\"local:BAAI/bge-small-en-v1.5\",\n",
    "    sentence_window_size=3,\n",
    "    save_dir=\"sentence_index\",\n",
    "):\n",
    "    # create the sentence window node parser w/ default settings\n",
    "    node_parser = SentenceWindowNodeParser.from_defaults(\n",
    "        window_size=sentence_window_size,\n",
    "        window_metadata_key=\"window\",\n",
    "        original_text_metadata_key=\"original_text\",\n",
    "    )\n",
    "    sentence_context = ServiceContext.from_defaults(\n",
    "        llm=llm,\n",
    "        embed_model=embed_model,\n",
    "        node_parser=node_parser,\n",
    "    )\n",
    "    if not os.path.exists(save_dir):\n",
    "        sentence_index = VectorStoreIndex.from_documents(\n",
    "            documents, service_context=sentence_context\n",
    "        )\n",
    "        sentence_index.storage_context.persist(persist_dir=save_dir)\n",
    "    else:\n",
    "        sentence_index = load_index_from_storage(\n",
    "            StorageContext.from_defaults(persist_dir=save_dir),\n",
    "            service_context=sentence_context,\n",
    "        )\n",
    "\n",
    "    return sentence_index\n",
    "\n",
    "\n",
    "def get_sentence_window_query_engine(\n",
    "    sentence_index, similarity_top_k=6, rerank_top_n=2\n",
    "):\n",
    "    # define postprocessors\n",
    "    postproc = MetadataReplacementPostProcessor(target_metadata_key=\"window\")\n",
    "    rerank = SentenceTransformerRerank(\n",
    "        top_n=rerank_top_n, model=\"BAAI/bge-reranker-base\"\n",
    "    )\n",
    "\n",
    "    sentence_window_engine = sentence_index.as_query_engine(\n",
    "        similarity_top_k=similarity_top_k, node_postprocessors=[postproc, rerank]\n",
    "    )\n",
    "    return sentence_window_engine"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "3b5bb36f-97f9-4e03-bd4e-3ceb635a868b",
   "metadata": {
    "height": 149,
    "ExecuteTime": {
     "end_time": "2023-12-07T08:20:16.083058Z",
     "start_time": "2023-12-07T08:20:14.252644Z"
    }
   },
   "outputs": [],
   "source": [
    "index = build_sentence_window_index(\n",
    "    [document],\n",
    "    llm=llm,\n",
    "    save_dir=\"./sentence_index\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "8d0e78e0-7765-4037-a5fa-63b711dab5aa",
   "metadata": {
    "height": 47,
    "ExecuteTime": {
     "end_time": "2023-12-07T08:20:24.554067Z",
     "start_time": "2023-12-07T08:20:16.078267Z"
    }
   },
   "outputs": [],
   "source": [
    "query_engine = get_sentence_window_query_engine(index, similarity_top_k=6)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eb855430-8fcb-4c25-8d83-a30160449acf",
   "metadata": {},
   "source": [
    "## TruLens Evaluation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "07aa5320-5c1c-4636-aa72-e6a786a58c8a",
   "metadata": {
    "height": 132,
    "ExecuteTime": {
     "end_time": "2023-12-07T08:20:24.556962Z",
     "start_time": "2023-12-07T08:20:24.548669Z"
    }
   },
   "outputs": [],
   "source": [
    "eval_questions = []\n",
    "for path in [\"01_05\", \"06_10\", \"11_15\", \"16_20\", \"21_24\"]:\n",
    "    with open(f'generated_questions_{path}.text', 'r') as file:\n",
    "        for line in file:\n",
    "            # Remove newline character and convert to integer\n",
    "            item = line.strip()\n",
    "            eval_questions.append(item)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "7f9bdeb7-8c37-46b0-80f2-eb8ef1cc7b59",
   "metadata": {
    "height": 132,
    "ExecuteTime": {
     "end_time": "2023-12-07T08:20:24.614617Z",
     "start_time": "2023-12-07T08:20:24.555042Z"
    }
   },
   "outputs": [],
   "source": [
    "from trulens_eval import Tru\n",
    "from tqdm import tqdm\n",
    "\n",
    "def run_evals(eval_questions, tru_recorder, query_engine):\n",
    "    for question in tqdm(eval_questions):\n",
    "        with tru_recorder as recording:\n",
    "            response = query_engine.query(question)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "c7818a9b-04cc-4a6e-98ae-d3d0ab309998",
   "metadata": {
    "height": 98,
    "ExecuteTime": {
     "end_time": "2023-12-07T08:20:24.666849Z",
     "start_time": "2023-12-07T08:20:24.560201Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "🦑 Tru initialized with db url sqlite:///default.sqlite .\n",
      "🛑 Secret keys may be written to the database. See the `database_redact_keys` option of `Tru` to prevent this.\n"
     ]
    }
   ],
   "source": [
    "from utils import get_prebuilt_trulens_recorder\n",
    "\n",
    "from trulens_eval import Tru\n",
    "\n",
    "Tru().reset_database()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "59ef4d27-2c4e-4dd0-85ed-a69a0f8eea00",
   "metadata": {},
   "source": [
    "### Sentence window size = 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "93e6c651-838d-4c36-b806-efd0d8c2d5f0",
   "metadata": {
    "height": 132,
    "ExecuteTime": {
     "end_time": "2023-12-07T08:20:39.291037Z",
     "start_time": "2023-12-07T08:20:24.661593Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "✅ In Answer Relevance, input prompt will be set to __record__.main_input or `Select.RecordInput` .\n",
      "✅ In Answer Relevance, input response will be set to __record__.main_output or `Select.RecordOutput` .\n",
      "✅ In Context Relevance, input prompt will be set to __record__.main_input or `Select.RecordInput` .\n",
      "✅ In Context Relevance, input response will be set to __record__.app.query.rets.source_nodes[:].node.text .\n",
      "✅ In Groundedness, input source will be set to __record__.app.query.rets.source_nodes[:].node.text .\n",
      "✅ In Groundedness, input statement will be set to __record__.main_output or `Select.RecordOutput` .\n"
     ]
    }
   ],
   "source": [
    "sentence_index_1 = build_sentence_window_index(\n",
    "    documents,\n",
    "    llm=llm,\n",
    "    embed_model=embed_model,\n",
    "    sentence_window_size=1,\n",
    "    save_dir=\"sentence_index_1\",\n",
    ")\n",
    "\n",
    "sentence_window_engine_1 = get_sentence_window_query_engine(\n",
    "    sentence_index_1\n",
    ")\n",
    "\n",
    "tru_recorder_1 = get_prebuilt_trulens_recorder(\n",
    "    sentence_window_engine_1,\n",
    "    app_id='sentence window engine 1'\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "ca942b16-52be-4107-861c-33daeca1f619",
   "metadata": {
    "height": 30,
    "ExecuteTime": {
     "end_time": "2023-12-07T08:24:59.550997Z",
     "start_time": "2023-12-07T08:20:39.291494Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 0/24 [00:00<?, ?it/s]huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...\n",
      "To disable this warning, you can either:\n",
      "\t- Avoid using `tokenizers` before the fork if possible\n",
      "\t- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "boto3,botocore is/are required for using BedrockEndpoint. You should be able to install it/them with\n",
      "\tpip install boto3 botocore\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...\n",
      "To disable this warning, you can either:\n",
      "\t- Avoid using `tokenizers` before the fork if possible\n",
      "\t- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)\n",
      "huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...\n",
      "To disable this warning, you can either:\n",
      "\t- Avoid using `tokenizers` before the fork if possible\n",
      "\t- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)\n",
      "100%|██████████| 24/24 [04:20<00:00, 10.84s/it]\n"
     ]
    }
   ],
   "source": [
    "run_evals(eval_questions, tru_recorder_1, sentence_window_engine_1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Starting dashboard ...\n",
      "Config file already exists. Skipping writing process.\n",
      "Credentials file already exists. Skipping writing process.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...\n",
      "To disable this warning, you can either:\n",
      "\t- Avoid using `tokenizers` before the fork if possible\n",
      "\t- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)\n"
     ]
    },
    {
     "data": {
      "text/plain": "Accordion(children=(VBox(children=(VBox(children=(Label(value='STDOUT'), Output())), VBox(children=(Label(valu…",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "4549b1f764ea4c8ab36a130ea9d48c74"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Dashboard started at http://10.105.0.130:8501 .\n"
     ]
    },
    {
     "data": {
      "text/plain": "<Popen: returncode: None args: ['streamlit', 'run', '--server.headless=True'...>"
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Tru().run_dashboard()"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2023-12-07T08:25:04.476106Z",
     "start_time": "2023-12-07T08:24:59.556955Z"
    }
   },
   "id": "e91618f0ac6c9c15"
  },
  {
   "cell_type": "markdown",
   "id": "e54c9977-c606-41bb-b87a-ec3cc17eabaf",
   "metadata": {},
   "source": [
    "### Sentence window size = 3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0c8bea9d-9805-4302-b8c0-6e0f3eeec030",
   "metadata": {
    "height": 285,
    "ExecuteTime": {
     "start_time": "2023-12-06T16:32:11.058553Z"
    }
   },
   "outputs": [],
   "source": [
    "sentence_index_3 = build_sentence_window_index(\n",
    "    documents,\n",
    "    llm=llm,\n",
    "    embed_model=\"local:BAAI/bge-small-en-v1.5\",\n",
    "    sentence_window_size=3,\n",
    "    save_dir=\"sentence_index_3\",\n",
    ")\n",
    "sentence_window_engine_3 = get_sentence_window_query_engine(\n",
    "    sentence_index_3\n",
    ")\n",
    "\n",
    "tru_recorder_3 = get_prebuilt_trulens_recorder(\n",
    "    sentence_window_engine_3,\n",
    "    app_id='sentence window engine 3'\n",
    ")\n",
    "run_evals(eval_questions, tru_recorder_3, sentence_window_engine_3)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c5113f41",
   "metadata": {
    "height": 30
   },
   "source": [
    "### Sentence window size = 5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1223e7e9",
   "metadata": {
    "height": 285,
    "ExecuteTime": {
     "end_time": "2023-12-06T16:32:11.067672Z",
     "start_time": "2023-12-06T16:32:11.065363Z"
    }
   },
   "outputs": [],
   "source": [
    "sentence_index_5 = build_sentence_window_index(\n",
    "    documents,\n",
    "    llm=llm,\n",
    "    embed_model=\"local:BAAI/bge-small-en-v1.5\",\n",
    "    sentence_window_size=5,\n",
    "    save_dir=\"sentence_index_5\",\n",
    ")\n",
    "sentence_window_engine_5 = get_sentence_window_query_engine(\n",
    "    sentence_index_5\n",
    ")\n",
    "\n",
    "tru_recorder_5 = get_prebuilt_trulens_recorder(\n",
    "    sentence_window_engine_5,\n",
    "    app_id='sentence window engine 5'\n",
    ")\n",
    "run_evals(eval_questions, tru_recorder_5, sentence_window_engine_5)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8f7a893b",
   "metadata": {
    "height": 30
   },
   "source": [
    "### Sentence window size = 8"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f5c9961a",
   "metadata": {
    "height": 285,
    "ExecuteTime": {
     "end_time": "2023-12-06T16:32:11.077429Z",
     "start_time": "2023-12-06T16:32:11.069600Z"
    }
   },
   "outputs": [],
   "source": [
    "sentence_index_8 = build_sentence_window_index(\n",
    "    documents,\n",
    "    llm=llm,\n",
    "    embed_model=\"local:BAAI/bge-small-en-v1.5\",\n",
    "    sentence_window_size=8,\n",
    "    save_dir=\"sentence_index_8\",\n",
    ")\n",
    "sentence_window_engine_8 = get_sentence_window_query_engine(\n",
    "    sentence_index_8\n",
    ")\n",
    "\n",
    "tru_recorder_8 = get_prebuilt_trulens_recorder(\n",
    "    sentence_window_engine_8,\n",
    "    app_id='sentence window engine 8'\n",
    ")\n",
    "run_evals(eval_questions, tru_recorder_8, sentence_window_engine_8)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b483fd08",
   "metadata": {
    "height": 30
   },
   "source": [
    "### Show the dashboard"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bb1373b9",
   "metadata": {
    "height": 30,
    "ExecuteTime": {
     "start_time": "2023-12-06T16:32:11.072272Z"
    }
   },
   "outputs": [],
   "source": [
    "Tru().run_dashboard()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "98b52367",
   "metadata": {
    "height": 30
   },
   "source": [
    "### Note about the dataset of questions\n",
    "- For evaluating a personal project, an eval set of 20 is reasonable.\n",
    "- For evaluating business applications, you may need a set of 100+ in order to cover all the use cases thoroughly.\n",
    "- Note that since API calls can sometimes fail, you may occasionally see null responses, and would want to re-run your evaluations.  So running your evaluations in smaller batches can also help you save time and cost by only re-running the evaluation on the batches with issues."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [],
   "metadata": {
    "collapsed": false
   },
   "id": "4cc93fea6ab47550"
  }
 ],
 "metadata": {
  "kernelspec": {
   "name": "prometeia",
   "language": "python",
   "display_name": "prometeia"
  },
  "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
}