LlamaIndex_Yi-34B-RAG.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Yi 34B Chat\n",
    "#### RAG with LlamaIndex - Nvidia CUDA + WSL (Windows Subsystem for Linux) + Word documents + Local LLM\n",
    "\n",
    "This notebook demonstrates the use of LlamaIndex for Retrieval Augmented Generation using Windows' WSL and an Nvidia's CUDA.\n",
    "\n",
    "See the [README.md](README.md) file for help on how to run this."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 1. Prepare Llama Index for use"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/markwsl/miniconda3/envs/LlamaIndexRAGLinux/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "import logging\n",
    "import sys\n",
    "\n",
    "logging.basicConfig(stream=sys.stdout, level=logging.INFO)\n",
    "logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))\n",
    "\n",
    "from llama_index.core import VectorStoreIndex, SimpleDirectoryReader, ServiceContext"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 2. Load the Word document(s)\n",
    "\n",
    "Note: A fictitious story about Thundertooth a dinosaur who has travelled to the future. Thanks ChatGPT!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "documents = SimpleDirectoryReader(\"./Data/\").load_data()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 3. Instantiate the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "ggml_init_cublas: GGML_CUDA_FORCE_MMQ:   no\n",
      "ggml_init_cublas: CUDA_USE_TENSOR_CORES: yes\n",
      "ggml_init_cublas: found 1 CUDA devices:\n",
      "  Device 0: NVIDIA GeForce RTX 3090, compute capability 8.6, VMM: yes\n",
      "llama_model_loader: loaded meta data with 23 key-value pairs and 543 tensors from ./Models/yi-34b-chat.Q4_K_M.gguf (version GGUF V3 (latest))\n",
      "llama_model_loader: Dumping metadata keys/values. Note: KV overrides do not apply in this output.\n",
      "llama_model_loader: - kv   0:                       general.architecture str              = llama\n",
      "llama_model_loader: - kv   1:                               general.name str              = LLaMA v2\n",
      "llama_model_loader: - kv   2:                       llama.context_length u32              = 4096\n",
      "llama_model_loader: - kv   3:                     llama.embedding_length u32              = 7168\n",
      "llama_model_loader: - kv   4:                          llama.block_count u32              = 60\n",
      "llama_model_loader: - kv   5:                  llama.feed_forward_length u32              = 20480\n",
      "llama_model_loader: - kv   6:                 llama.rope.dimension_count u32              = 128\n",
      "llama_model_loader: - kv   7:                 llama.attention.head_count u32              = 56\n",
      "llama_model_loader: - kv   8:              llama.attention.head_count_kv u32              = 8\n",
      "llama_model_loader: - kv   9:     llama.attention.layer_norm_rms_epsilon f32              = 0.000010\n",
      "llama_model_loader: - kv  10:                       llama.rope.freq_base f32              = 5000000.000000\n",
      "llama_model_loader: - kv  11:                          general.file_type u32              = 15\n",
      "llama_model_loader: - kv  12:                       tokenizer.ggml.model str              = llama\n",
      "llama_model_loader: - kv  13:                      tokenizer.ggml.tokens arr[str,64000]   = [\"<unk>\", \"<|startoftext|>\", \"<|endof...\n",
      "llama_model_loader: - kv  14:                      tokenizer.ggml.scores arr[f32,64000]   = [0.000000, 0.000000, 0.000000, 0.0000...\n",
      "llama_model_loader: - kv  15:                  tokenizer.ggml.token_type arr[i32,64000]   = [2, 3, 3, 3, 3, 3, 1, 1, 1, 3, 3, 3, ...\n",
      "llama_model_loader: - kv  16:                tokenizer.ggml.bos_token_id u32              = 1\n",
      "llama_model_loader: - kv  17:                tokenizer.ggml.eos_token_id u32              = 2\n",
      "llama_model_loader: - kv  18:            tokenizer.ggml.padding_token_id u32              = 0\n",
      "llama_model_loader: - kv  19:               tokenizer.ggml.add_bos_token bool             = false\n",
      "llama_model_loader: - kv  20:               tokenizer.ggml.add_eos_token bool             = false\n",
      "llama_model_loader: - kv  21:                    tokenizer.chat_template str              = {% if not add_generation_prompt is de...\n",
      "llama_model_loader: - kv  22:               general.quantization_version u32              = 2\n",
      "llama_model_loader: - type  f32:  121 tensors\n",
      "llama_model_loader: - type q4_K:  361 tensors\n",
      "llama_model_loader: - type q6_K:   61 tensors\n",
      "llm_load_vocab: mismatch in special tokens definition ( 498/64000 vs 267/64000 ).\n",
      "llm_load_print_meta: format           = GGUF V3 (latest)\n",
      "llm_load_print_meta: arch             = llama\n",
      "llm_load_print_meta: vocab type       = SPM\n",
      "llm_load_print_meta: n_vocab          = 64000\n",
      "llm_load_print_meta: n_merges         = 0\n",
      "llm_load_print_meta: n_ctx_train      = 4096\n",
      "llm_load_print_meta: n_embd           = 7168\n",
      "llm_load_print_meta: n_head           = 56\n",
      "llm_load_print_meta: n_head_kv        = 8\n",
      "llm_load_print_meta: n_layer          = 60\n",
      "llm_load_print_meta: n_rot            = 128\n",
      "llm_load_print_meta: n_embd_head_k    = 128\n",
      "llm_load_print_meta: n_embd_head_v    = 128\n",
      "llm_load_print_meta: n_gqa            = 7\n",
      "llm_load_print_meta: n_embd_k_gqa     = 1024\n",
      "llm_load_print_meta: n_embd_v_gqa     = 1024\n",
      "llm_load_print_meta: f_norm_eps       = 0.0e+00\n",
      "llm_load_print_meta: f_norm_rms_eps   = 1.0e-05\n",
      "llm_load_print_meta: f_clamp_kqv      = 0.0e+00\n",
      "llm_load_print_meta: f_max_alibi_bias = 0.0e+00\n",
      "llm_load_print_meta: n_ff             = 20480\n",
      "llm_load_print_meta: n_expert         = 0\n",
      "llm_load_print_meta: n_expert_used    = 0\n",
      "llm_load_print_meta: rope scaling     = linear\n",
      "llm_load_print_meta: freq_base_train  = 5000000.0\n",
      "llm_load_print_meta: freq_scale_train = 1\n",
      "llm_load_print_meta: n_yarn_orig_ctx  = 4096\n",
      "llm_load_print_meta: rope_finetuned   = unknown\n",
      "llm_load_print_meta: model type       = 30B\n",
      "llm_load_print_meta: model ftype      = Q4_K - Medium\n",
      "llm_load_print_meta: model params     = 34.39 B\n",
      "llm_load_print_meta: model size       = 19.24 GiB (4.81 BPW) \n",
      "llm_load_print_meta: general.name     = LLaMA v2\n",
      "llm_load_print_meta: BOS token        = 1 '<|startoftext|>'\n",
      "llm_load_print_meta: EOS token        = 2 '<|endoftext|>'\n",
      "llm_load_print_meta: UNK token        = 0 '<unk>'\n",
      "llm_load_print_meta: PAD token        = 0 '<unk>'\n",
      "llm_load_print_meta: LF token         = 315 '<0x0A>'\n",
      "llm_load_tensors: ggml ctx size =    0.42 MiB\n",
      "llm_load_tensors: offloading 30 repeating layers to GPU\n",
      "llm_load_tensors: offloaded 30/61 layers to GPU\n",
      "llm_load_tensors:        CPU buffer size = 19700.24 MiB\n",
      "llm_load_tensors:      CUDA0 buffer size =  9585.52 MiB\n",
      "...................................................................................................\n",
      "llama_new_context_with_model: n_ctx      = 32768\n",
      "llama_new_context_with_model: freq_base  = 5000000.0\n",
      "llama_new_context_with_model: freq_scale = 1\n",
      "llama_kv_cache_init:  CUDA_Host KV buffer size =  3840.00 MiB\n",
      "llama_kv_cache_init:      CUDA0 KV buffer size =  3840.00 MiB\n",
      "llama_new_context_with_model: KV self size  = 7680.00 MiB, K (f16): 3840.00 MiB, V (f16): 3840.00 MiB\n",
      "llama_new_context_with_model:  CUDA_Host input buffer size   =    79.26 MiB\n",
      "llama_new_context_with_model:      CUDA0 compute buffer size =  3704.13 MiB\n",
      "llama_new_context_with_model:  CUDA_Host compute buffer size =  3654.00 MiB\n",
      "llama_new_context_with_model: graph splits (measure): 5\n",
      "AVX = 1 | AVX_VNNI = 0 | AVX2 = 1 | AVX512 = 0 | AVX512_VBMI = 0 | AVX512_VNNI = 0 | FMA = 1 | NEON = 0 | ARM_FMA = 0 | F16C = 1 | FP16_VA = 0 | WASM_SIMD = 0 | BLAS = 1 | SSE3 = 1 | SSSE3 = 1 | VSX = 0 | MATMUL_INT8 = 0 | \n",
      "Model metadata: {'tokenizer.chat_template': \"{% if not add_generation_prompt is defined %}{% set add_generation_prompt = false %}{% endif %}{% for message in messages %}{{'<|im_start|>' + message['role'] + '\\n' + message['content'] + '<|im_end|>' + '\\n'}}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant\\n' }}{% endif %}\", 'tokenizer.ggml.add_eos_token': 'false', 'tokenizer.ggml.padding_token_id': '0', 'tokenizer.ggml.eos_token_id': '2', 'general.architecture': 'llama', 'llama.rope.freq_base': '5000000.000000', 'llama.context_length': '4096', 'general.name': 'LLaMA v2', 'tokenizer.ggml.add_bos_token': 'false', 'llama.embedding_length': '7168', 'llama.feed_forward_length': '20480', 'llama.attention.layer_norm_rms_epsilon': '0.000010', 'llama.rope.dimension_count': '128', 'tokenizer.ggml.bos_token_id': '1', 'llama.attention.head_count': '56', 'llama.block_count': '60', 'llama.attention.head_count_kv': '8', 'general.quantization_version': '2', 'tokenizer.ggml.model': 'llama', 'general.file_type': '15'}\n",
      "Using chat template: {% if not add_generation_prompt is defined %}{% set add_generation_prompt = false %}{% endif %}{% for message in messages %}{{'<|im_start|>' + message['role'] + '\n",
      "' + message['content'] + '<|im_end|>' + '\n",
      "'}}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant\n",
      "' }}{% endif %}\n",
      "Using chat eos_token: \n",
      "Using chat bos_token: \n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "\n",
    "from llama_index.llms.llama_cpp import LlamaCPP\n",
    "from llama_index.llms.llama_cpp.llama_utils import messages_to_prompt, completion_to_prompt\n",
    "llm = LlamaCPP(\n",
    "    model_url=None, # We'll load locally.\n",
    "    model_path='./Models/yi-34b-chat.Q4_K_M.gguf',\n",
    "    temperature=0.1,\n",
    "    max_new_tokens=1024, # Increasing to support longer responses\n",
    "    context_window=32768, # Yi 34B 32K context window!\n",
    "    generate_kwargs={},\n",
    "    # set to at least 1 to use GPU\n",
    "    model_kwargs={\"n_gpu_layers\": 20}, # Had to drop to 20 layers for 3090 or it was running out of memory\n",
    "    messages_to_prompt=messages_to_prompt,\n",
    "    completion_to_prompt=completion_to_prompt,\n",
    "    verbose=True\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 4. Checkpoint\n",
    "\n",
    "Are you running on GPU? The above output should include near the top something like:\n",
    "> ggml_init_cublas: found 1 CUDA devices:\n",
    "\n",
    "And in the full text near the bottom should be:\n",
    "> llm_load_tensors: using CUDA for GPU acceleration"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 5. Embeddings\n",
    "\n",
    "Convert your source document text into embeddings.\n",
    "\n",
    "The embedding model is from huggingface, this one performs well.\n",
    "\n",
    "> https://huggingface.co/thenlper/gte-large\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.embeddings.huggingface import HuggingFaceEmbedding\n",
    "\n",
    "embed_model = HuggingFaceEmbedding(model_name=\"thenlper/gte-large\", cache_folder=None)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 6. Prompt Template\n",
    "\n",
    "Prompt template for Yi 34B is the ChatML template:\n",
    "\n",
    "<|im_start|>system<br>\n",
    "{system_message}<|im_end|><br>\n",
    "<|im_start|>user<br>\n",
    "{prompt}<|im_end|><br>\n",
    "<|im_start|>assistant"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Produces a prompt for the Llama2 model\n",
    "def chatml_prompt(systemmessage, promptmessage):\n",
    "    return f\"<|im_start|>system\\n{systemmessage}<|im_end|>\\n<|im_start|>user\\n{promptmessage}<|im_end|>\\n<|im_start|>assistant\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 7. Service Context\n",
    "\n",
    "For chunking the document into tokens using the embedding model and our LLM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core import Settings\n",
    "\n",
    "Settings.llm = llm\n",
    "Settings.embed_model = embed_model\n",
    "Settings.chunk_size=256 # Number of tokens in each chunk"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 8. Index documents"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "index = VectorStoreIndex.from_documents(documents)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 9. Query Engine\n",
    "\n",
    "Create a query engine, specifying how many citations we want to get back from the searched text (in this case 3).\n",
    "\n",
    "The DB_DOC_ID_KEY is used to get back the filename of the original document"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.query_engine import CitationQueryEngine\n",
    "query_engine = CitationQueryEngine.from_args(\n",
    "    index,\n",
    "    similarity_top_k=3,\n",
    "    # here we can control how granular citation sources are, the default is 512\n",
    "    citation_chunk_size=256,\n",
    ")\n",
    "\n",
    "# For citations we get the document info\n",
    "DB_DOC_ID_KEY = \"db_document_id\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 10. Prompt and Response function\n",
    "\n",
    "Pass in a question, get a response back.\n",
    "\n",
    "IMPORTANT: The prompt is set here, adjust it to match what you want the LLM to act like and do."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "def RunQuestion(questionText):\n",
    "    systemmessage = \"You are a story teller who likes to elaborate. Answer questions in a positive, helpful and interesting way. If the answer is not in the following context return ONLY 'Sorry, I don't know the answer to that'.\"\n",
    "\n",
    "    queryQuestion = chatml_prompt(systemmessage, questionText)\n",
    "\n",
    "    response = query_engine.query(queryQuestion)\n",
    "\n",
    "    return response"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 11. Questions to test with"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "TestQuestions = [\n",
    "    \"Summarise the story for me\",\n",
    "    \"Who was the main protagonist?\",\n",
    "    \"Did they have any children? If so, what were their names?\",\n",
    "    \"Did anything eventful happen?\",\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 12. Run Questions through model (this can take a while) and see citations\n",
    "\n",
    "Runs each test question, saves it to a dictionary for output in the last step.\n",
    "\n",
    "Note: Citations are the source documents used and the text the response is based on. This is important for RAG so you can reference these documents for the user, and to ensure it's utilising the right documents."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "1/4: Summarise the story for me\n"
     ]
    },
    {
     "ename": "",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31mThe Kernel crashed while executing code in the current cell or a previous cell. \n",
      "\u001b[1;31mPlease review the code in the cell(s) to identify a possible cause of the failure. \n",
      "\u001b[1;31mClick <a href='https://aka.ms/vscodeJupyterKernelCrash'>here</a> for more info. \n",
      "\u001b[1;31mView Jupyter <a href='command:jupyter.viewOutput'>log</a> for further details."
     ]
    }
   ],
   "source": [
    "qa_pairs = []\n",
    "\n",
    "for index, question in enumerate(TestQuestions, start=1):\n",
    "    question = question.strip() # Clean up\n",
    "\n",
    "    print(f\"\\n{index}/{len(TestQuestions)}: {question}\")\n",
    "\n",
    "    response = RunQuestion(question) # Query and get  response\n",
    "\n",
    "    qa_pairs.append((question.strip(), str(response).strip())) # Add to our output array\n",
    "\n",
    "    # Displays the citations\n",
    "    for index, node in enumerate(response.source_nodes, start=1):\n",
    "        print(f\"{index}/{len(response.source_nodes)}: |{node.node.metadata['file_name']}| {node.node.get_text()}\")\n",
    "\n",
    "    # Uncomment the following line if you want to test just the first question\n",
    "    # break "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 13. Output responses"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for index, (question, answer) in enumerate(qa_pairs, start=1):\n",
    "    print(f\"{index}/{len(qa_pairs)} {question}\\n\\n{answer}\\n\\n--------\\n\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "llamaindexgeneric",
   "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": 2
}