generating_qa_dataset.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install -r requirements.txt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from huggingface_hub import login\n",
    "\n",
    "login(\"api_key_here\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core import Settings, SimpleDirectoryReader, VectorStoreIndex, set_global_tokenizer\n",
    "\n",
    "from llama_index.embeddings.huggingface import HuggingFaceEmbedding\n",
    "from llama_index.core.base.llms.types import CompletionResponse\n",
    "from llama_index.llms.ollama import Ollama as LlamaIndexOllama\n",
    "from llama_index.core.retrievers import QueryFusionRetriever\n",
    "from transformers import AutoTokenizer, BitsAndBytesConfig\n",
    "from llama_index.core.node_parser import SentenceSplitter\n",
    "from llama_index.embeddings.ollama import OllamaEmbedding\n",
    "from llama_index.core.llms import CustomLLM, LLMMetadata\n",
    "from llama_index.llms.huggingface import HuggingFaceLLM\n",
    "from llama_index.retrievers.bm25 import BM25Retriever\n",
    "from llama_index.core.evaluation import (\n",
    "    EmbeddingQAFinetuneDataset,\n",
    "    RetrieverEvaluator,\n",
    "    generate_question_context_pairs\n",
    ")\n",
    "import nest_asyncio\n",
    "import warnings\n",
    "import argparse\n",
    "import torch\n",
    "\n",
    "nest_asyncio.apply()\n",
    "\n",
    "warnings.filterwarnings('ignore')\n",
    "\n",
    "\n",
    "\n",
    "system_prompt = \"You are a Q&A assistant. Your goal is to answer questions as accurately as possible based on the instructions and context provided. Only use the context provided and STRICTLY say you dont know if you dont know.\"\n",
    "query_wrapper_prompt = \"<|USER|>{query_str}<|ASSISTANT|>\"\n",
    "\n",
    "bnb_config = BitsAndBytesConfig(\n",
    "load_in_4bit=True,\n",
    "bnb_4bit_use_double_quant=True,\n",
    "bnb_4bit_quant_type=\"nf4\",\n",
    "bnb_4bit_compute_dtype=torch.bfloat16\n",
    ")\n",
    "\n",
    "llm_quest = HuggingFaceLLM(\n",
    "        model_name=\"Qwen/Qwen2.5-1.5B-Instruct\",\n",
    "        tokenizer_name=\"Qwen/Qwen2.5-1.5B-Instruct\",\n",
    "        system_prompt=system_prompt,\n",
    "        query_wrapper_prompt=query_wrapper_prompt,\n",
    "        context_window=3900,\n",
    "        max_new_tokens=256,\n",
    "        model_kwargs={\"quantization_config\": bnb_config},\n",
    "        generate_kwargs={\"temperature\": 0.1},\n",
    "        device_map=\"auto\",\n",
    "    )\n",
    "\n",
    "set_global_tokenizer(\n",
    "    AutoTokenizer.from_pretrained(\"Qwen/Qwen2.5-1.5B-Instruct\").encode\n",
    ")\n",
    "\n",
    "embed_model = HuggingFaceEmbedding(model_name = \"BAAI/bge-small-en-v1.5\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def parse_args():\n",
    "    parser = argparse.ArgumentParser(description=\"Tune RAG Model MLflow\")\n",
    "    parser.add_argument('--dataset_dir', type=str, default=\"./data\", help=\"Directory for the dataset\")\n",
    "    parser.add_argument('--chunk_size', type=int, default=512, help=\"Chunk size for splitting documents\")\n",
    "    parser.add_argument('--top_k', type=int, default=5, help=\"Top K similar nodes to retrieve\")\n",
    "    parser.add_argument('--model_name', type=str, default='llama3.2:1b', help=\"Model name\")\n",
    "    parser.add_argument('--embedder_name', type=str, default='nomic-embed-text:latest', help=\"Embedder name\")\n",
    "    parser.add_argument('--dataset_name', type=str, default='pg_eval_dataset_index.json', help=\"Dataset name\")\n",
    "    parser.add_argument('--chunk_questions', type=int, default=2, help=\"Number of questions per chunk\")\n",
    "    return parser.parse_args()\n",
    "\n",
    "async def tune_rag(args):\n",
    "    documents = SimpleDirectoryReader(args.dataset_dir).load_data()\n",
    "    node_parser = SentenceSplitter(chunk_size=args.chunk_size, chunk_overlap=100)\n",
    "    nodes = node_parser.get_nodes_from_documents(documents)\n",
    "\n",
    "    for idx, node in enumerate(nodes):\n",
    "        node.id_ = f\"node_{idx}\"\n",
    "\n",
    "    vector_index = VectorStoreIndex(nodes)\n",
    "\n",
    "    retriever = vector_index.as_retriever(similarity_top_k=args.top_k, similarity_cutoff=0.7)\n",
    "\n",
    "    # Generate new dataset\n",
    "    qa_dataset = generate_question_context_pairs(\n",
    "        nodes, \n",
    "        llm=llm_quest, \n",
    "        num_questions_per_chunk=args.chunk_questions\n",
    "    )\n",
    "    print(f\"this is the qa_dataset outside the function: {qa_dataset}\")\n",
    "    \n",
    "    # Save the newly generated dataset\n",
    "    qa_dataset.save_json(args.dataset_name)\n",
    "\n",
    "    print(f\"Generated and saved new dataset to {args.dataset_name}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Manually define args in Colab\n",
    "from types import SimpleNamespace\n",
    "\n",
    "Settings.embed_model = embed_model\n",
    "args = SimpleNamespace(\n",
    "    dataset_dir=\"./data\",\n",
    "    chunk_size=512,\n",
    "    top_k=5,\n",
    "    model_name=\"Qwen/Qwen2.5-1.5B-Instruct\",\n",
    "    embedder_name=\"nomic-embed-text:latest\",\n",
    "    dataset_name=\"pg_eval_dataset_index_BERT.json\",\n",
    "    chunk_questions=2,\n",
    "\n",
    ")\n",
    "await tune_rag(args)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "rag_afdel",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.10.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}