add generate_qa_dataset notebook

scripts/generate_qa_dataset.ipynb

@@ -0,0 +1,287 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import nest_asyncio\n",
+    "\n",
+    "nest_asyncio.apply()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Generate synthetic dataset of questions + doc ids\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# DEFAULT_QA_GENERATE_PROMPT_TMPL = \"\"\"\\\n",
+    "# Context information is below.\n",
+    "\n",
+    "# ---------------------\n",
+    "# {context_str}\n",
+    "# ---------------------\n",
+    "\n",
+    "# Given the context information and not prior knowledge.\n",
+    "# generate only questions based on the below query.\n",
+    "\n",
+    "# You are a Teacher/ Professor with expertise in the field of AI. Your task is to setup \\\n",
+    "# {num_questions_per_chunk} questions for an upcoming \\\n",
+    "# quiz/examination. The questions should be diverse in nature \\\n",
+    "# across the document. Restrict the questions to the \\\n",
+    "# context information provided.\"\n",
+    "# \"\"\"\n",
+    "\n",
+    "STUDENT_AI_QUESTIONS_GENERATE_PROMPT_TMPL = \"\"\"\\\n",
+    "Context information is below.\n",
+    "\n",
+    "---------------------\n",
+    "{context_str}\n",
+    "---------------------\n",
+    "\n",
+    "Given the context information above, generate {num_questions_per_chunk} questions that a student might ask about AI, specifically related to the information provided in the context, \n",
+    "but without the questions mentioning the context information.\n",
+    "\n",
+    "You are simulating curious students who are learning about AI. Your task is to create questions that:\n",
+    "1. Reflect genuine curiosity about the topics covered in the context.\n",
+    "2. Vary in complexity, from basic clarifications to more advanced inquiries.\n",
+    "3. Demonstrate a student's desire to understand AI concepts better.\n",
+    "4. Can be answered using the information provided in the context.\n",
+    "\n",
+    "The questions should be diverse and cover different aspects of the content. Do not use any prior knowledge beyond what's given in the context. Ensure that each question you generate can be answered using the information provided in the context.\n",
+    "\"\"\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from llama_index.llms.openai import OpenAI\n",
+    "from llama_index.llms.gemini import Gemini\n",
+    "\n",
+    "# llm = Gemini(model=\"models/gemini-1.5-flash-latest\", temperature=1, max_tokens=1000)\n",
+    "llm = OpenAI(\n",
+    "    api_key=\"\",\n",
+    "    model=\"gpt-4o-mini\",\n",
+    "    temperature=1,\n",
+    "    max_tokens=1000,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import json\n",
+    "import re\n",
+    "import uuid\n",
+    "import warnings\n",
+    "import asyncio\n",
+    "from typing import Dict, List, Tuple\n",
+    "\n",
+    "from llama_index.core.llms.llm import LLM\n",
+    "from llama_index.core.schema import MetadataMode, TextNode\n",
+    "from tqdm.asyncio import tqdm as async_tqdm\n",
+    "from llama_index.core.llama_dataset.legacy.embedding import EmbeddingQAFinetuneDataset\n",
+    "from llama_index.core import Document\n",
+    "\n",
+    "\n",
+    "async def generate_qa_embedding_pairs(\n",
+    "    nodes: List[TextNode],\n",
+    "    llm: LLM,\n",
+    "    qa_generate_prompt_tmpl: str = STUDENT_AI_QUESTIONS_GENERATE_PROMPT_TMPL,\n",
+    "    num_questions_per_chunk: int = 1,\n",
+    "    max_concurrent: int = 10,\n",
+    "    delay: float = 0.5,\n",
+    ") -> EmbeddingQAFinetuneDataset:\n",
+    "    \"\"\"Generate examples given a set of nodes.\"\"\"\n",
+    "    node_dict = {\n",
+    "        node.node_id: node.get_content(metadata_mode=MetadataMode.NONE)\n",
+    "        for node in nodes\n",
+    "    }\n",
+    "\n",
+    "    queries = {}\n",
+    "    relevant_docs = {}\n",
+    "    semaphore = asyncio.Semaphore(max_concurrent)\n",
+    "\n",
+    "    async def process_node(node_id: str, text: str):\n",
+    "        async with semaphore:\n",
+    "            query = qa_generate_prompt_tmpl.format(\n",
+    "                context_str=text, num_questions_per_chunk=num_questions_per_chunk\n",
+    "            )\n",
+    "            response = await llm.acomplete(\n",
+    "                query\n",
+    "            )  # Assuming the LLM has an async method\n",
+    "\n",
+    "            result = str(response).strip().split(\"\\n\")\n",
+    "            questions = [\n",
+    "                re.sub(r\"^\\d+[\\).\\s]\", \"\", question).strip() for question in result\n",
+    "            ]\n",
+    "            questions = [question for question in questions if len(question) > 0][\n",
+    "                :num_questions_per_chunk\n",
+    "            ]\n",
+    "\n",
+    "            num_questions_generated = len(questions)\n",
+    "            if num_questions_generated < num_questions_per_chunk:\n",
+    "                warnings.warn(\n",
+    "                    f\"Fewer questions generated ({num_questions_generated}) \"\n",
+    "                    f\"than requested ({num_questions_per_chunk}).\"\n",
+    "                )\n",
+    "\n",
+    "            for question in questions:\n",
+    "                question_id = str(uuid.uuid4())\n",
+    "                queries[question_id] = question\n",
+    "                relevant_docs[question_id] = [node_id]\n",
+    "\n",
+    "            await asyncio.sleep(delay)\n",
+    "\n",
+    "    # Use asyncio.gather to process nodes concurrently\n",
+    "    await async_tqdm.gather(\n",
+    "        *[process_node(node_id, text) for node_id, text in node_dict.items()]\n",
+    "    )\n",
+    "\n",
+    "    # construct dataset\n",
+    "    return EmbeddingQAFinetuneDataset(\n",
+    "        queries=queries, corpus=node_dict, relevant_docs=relevant_docs\n",
+    "    )"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import pickle\n",
+    "\n",
+    "\n",
+    "async def generate_questions(path, source_name):\n",
+    "    nodes = []\n",
+    "    with open(path, \"rb\") as f:\n",
+    "        document_dict = pickle.load(f)\n",
+    "        for doc_id in document_dict.keys():\n",
+    "            doc: Document = document_dict[doc_id]\n",
+    "            if doc.metadata[\"tokens\"] >= 100_000:\n",
+    "                print(\"skipping\", doc.metadata[\"tokens\"])\n",
+    "                continue\n",
+    "            node = TextNode(text=doc.text, metadata=doc.metadata, id_=doc_id)\n",
+    "            nodes.append(node)\n",
+    "\n",
+    "    rag_eval_dataset: EmbeddingQAFinetuneDataset = await generate_qa_embedding_pairs(\n",
+    "        nodes,\n",
+    "        llm=llm,\n",
+    "        num_questions_per_chunk=1,\n",
+    "        qa_generate_prompt_tmpl=STUDENT_AI_QUESTIONS_GENERATE_PROMPT_TMPL,\n",
+    "        max_concurrent=20,  # Adjust this to control concurrency\n",
+    "        delay=0.5,  # Adjust this to add delay between API calls\n",
+    "    )\n",
+    "    rag_eval_dataset.save_json(f\"./rag_eval_{source_name}.json\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# await generate_questions(\n",
+    "#     \"../data/chroma-db-langchain/document_dict_langchain.pkl\", \"langchain\"\n",
+    "# )\n",
+    "# await generate_questions(\n",
+    "#     \"../data/chroma-db-llama_index/document_dict_llama_index.pkl\", \"llama_index\"\n",
+    "# )\n",
+    "# await generate_questions(\n",
+    "#     \"../data/chroma-db-openai_cookbooks/document_dict_openai_cookbooks.pkl\",\n",
+    "#     \"openai_cookbooks\",\n",
+    "# )\n",
+    "# await generate_questions(\"../data/chroma-db-peft/document_dict_peft.pkl\", \"peft\")\n",
+    "# await generate_questions(\"../data/chroma-db-trl/document_dict_trl.pkl\", \"trl\")\n",
+    "\n",
+    "await generate_questions(\n",
+    "    \"../data/chroma-db-tai_blog/document_dict_tai_blog.pkl\", \"tai_blog\"\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# # We can also load the dataset from a previously saved json file.\n",
+    "# from llama_index.core.evaluation import EmbeddingQAFinetuneDataset\n",
+    "\n",
+    "# rag_eval_dataset = EmbeddingQAFinetuneDataset.from_json(\"./rag_eval_transformers.json\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "env",
+   "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.12.5"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}