baseline.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The autoreload extension is already loaded. To reload it, use:\n",
      "  %reload_ext autoreload\n"
     ]
    }
   ],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "\n",
    "from rag_components import *\n",
    "from llama_index.core.retrievers import VectorIndexRetriever\n",
    "\n",
    "import re"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "def answer_reading_comprehension_baseline(\n",
    "    question, context_title, context_text=None, top_k=2, chunk_size=1024, chunk_overlap=200, qa_llm=gpt4\n",
    "):\n",
    "    \"\"\"\n",
    "    Answer a question given a context, using question generation.\n",
    "\n",
    "    Args:\n",
    "        question (str): The question to answer\n",
    "        context_title (str): The title of the context\n",
    "        context_text (str): The text of the context\n",
    "        top_k (int): The number of top matching chunks to retrieve\n",
    "\n",
    "    Returns:\n",
    "        str: The answer to the question\n",
    "    \"\"\"\n",
    "    raw_index = get_index_by_title(context_title)\n",
    "    combined_index = index_union([raw_index])\n",
    "    retriever = VectorIndexRetriever(\n",
    "        index=combined_index,\n",
    "        similarity_top_k=top_k * 32,\n",
    "    )\n",
    "    top_nodes = retriever.retrieve(question)\n",
    "\n",
    "    def text_chunk_id(feature_id):\n",
    "        if \"text_chunk_\" in feature_id:\n",
    "            return feature_id\n",
    "        elif \"summary_\" in feature_id:\n",
    "            return feature_id.replace(\"summary_\", \"text_chunk_\")\n",
    "        elif \"question_\" in feature_id:\n",
    "            return \"_\".join(feature_id.replace(\"question_\", \"text_chunk_\").split(\"_\")[:-1])\n",
    "\n",
    "    unique_top_k_raw_text_ids = []\n",
    "    unique_top_k_nodes = []\n",
    "    for node in top_nodes:\n",
    "        raw_text_id = text_chunk_id(node.node.id_)\n",
    "        if raw_text_id not in unique_top_k_raw_text_ids:\n",
    "            unique_top_k_raw_text_ids.append(raw_text_id)\n",
    "            unique_top_k_nodes.append(node)\n",
    "            if len(unique_top_k_nodes) == top_k:\n",
    "                break\n",
    "\n",
    "    nodes_to_include = [get_node_by_id(raw_index, raw_text_id) for raw_text_id in unique_top_k_raw_text_ids]\n",
    "    texts_to_include = [node.text for node in nodes_to_include]\n",
    "    texts_to_include_combined = \" \".join(texts_to_include)\n",
    "\n",
    "    response = answer_reading_comprehension(question, texts_to_include_combined, qa_llm=qa_llm)\n",
    "\n",
    "    top_chunks_info = [\n",
    "        {\n",
    "            \"feature_type\": (\n",
    "                \"summary\"\n",
    "                if \"summary_\" in unique_top_k_nodes[i].node.id_\n",
    "                else \"question\" if \"question_\" in unique_top_k_nodes[i].node.id_ else \"text\"\n",
    "            ),\n",
    "            \"feature_score\": unique_top_k_nodes[i].score,\n",
    "            \"feature_rank\": i,\n",
    "            \"text_score\": similarity_score(question, context_title, text_chunk_id(node.id_)),\n",
    "            \"feature\": unique_top_k_nodes[i].node.text,\n",
    "            \"text\": nodes_to_include[i].text,\n",
    "        }\n",
    "        for i in range(len(nodes_to_include))\n",
    "    ]\n",
    "\n",
    "    additional_info = {\n",
    "        \"top_chunks_info\": top_chunks_info,\n",
    "    }\n",
    "    return response, additional_info"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Generating output for chunk size 256\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "022c6ad92a834594b942e4ffed8b6521",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Answering questions:   0%|          | 0/100 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "debug_lim = 100\n",
    "chunk_sizes = [64, 128, 256, 512, 1024, 2048]\n",
    "topks = [32, 16, 8, 4, 2, 1]\n",
    "chunk_overlaps = [10, 25, 50, 100, 200, 400]\n",
    "for idx in [2]:\n",
    "    print(\"Generating output for chunk size\", chunk_sizes[idx])\n",
    "    output_file = f\"output/testset_baseline_mistral_large_chunksize{chunk_sizes[idx]}.jsonl\"\n",
    "    test_longdep_qa(\n",
    "        answer_reading_comprehension_baseline,\n",
    "        output_file=output_file,\n",
    "        debug_start=105,\n",
    "        debug_lim=debug_lim,\n",
    "        qa_llm=mistral_large,\n",
    "        chunk_size=chunk_sizes[idx],\n",
    "        top_k=topks[idx],\n",
    "        chunk_overlap=chunk_overlaps[idx],\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "results for chunk size 256\n",
      "Rouge Metrics: {'rouge1': 0.19060201063602866, 'rouge2': 0.07962804700229226, 'rougeL': 0.16280161569167925, 'rougeLsum': 0.16693164068469166}\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3ba23b3a426f49da97c78b0cca508c07",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/100 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "LLM Self-Score: 0.48\n"
     ]
    }
   ],
   "source": [
    "for idx in [2]:\n",
    "    output_file = f\"output/testset_baseline_mistral_large_chunksize{chunk_sizes[idx]}.jsonl\"\n",
    "    rouge_metrics = get_rouge_metrics(output_file)\n",
    "    print(\"results for chunk size\", chunk_sizes[idx])\n",
    "    print(\"Rouge Metrics:\", rouge_metrics)\n",
    "\n",
    "    self_score = llm_self_score(output_file, llm=gpt4)\n",
    "    print(\"LLM Self-Score:\", self_score)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "LongRAG",
   "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.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}