notebooks/llm-rag-langchain/llm-rag-langchain.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "02a561f4",
   "metadata": {},
   "source": [
    "# Create a RAG system using OpenVINO and LangChain\n",
    "\n",
    "**Retrieval-augmented generation (RAG)** is a technique for augmenting LLM knowledge with additional, often private or real-time, data. LLMs can reason about wide-ranging topics, but their knowledge is limited to the public data up to a specific point in time that they were trained on. If you want to build AI applications that can reason about private data or data introduced after a model’s cutoff date, you need to augment the knowledge of the model with the specific information it needs. The process of bringing the appropriate information and inserting it into the model prompt is known as Retrieval Augmented Generation (RAG).\n",
    "\n",
    "[LangChain](https://python.langchain.com/docs/get_started/introduction) is a framework for developing applications powered by language models. It has a number of components specifically designed to help build RAG applications. In this tutorial, we’ll build a simple question-answering application over a text data source.\n",
    "\n",
    "The tutorial consists of the following steps:\n",
    "\n",
    "- Install prerequisites\n",
    "- Download and convert the model from a public source using the [OpenVINO integration with Hugging Face Optimum](https://huggingface.co/blog/openvino).\n",
    "- Compress model weights to 4-bit or 8-bit data types using [NNCF](https://github.com/openvinotoolkit/nncf)\n",
    "- Create a RAG chain pipeline\n",
    "- Run Q&A pipeline\n",
    "\n",
    "In this example, the customized RAG pipeline consists of following components in order, where embedding, rerank and LLM will be deployed with OpenVINO to optimize their inference performance.\n",
    "\n",
    "![RAG](https://github.com/openvinotoolkit/openvino_notebooks/assets/91237924/0076f6c7-75e4-4c2e-9015-87b355e5ca28)\n",
    "\n",
    "#### Table of contents:\n",
    "\n",
    "- [Prerequisites](#Prerequisites)\n",
    "- [Select model for inference](#Select-model-for-inference)\n",
    "- [login to huggingfacehub to get access to pretrained model](#login-to-huggingfacehub-to-get-access-to-pretrained-model)\n",
    "- [Convert model and compress model weights](#convert-model-and-compress-model-weights)\n",
    "  - [LLM conversion and Weights Compression using Optimum-CLI](#LLM-conversion-and-Weights-Compression-using-Optimum-CLI)\n",
    "  - [Convert embedding model using Optimum-CLI](#Convert-embedding-model-using-Optimum-CLI)\n",
    "  - [Convert rerank model using Optimum-CLI](#Convert-rerank-model-using-Optimum-CLI)\n",
    "- [Select device for inference and model variant](#Select-device-for-inference-and-model-variant)\n",
    "  - [Select device for embedding model inference](#Select-device-for-embedding-model-inference)\n",
    "  - [Select device for rerank model inference](#Select-device-for-rerank-model-inference)\n",
    "  - [Select device for LLM model inference](#Select-device-for-LLM-model-inference)\n",
    "- [Load model](#Load-model)\n",
    "  - [Load embedding model](#Load-embedding-model)\n",
    "  - [Load rerank model](#Load-rerank-model)\n",
    "  - [Load LLM model](#Load-LLM-model)\n",
    "- [Run QA over Document](#Run-QA-over-Document)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7c09cb8f",
   "metadata": {},
   "source": [
    "## Prerequisites\n",
    "\n",
    "[back to top ⬆️](#Table-of-contents:)\n",
    "\n",
    "Install required dependencies\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "id": "1f077b32-5d36-44b0-9041-407e996283a3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[33mWARNING: Skipping openvino-dev as it is not installed.\u001b[0m\u001b[33m\n",
      "\u001b[0m\u001b[33mWARNING: Skipping openvino as it is not installed.\u001b[0m\u001b[33m\n",
      "\u001b[0mNote: you may need to restart the kernel to use updated packages.\n",
      "Note: you may need to restart the kernel to use updated packages.\n"
     ]
    }
   ],
   "source": [
    "%pip uninstall -q -y openvino-dev openvino openvino-nightly optimum optimum-intel\n",
    "%pip install -q --extra-index-url https://download.pytorch.org/whl/cpu\\\n",
    "\"git+https://github.com/huggingface/optimum-intel.git\"\\\n",
    "\"git+https://github.com/openvinotoolkit/nncf.git\"\\\n",
    "\"datasets\"\\\n",
    "\"accelerate\"\\\n",
    "\"openvino-nightly\"\\\n",
    "\"gradio\"\\\n",
    "\"onnx\" \"einops\" \"transformers_stream_generator\" \"tiktoken\" \"transformers>=4.38.1\" \"bitsandbytes\" \"chromadb\" \"sentence_transformers\" \"langchain>=0.1.15\" \"langchainhub\" \"unstructured\" \"scikit-learn\" \"python-docx\" \"pypdf\" "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "1b2c3f4e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from pathlib import Path\n",
    "import requests\n",
    "import io\n",
    "\n",
    "# fetch model configuration\n",
    "\n",
    "config_shared_path = Path(\"../../utils/llm_config.py\")\n",
    "config_dst_path = Path(\"llm_config.py\")\n",
    "text_example_en_path = Path(\"text_example_en.pdf\")\n",
    "text_example_cn_path = Path(\"text_example_cn.pdf\")\n",
    "text_example_en = \"https://github.com/openvinotoolkit/openvino_notebooks/files/15039728/Platform.Brief_Intel.vPro.with.Intel.Core.Ultra_Final.pdf\"\n",
    "text_example_cn = \"https://github.com/openvinotoolkit/openvino_notebooks/files/15039713/Platform.Brief_Intel.vPro.with.Intel.Core.Ultra_Final_CH.pdf\"\n",
    "\n",
    "if not config_dst_path.exists():\n",
    "    if config_shared_path.exists():\n",
    "        os.symlink(config_shared_path, config_dst_path)\n",
    "    else:\n",
    "        r = requests.get(url=\"https://raw.githubusercontent.com/openvinotoolkit/openvino_notebooks/latest/utils/llm_config.py\")\n",
    "        with open(\"llm_config.py\", \"w\") as f:\n",
    "            f.write(r.text)\n",
    "\n",
    "if not text_example_en_path.exists():\n",
    "    r = requests.get(url=text_example_en)\n",
    "    content = io.BytesIO(r.content)\n",
    "    with open(\"text_example_en.pdf\", \"wb\") as f:\n",
    "        f.write(content.read())\n",
    "\n",
    "if not text_example_cn_path.exists():\n",
    "    r = requests.get(url=text_example_cn)\n",
    "    content = io.BytesIO(r.content)\n",
    "    with open(\"text_example_cn.pdf\", \"wb\") as f:\n",
    "        f.write(content.read())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c8e7965f",
   "metadata": {},
   "source": [
    "## Select model for inference\n",
    "\n",
    "[back to top ⬆️](#Table-of-contents:)\n",
    "\n",
    "The tutorial supports different models, you can select one from the provided options to compare the quality of open source LLM solutions.\n",
    "\n",
    "> **Note**: conversion of some models can require additional actions from user side and at least 64GB RAM for conversion.\n",
    "\n",
    "The available embedding model options are:\n",
    "\n",
    "- [**bge-small-en-v1.5**](https://huggingface.co/BAAI/bge-small-en-v1.5)\n",
    "- [**bge-small-zh-v1.5**](https://huggingface.co/BAAI/bge-small-zh-v1.5)\n",
    "- [**bge-large-en-v1.5**](https://huggingface.co/BAAI/bge-large-en-v1.5)\n",
    "- [**bge-large-zh-v1.5**](https://huggingface.co/BAAI/bge-large-zh-v1.5)\n",
    "\n",
    "BGE embedding is a general Embedding Model. The model is pre-trained using RetroMAE and trained on large-scale pair data using contrastive learning.\n",
    "\n",
    "The available rerank model options are:\n",
    "\n",
    "- [**bge-reranker-large**](https://huggingface.co/BAAI/bge-reranker-large)\n",
    "- [**bge-reranker-base**](https://huggingface.co/BAAI/bge-reranker-base)\n",
    "\n",
    "Reranker model with cross-encoder will perform full-attention over the input pair, which is more accurate than embedding model (i.e., bi-encoder) but more time-consuming than embedding model. Therefore, it can be used to re-rank the top-k documents returned by embedding model.\n",
    "\n",
    "You can also find available LLM model options in [llm-chatbot](../llm-chatbot/README.md) notebook.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "d3b57cfb-e727-43a5-b2c9-8f1b1ba72061",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "import openvino as ov\n",
    "import torch\n",
    "import ipywidgets as widgets\n",
    "from transformers import (\n",
    "    TextIteratorStreamer,\n",
    "    StoppingCriteria,\n",
    "    StoppingCriteriaList,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b51ff5a9",
   "metadata": {},
   "source": [
    "## Convert model and compress model weights\n",
    "\n",
    "[back to top ⬆️](#Table-of-contents:)\n",
    "\n",
    "The Weights Compression algorithm is aimed at compressing the weights of the models and can be used to optimize the model footprint and performance of large models where the size of weights is relatively larger than the size of activations, for example, Large Language Models (LLM). Compared to INT8 compression, INT4 compression improves performance even more, but introduces a minor drop in prediction quality."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "37bf49d7",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "19014cafd0614e31838477b41fb405ef",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Dropdown(description='Model Language:', options=('English', 'Chinese', 'Japanese'), value='English')"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from llm_config import (\n",
    "    SUPPORTED_EMBEDDING_MODELS,\n",
    "    SUPPORTED_RERANK_MODELS,\n",
    "    SUPPORTED_LLM_MODELS,\n",
    ")\n",
    "\n",
    "model_languages = list(SUPPORTED_LLM_MODELS)\n",
    "\n",
    "model_language = widgets.Dropdown(\n",
    "    options=model_languages,\n",
    "    value=model_languages[0],\n",
    "    description=\"Model Language:\",\n",
    "    disabled=False,\n",
    ")\n",
    "\n",
    "model_language"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "184d1678-0e73-4f35-8af5-1a7d291c2e6e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e12d1bdee7894f82a015cb3e0a85c4f5",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Dropdown(description='Model:', index=9, options=('tiny-llama-1b-chat', 'gemma-2b-it', 'red-pajama-3b-chat', 'g…"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "llm_model_ids = [model_id for model_id, model_config in SUPPORTED_LLM_MODELS[model_language.value].items() if model_config.get(\"rag_prompt_template\")]\n",
    "\n",
    "llm_model_id = widgets.Dropdown(\n",
    "    options=llm_model_ids,\n",
    "    value=llm_model_ids[-1],\n",
    "    description=\"Model:\",\n",
    "    disabled=False,\n",
    ")\n",
    "\n",
    "llm_model_id"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "49ea95f8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Selected LLM model neural-chat-7b-v3-1\n"
     ]
    }
   ],
   "source": [
    "llm_model_configuration = SUPPORTED_LLM_MODELS[model_language.value][llm_model_id.value]\n",
    "print(f\"Selected LLM model {llm_model_id.value}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5d370f13",
   "metadata": {},
   "source": [
    "🤗 [Optimum Intel](https://huggingface.co/docs/optimum/intel/index) is the interface between the 🤗 [Transformers](https://huggingface.co/docs/transformers/index) and [Diffusers](https://huggingface.co/docs/diffusers/index) libraries and OpenVINO to accelerate end-to-end pipelines on Intel architectures. It provides ease-to-use cli interface for exporting models to [OpenVINO Intermediate Representation (IR)](https://docs.openvino.ai/2024/documentation/openvino-ir-format.html) format.\n",
    "\n",
    "The command bellow demonstrates basic command for model export with `optimum-cli`\n",
    "\n",
    "```\n",
    "optimum-cli export openvino --model <model_id_or_path> --task <task> <out_dir>\n",
    "```\n",
    "\n",
    "where `--model` argument is model id from HuggingFace Hub or local directory with model (saved using `.save_pretrained` method), `--task ` is one of [supported task](https://huggingface.co/docs/optimum/exporters/task_manager) that exported model should solve. For LLMs it will be `text-generation-with-past`. If model initialization requires to use remote code, `--trust-remote-code` flag additionally should be passed.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6337eab8",
   "metadata": {},
   "source": [
    "### LLM conversion and Weights Compression using Optimum-CLI\n",
    "\n",
    "[back to top ⬆️](#Table-of-contents:)\n",
    "\n",
    "You can also apply fp16, 8-bit or 4-bit weight compression on the Linear, Convolutional and Embedding layers when exporting your model with the CLI by setting `--weight-format` to respectively fp16, int8 or int4. This type of optimization allows to reduce the memory footprint and inference latency.\n",
    "By default the quantization scheme for int8/int4 will be [asymmetric](https://github.com/openvinotoolkit/nncf/blob/develop/docs/compression_algorithms/Quantization.md#asymmetric-quantization), to make it [symmetric](https://github.com/openvinotoolkit/nncf/blob/develop/docs/compression_algorithms/Quantization.md#symmetric-quantization) you can add `--sym`.\n",
    "\n",
    "For INT4 quantization you can also specify the following arguments :\n",
    "\n",
    "- The `--group-size` parameter will define the group size to use for quantization, -1 it will results in per-column quantization.\n",
    "- The `--ratio` parameter controls the ratio between 4-bit and 8-bit quantization. If set to 0.9, it means that 90% of the layers will be quantized to int4 while 10% will be quantized to int8.\n",
    "\n",
    "Smaller group_size and ratio values usually improve accuracy at the sacrifice of the model size and inference latency.\n",
    "\n",
    "> **Note**: There may be no speedup for INT4/INT8 compressed models on dGPU.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "c6a38153",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "57d17e276ae14a66b72326067dcfe694",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Checkbox(value=True, description='Prepare INT4 model')"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3eac17e4d1a44d3a899b191de049cab7",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Checkbox(value=False, description='Prepare INT8 model')"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "dc5a4a193e3f46aebe68a212a8e88e52",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Checkbox(value=False, description='Prepare FP16 model')"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from IPython.display import Markdown, display\n",
    "\n",
    "prepare_int4_model = widgets.Checkbox(\n",
    "    value=True,\n",
    "    description=\"Prepare INT4 model\",\n",
    "    disabled=False,\n",
    ")\n",
    "prepare_int8_model = widgets.Checkbox(\n",
    "    value=False,\n",
    "    description=\"Prepare INT8 model\",\n",
    "    disabled=False,\n",
    ")\n",
    "prepare_fp16_model = widgets.Checkbox(\n",
    "    value=False,\n",
    "    description=\"Prepare FP16 model\",\n",
    "    disabled=False,\n",
    ")\n",
    "\n",
    "display(prepare_int4_model)\n",
    "display(prepare_int8_model)\n",
    "display(prepare_fp16_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "2020d522",
   "metadata": {},
   "outputs": [],
   "source": [
    "pt_model_id = llm_model_configuration[\"model_id\"]\n",
    "pt_model_name = llm_model_id.value.split(\"-\")[0]\n",
    "fp16_model_dir = Path(llm_model_id.value) / \"FP16\"\n",
    "int8_model_dir = Path(llm_model_id.value) / \"INT8_compressed_weights\"\n",
    "int4_model_dir = Path(llm_model_id.value) / \"INT4_compressed_weights\"\n",
    "\n",
    "\n",
    "def convert_to_fp16():\n",
    "    if (fp16_model_dir / \"openvino_model.xml\").exists():\n",
    "        return\n",
    "    remote_code = llm_model_configuration.get(\"remote_code\", False)\n",
    "    export_command_base = \"optimum-cli export openvino --model {} --task text-generation-with-past --weight-format fp16\".format(pt_model_id)\n",
    "    if remote_code:\n",
    "        export_command_base += \" --trust-remote-code\"\n",
    "    export_command = export_command_base + \" \" + str(fp16_model_dir)\n",
    "    display(Markdown(\"**Export command:**\"))\n",
    "    display(Markdown(f\"`{export_command}`\"))\n",
    "    ! $export_command\n",
    "\n",
    "\n",
    "def convert_to_int8():\n",
    "    if (int8_model_dir / \"openvino_model.xml\").exists():\n",
    "        return\n",
    "    int8_model_dir.mkdir(parents=True, exist_ok=True)\n",
    "    remote_code = llm_model_configuration.get(\"remote_code\", False)\n",
    "    export_command_base = \"optimum-cli export openvino --model {} --task text-generation-with-past --weight-format int8\".format(pt_model_id)\n",
    "    if remote_code:\n",
    "        export_command_base += \" --trust-remote-code\"\n",
    "    export_command = export_command_base + \" \" + str(int8_model_dir)\n",
    "    display(Markdown(\"**Export command:**\"))\n",
    "    display(Markdown(f\"`{export_command}`\"))\n",
    "    ! $export_command\n",
    "\n",
    "\n",
    "def convert_to_int4():\n",
    "    compression_configs = {\n",
    "        \"zephyr-7b-beta\": {\n",
    "            \"sym\": True,\n",
    "            \"group_size\": 64,\n",
    "            \"ratio\": 0.6,\n",
    "        },\n",
    "        \"mistral-7b\": {\n",
    "            \"sym\": True,\n",
    "            \"group_size\": 64,\n",
    "            \"ratio\": 0.6,\n",
    "        },\n",
    "        \"minicpm-2b-dpo\": {\n",
    "            \"sym\": True,\n",
    "            \"group_size\": 64,\n",
    "            \"ratio\": 0.6,\n",
    "        },\n",
    "        \"gemma-2b-it\": {\n",
    "            \"sym\": True,\n",
    "            \"group_size\": 64,\n",
    "            \"ratio\": 0.6,\n",
    "        },\n",
    "        \"notus-7b-v1\": {\n",
    "            \"sym\": True,\n",
    "            \"group_size\": 64,\n",
    "            \"ratio\": 0.6,\n",
    "        },\n",
    "        \"neural-chat-7b-v3-1\": {\n",
    "            \"sym\": True,\n",
    "            \"group_size\": 64,\n",
    "            \"ratio\": 0.6,\n",
    "        },\n",
    "        \"llama-2-chat-7b\": {\n",
    "            \"sym\": True,\n",
    "            \"group_size\": 128,\n",
    "            \"ratio\": 0.8,\n",
    "        },\n",
    "        \"llama-3-8b-instruct\": {\n",
    "            \"sym\": True,\n",
    "            \"group_size\": 128,\n",
    "            \"ratio\": 0.8,\n",
    "        },\n",
    "        \"gemma-7b-it\": {\n",
    "            \"sym\": True,\n",
    "            \"group_size\": 128,\n",
    "            \"ratio\": 0.8,\n",
    "        },\n",
    "        \"chatglm2-6b\": {\n",
    "            \"sym\": True,\n",
    "            \"group_size\": 128,\n",
    "            \"ratio\": 0.72,\n",
    "        },\n",
    "        \"qwen-7b-chat\": {\"sym\": True, \"group_size\": 128, \"ratio\": 0.6},\n",
    "        \"red-pajama-3b-chat\": {\n",
    "            \"sym\": False,\n",
    "            \"group_size\": 128,\n",
    "            \"ratio\": 0.5,\n",
    "        },\n",
    "        \"default\": {\n",
    "            \"sym\": False,\n",
    "            \"group_size\": 128,\n",
    "            \"ratio\": 0.8,\n",
    "        },\n",
    "    }\n",
    "\n",
    "    model_compression_params = compression_configs.get(llm_model_id.value, compression_configs[\"default\"])\n",
    "    if (int4_model_dir / \"openvino_model.xml\").exists():\n",
    "        return\n",
    "    remote_code = llm_model_configuration.get(\"remote_code\", False)\n",
    "    export_command_base = \"optimum-cli export openvino --model {} --task text-generation-with-past --weight-format int4\".format(pt_model_id)\n",
    "    int4_compression_args = \" --group-size {} --ratio {}\".format(model_compression_params[\"group_size\"], model_compression_params[\"ratio\"])\n",
    "    if model_compression_params[\"sym\"]:\n",
    "        int4_compression_args += \" --sym\"\n",
    "    export_command_base += int4_compression_args\n",
    "    if remote_code:\n",
    "        export_command_base += \" --trust-remote-code\"\n",
    "    export_command = export_command_base + \" \" + str(int4_model_dir)\n",
    "    display(Markdown(\"**Export command:**\"))\n",
    "    display(Markdown(f\"`{export_command}`\"))\n",
    "    ! $export_command\n",
    "\n",
    "\n",
    "if prepare_fp16_model.value:\n",
    "    convert_to_fp16()\n",
    "if prepare_int8_model.value:\n",
    "    convert_to_int8()\n",
    "if prepare_int4_model.value:\n",
    "    convert_to_int4()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f1d1d1a2",
   "metadata": {},
   "source": [
    "Let's compare model size for different compression types\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "8e127215",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Size of model with INT4 compressed weights is 5069.90 MB\n"
     ]
    }
   ],
   "source": [
    "fp16_weights = fp16_model_dir / \"openvino_model.bin\"\n",
    "int8_weights = int8_model_dir / \"openvino_model.bin\"\n",
    "int4_weights = int4_model_dir / \"openvino_model.bin\"\n",
    "\n",
    "if fp16_weights.exists():\n",
    "    print(f\"Size of FP16 model is {fp16_weights.stat().st_size / 1024 / 1024:.2f} MB\")\n",
    "for precision, compressed_weights in zip([8, 4], [int8_weights, int4_weights]):\n",
    "    if compressed_weights.exists():\n",
    "        print(f\"Size of model with INT{precision} compressed weights is {compressed_weights.stat().st_size / 1024 / 1024:.2f} MB\")\n",
    "    if compressed_weights.exists() and fp16_weights.exists():\n",
    "        print(f\"Compression rate for INT{precision} model: {fp16_weights.stat().st_size / compressed_weights.stat().st_size:.3f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4f943465",
   "metadata": {},
   "source": [
    "### Convert embedding model using Optimum-CLI\n",
    "\n",
    "[back to top ⬆️](#Table-of-contents:)\n",
    "\n",
    "Since some embedding models can only support limited languages, we can filter them out according the LLM you selected.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "49c28d3a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "12de29912eb44365971be56e128d3266",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Dropdown(description='Embedding Model:', options=('bge-small-en-v1.5', 'bge-large-en-v1.5'), value='bge-small-…"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "embedding_model_id = list(SUPPORTED_EMBEDDING_MODELS[model_language.value])\n",
    "\n",
    "embedding_model_id = widgets.Dropdown(\n",
    "    options=embedding_model_id,\n",
    "    value=embedding_model_id[0],\n",
    "    description=\"Embedding Model:\",\n",
    "    disabled=False,\n",
    ")\n",
    "\n",
    "embedding_model_id"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "3594116d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Selected bge-small-en-v1.5 model\n"
     ]
    }
   ],
   "source": [
    "embedding_model_configuration = SUPPORTED_EMBEDDING_MODELS[model_language.value][embedding_model_id.value]\n",
    "print(f\"Selected {embedding_model_id.value} model\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b2fbc403-1cac-4864-8965-ad2fea0fd1ca",
   "metadata": {},
   "source": [
    "OpenVINO embedding model and tokenizer can be exported by `feature-extraction` task with `optimum-cli`.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "ff80e6eb-7923-40ef-93d8-5e6c56e50667",
   "metadata": {},
   "outputs": [],
   "source": [
    "export_command_base = \"optimum-cli export openvino --model {} --task feature-extraction\".format(embedding_model_configuration[\"model_id\"])\n",
    "export_command = export_command_base + \" \" + str(embedding_model_id.value)\n",
    "\n",
    "if not Path(embedding_model_id.value).exists():\n",
    "    ! $export_command"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a2d818f0",
   "metadata": {},
   "source": [
    "### Convert rerank model using Optimum-CLI\n",
    "\n",
    "[back to top ⬆️](#Table-of-contents:)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "1b5b8840",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3fc1e02144ae4bc788a21a9804cd4e3f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Dropdown(description='Rerank Model:', options=('bge-reranker-large', 'bge-reranker-base'), value='bge-reranker…"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "rerank_model_id = list(SUPPORTED_RERANK_MODELS)\n",
    "\n",
    "rerank_model_id = widgets.Dropdown(\n",
    "    options=rerank_model_id,\n",
    "    value=rerank_model_id[0],\n",
    "    description=\"Rerank Model:\",\n",
    "    disabled=False,\n",
    ")\n",
    "\n",
    "rerank_model_id"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "1ca6fe04",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Selected bge-reranker-large model\n"
     ]
    }
   ],
   "source": [
    "rerank_model_configuration = SUPPORTED_RERANK_MODELS[rerank_model_id.value]\n",
    "print(f\"Selected {rerank_model_id.value} model\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "400dbff2-9915-4df0-a1ef-01cda3543b27",
   "metadata": {},
   "source": [
    "Since `rerank` model is sort of sentence classification task, its OpenVINO IR and tokenizer can be exported by `text-classification` task with `optimum-cli`.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "d0bab20b",
   "metadata": {},
   "outputs": [],
   "source": [
    "export_command_base = \"optimum-cli export openvino --model {} --task text-classification\".format(rerank_model_configuration[\"model_id\"])\n",
    "export_command = export_command_base + \" \" + str(rerank_model_id.value)\n",
    "\n",
    "if not Path(rerank_model_id.value).exists():\n",
    "    ! $export_command"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "749b5bbd",
   "metadata": {},
   "source": [
    "## Select device for inference and model variant\n",
    "\n",
    "[back to top ⬆️](#Table-of-contents:)\n",
    "\n",
    "> **Note**: There may be no speedup for INT4/INT8 compressed models on dGPU.\n",
    "\n",
    "### Select device for embedding model inference\n",
    "\n",
    "[back to top ⬆️](#Table-of-contents:)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "e11e73cf",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9f6ce6f571c244e79e838c28abdc52bd",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Dropdown(description='Device:', options=('CPU', 'GPU', 'AUTO'), value='CPU')"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "core = ov.Core()\n",
    "\n",
    "support_devices = core.available_devices\n",
    "if \"NPU\" in support_devices:\n",
    "    support_devices.remove(\"NPU\")\n",
    "\n",
    "embedding_device = widgets.Dropdown(\n",
    "    options=support_devices + [\"AUTO\"],\n",
    "    value=\"CPU\",\n",
    "    description=\"Device:\",\n",
    "    disabled=False,\n",
    ")\n",
    "\n",
    "embedding_device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "9ab29b85",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Embedding model will be loaded to CPU device for text embedding\n"
     ]
    }
   ],
   "source": [
    "print(f\"Embedding model will be loaded to {embedding_device.value} device for text embedding\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "81b2644c",
   "metadata": {},
   "source": [
    "### Select device for rerank model inference\n",
    "\n",
    "[back to top ⬆️](#Table-of-contents:)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "e0a2586b-5811-420f-834a-2b68de207df7",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2fcb244b0ec04e50a313f40e96edc949",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Dropdown(description='Device:', options=('CPU', 'GPU', 'AUTO'), value='CPU')"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "rerank_device = widgets.Dropdown(\n",
    "    options=support_devices + [\"AUTO\"],\n",
    "    value=\"CPU\",\n",
    "    description=\"Device:\",\n",
    "    disabled=False,\n",
    ")\n",
    "\n",
    "rerank_device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "7b7a76b2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Rerenk model will be loaded to CPU device for text reranking\n"
     ]
    }
   ],
   "source": [
    "print(f\"Rerenk model will be loaded to {rerank_device.value} device for text reranking\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ef31656a",
   "metadata": {},
   "source": [
    "### Select device for LLM model inference\n",
    "\n",
    "[back to top ⬆️](#Table-of-contents:)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "6d044d01",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "dd67900c499d4e5bb7bf1aec67d8d47f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Dropdown(description='Device:', options=('CPU', 'GPU', 'AUTO'), value='CPU')"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "llm_device = widgets.Dropdown(\n",
    "    options=support_devices + [\"AUTO\"],\n",
    "    value=\"CPU\",\n",
    "    description=\"Device:\",\n",
    "    disabled=False,\n",
    ")\n",
    "\n",
    "llm_device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "348b90fe",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "LLM model will be loaded to CPU device for response generation\n"
     ]
    }
   ],
   "source": [
    "print(f\"LLM model will be loaded to {llm_device.value} device for response generation\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bc225391",
   "metadata": {},
   "source": [
    "## Load models\n",
    "\n",
    "[back to top ⬆️](#Table-of-contents:)\n",
    "\n",
    "### Load embedding model\n",
    "\n",
    "[back to top ⬆️](#Table-of-contents:)\n",
    "\n",
    "Now a Hugging Face embedding model can be supported by OpenVINO through [`OpenVINOEmbeddings`](https://python.langchain.com/docs/integrations/text_embedding/openvino) and [`OpenVINOBgeEmbeddings`](https://python.langchain.com/docs/integrations/text_embedding/openvino#bge-with-openvino)classes of LangChain.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "df3e8fd1-d4c1-4e33-b46e-7840e392f8ee",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:nncf:NNCF initialized successfully. Supported frameworks detected: torch, tensorflow, onnx, openvino\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-04-28 21:05:33.318682: I tensorflow/core/util/port.cc:111] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.\n",
      "2024-04-28 21:05:33.322370: I tensorflow/tsl/cuda/cudart_stub.cc:28] Could not find cuda drivers on your machine, GPU will not be used.\n",
      "2024-04-28 21:05:33.366644: E tensorflow/compiler/xla/stream_executor/cuda/cuda_dnn.cc:9342] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered\n",
      "2024-04-28 21:05:33.366676: E tensorflow/compiler/xla/stream_executor/cuda/cuda_fft.cc:609] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered\n",
      "2024-04-28 21:05:33.366714: E tensorflow/compiler/xla/stream_executor/cuda/cuda_blas.cc:1518] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered\n",
      "2024-04-28 21:05:33.376052: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.\n",
      "To enable the following instructions: AVX2 AVX512F AVX512_VNNI FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.\n",
      "2024-04-28 21:05:34.068587: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT\n",
      "Compiling the model to CPU ...\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[-0.04208654910326004, 0.06681869924068451, 0.007916687056422234]"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from langchain_community.embeddings import OpenVINOBgeEmbeddings\n",
    "\n",
    "embedding_model_name = embedding_model_id.value\n",
    "embedding_model_kwargs = {\"device\": embedding_device.value}\n",
    "encode_kwargs = {\n",
    "    \"mean_pooling\": embedding_model_configuration[\"mean_pooling\"],\n",
    "    \"normalize_embeddings\": embedding_model_configuration[\"normalize_embeddings\"],\n",
    "}\n",
    "\n",
    "embedding = OpenVINOBgeEmbeddings(\n",
    "    model_name_or_path=embedding_model_name,\n",
    "    model_kwargs=embedding_model_kwargs,\n",
    "    encode_kwargs=encode_kwargs,\n",
    ")\n",
    "\n",
    "text = \"This is a test document.\"\n",
    "embedding_result = embedding.embed_query(text)\n",
    "embedding_result[:3]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d1a1fd58",
   "metadata": {},
   "source": [
    "### Load rerank model\n",
    "\n",
    "[back to top ⬆️](#Table-of-contents:)\n",
    "\n",
    "Now a Hugging Face embedding model can be supported by OpenVINO through [`OpenVINOReranker`](https://python.langchain.com/docs/integrations/document_transformers/openvino_rerank) class of LangChain.\n",
    "\n",
    "> **Note**: Rerank can be skipped in RAG.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "b67b39f2-8394-45fb-9b2b-ea63e267a2d3",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Compiling the model to CPU ...\n"
     ]
    }
   ],
   "source": [
    "from langchain_community.document_compressors.openvino_rerank import OpenVINOReranker\n",
    "\n",
    "rerank_model_name = rerank_model_id.value\n",
    "rerank_model_kwargs = {\"device\": rerank_device.value}\n",
    "rerank_top_n = 2\n",
    "\n",
    "reranker = OpenVINOReranker(\n",
    "    model_name_or_path=rerank_model_name,\n",
    "    model_kwargs=rerank_model_kwargs,\n",
    "    top_n=rerank_top_n,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "79fe990a",
   "metadata": {},
   "source": [
    "### Load LLM model\n",
    "\n",
    "[back to top ⬆️](#Table-of-contents:)\n",
    "\n",
    "OpenVINO models can be run locally through the `HuggingFacePipeline` class. To deploy a model with OpenVINO, you can specify the `backend=\"openvino\"` parameter to trigger OpenVINO as backend inference framework.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "90b968f3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b9610591570744258f024a79b31a39b1",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Dropdown(description='Model to run:', options=('INT4',), value='INT4')"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "available_models = []\n",
    "if int4_model_dir.exists():\n",
    "    available_models.append(\"INT4\")\n",
    "if int8_model_dir.exists():\n",
    "    available_models.append(\"INT8\")\n",
    "if fp16_model_dir.exists():\n",
    "    available_models.append(\"FP16\")\n",
    "\n",
    "model_to_run = widgets.Dropdown(\n",
    "    options=available_models,\n",
    "    value=available_models[0],\n",
    "    description=\"Model to run:\",\n",
    "    disabled=False,\n",
    ")\n",
    "\n",
    "model_to_run"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "20fcc33c",
   "metadata": {},
   "source": [
    "OpenVINO models can be run locally through the `HuggingFacePipeline` class in [LangChain](https://python.langchain.com/docs/integrations/llms/openvino/). To deploy a model with OpenVINO, you can specify the `backend=\"openvino\"` parameter to trigger OpenVINO as backend inference framework.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "f7f708db-8de1-4efd-94b2-fcabc48d52f4",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "The argument `trust_remote_code` is to be used along with export=True. It will be ignored.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading model from neural-chat-7b-v3-1/INT4_compressed_weights\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Compiling the model to CPU ...\n",
      "Setting `pad_token_id` to `eos_token_id`:2 for open-end generation.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'2 + 2 = 4'"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from langchain_community.llms.huggingface_pipeline import HuggingFacePipeline\n",
    "\n",
    "if model_to_run.value == \"INT4\":\n",
    "    model_dir = int4_model_dir\n",
    "elif model_to_run.value == \"INT8\":\n",
    "    model_dir = int8_model_dir\n",
    "else:\n",
    "    model_dir = fp16_model_dir\n",
    "print(f\"Loading model from {model_dir}\")\n",
    "\n",
    "ov_config = {\"PERFORMANCE_HINT\": \"LATENCY\", \"NUM_STREAMS\": \"1\", \"CACHE_DIR\": \"\"}\n",
    "\n",
    "# On a GPU device a model is executed in FP16 precision. For red-pajama-3b-chat model there known accuracy\n",
    "# issues caused by this, which we avoid by setting precision hint to \"f32\".\n",
    "if llm_model_id.value == \"red-pajama-3b-chat\" and \"GPU\" in core.available_devices and llm_device.value in [\"GPU\", \"AUTO\"]:\n",
    "    ov_config[\"INFERENCE_PRECISION_HINT\"] = \"f32\"\n",
    "\n",
    "llm = HuggingFacePipeline.from_model_id(\n",
    "    model_id=str(model_dir),\n",
    "    task=\"text-generation\",\n",
    "    backend=\"openvino\",\n",
    "    model_kwargs={\n",
    "        \"device\": llm_device.value,\n",
    "        \"ov_config\": ov_config,\n",
    "        \"trust_remote_code\": True,\n",
    "    },\n",
    "    pipeline_kwargs={\"max_new_tokens\": 2},\n",
    ")\n",
    "\n",
    "llm.invoke(\"2 + 2 =\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4fb8b0e4",
   "metadata": {},
   "source": [
    "## Run QA over Document\n",
    "\n",
    "[back to top ⬆️](#Table-of-contents:)\n",
    "\n",
    "Now, when model created, we can setup Chatbot interface using [Gradio](https://www.gradio.app/).\n",
    "\n",
    "A typical RAG application has two main components:\n",
    "\n",
    "- **Indexing**: a pipeline for ingesting data from a source and indexing it. This usually happen offline.\n",
    "\n",
    "- **Retrieval and generation**: the actual RAG chain, which takes the user query at run time and retrieves the relevant data from the index, then passes that to the model.\n",
    "\n",
    "The most common full sequence from raw data to answer looks like:\n",
    "\n",
    "**Indexing**\n",
    "\n",
    "1. `Load`: First we need to load our data. We’ll use DocumentLoaders for this.\n",
    "2. `Split`: Text splitters break large Documents into smaller chunks. This is useful both for indexing data and for passing it in to a model, since large chunks are harder to search over and won’t in a model’s finite context window.\n",
    "3. `Store`: We need somewhere to store and index our splits, so that they can later be searched over. This is often done using a VectorStore and Embeddings model.\n",
    "\n",
    "![Indexing pipeline](https://github.com/openvinotoolkit/openvino_notebooks/assets/91237924/dfed2ba3-0c3a-4e0e-a2a7-01638730486a)\n",
    "\n",
    "**Retrieval and generation**\n",
    "\n",
    "1. `Retrieve`: Given a user input, relevant splits are retrieved from storage using a Retriever.\n",
    "2. `Generate`: A LLM produces an answer using a prompt that includes the question and the retrieved data.\n",
    "\n",
    "![Retrieval and generation pipeline](https://github.com/openvinotoolkit/openvino_notebooks/assets/91237924/f0545ddc-c0cd-4569-8c86-9879fdab105a)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "5b97eeeb",
   "metadata": {},
   "outputs": [],
   "source": [
    "import re\n",
    "from typing import List\n",
    "from langchain.text_splitter import (\n",
    "    CharacterTextSplitter,\n",
    "    RecursiveCharacterTextSplitter,\n",
    "    MarkdownTextSplitter,\n",
    ")\n",
    "from langchain.document_loaders import (\n",
    "    CSVLoader,\n",
    "    EverNoteLoader,\n",
    "    PyPDFLoader,\n",
    "    TextLoader,\n",
    "    UnstructuredEPubLoader,\n",
    "    UnstructuredHTMLLoader,\n",
    "    UnstructuredMarkdownLoader,\n",
    "    UnstructuredODTLoader,\n",
    "    UnstructuredPowerPointLoader,\n",
    "    UnstructuredWordDocumentLoader,\n",
    ")\n",
    "\n",
    "\n",
    "class ChineseTextSplitter(CharacterTextSplitter):\n",
    "    def __init__(self, pdf: bool = False, **kwargs):\n",
    "        super().__init__(**kwargs)\n",
    "        self.pdf = pdf\n",
    "\n",
    "    def split_text(self, text: str) -> List[str]:\n",
    "        if self.pdf:\n",
    "            text = re.sub(r\"\\n{3,}\", \"\\n\", text)\n",
    "            text = text.replace(\"\\n\\n\", \"\")\n",
    "        sent_sep_pattern = re.compile('([﹒﹔﹖﹗．。！？][\"’”」』]{0,2}|(?=[\"‘“「『]{1,2}|$))')\n",
    "        sent_list = []\n",
    "        for ele in sent_sep_pattern.split(text):\n",
    "            if sent_sep_pattern.match(ele) and sent_list:\n",
    "                sent_list[-1] += ele\n",
    "            elif ele:\n",
    "                sent_list.append(ele)\n",
    "        return sent_list\n",
    "\n",
    "\n",
    "TEXT_SPLITERS = {\n",
    "    \"Character\": CharacterTextSplitter,\n",
    "    \"RecursiveCharacter\": RecursiveCharacterTextSplitter,\n",
    "    \"Markdown\": MarkdownTextSplitter,\n",
    "    \"Chinese\": ChineseTextSplitter,\n",
    "}\n",
    "\n",
    "\n",
    "LOADERS = {\n",
    "    \".csv\": (CSVLoader, {}),\n",
    "    \".doc\": (UnstructuredWordDocumentLoader, {}),\n",
    "    \".docx\": (UnstructuredWordDocumentLoader, {}),\n",
    "    \".enex\": (EverNoteLoader, {}),\n",
    "    \".epub\": (UnstructuredEPubLoader, {}),\n",
    "    \".html\": (UnstructuredHTMLLoader, {}),\n",
    "    \".md\": (UnstructuredMarkdownLoader, {}),\n",
    "    \".odt\": (UnstructuredODTLoader, {}),\n",
    "    \".pdf\": (PyPDFLoader, {}),\n",
    "    \".ppt\": (UnstructuredPowerPointLoader, {}),\n",
    "    \".pptx\": (UnstructuredPowerPointLoader, {}),\n",
    "    \".txt\": (TextLoader, {\"encoding\": \"utf8\"}),\n",
    "}\n",
    "\n",
    "chinese_examples = [\n",
    "    [\"英特尔®酷睿™ Ultra处理器可以降低多少功耗？\"],\n",
    "    [\"相比英特尔之前的移动处理器产品，英特尔®酷睿™ Ultra处理器的AI推理性能提升了多少？\"],\n",
    "    [\"英特尔博锐® Enterprise系统提供哪些功能？\"],\n",
    "]\n",
    "\n",
    "english_examples = [\n",
    "    [\"How much power consumption can Intel® Core™ Ultra Processors help save?\"],\n",
    "    [\"Compared to Intel’s previous mobile processor, what is the advantage of Intel® Core™ Ultra Processors for Artificial Intelligence?\"],\n",
    "    [\"What can Intel vPro® Enterprise systems offer?\"],\n",
    "]\n",
    "\n",
    "if model_language.value == \"English\":\n",
    "    text_example_path = \"text_example_en.pdf\"\n",
    "else:\n",
    "    text_example_path = \"text_example_cn.pdf\"\n",
    "\n",
    "examples = chinese_examples if (model_language.value == \"Chinese\") else english_examples"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "602f8ebd-789c-4eb2-b54d-b23d8f1d8e7b",
   "metadata": {},
   "source": [
    "We can build a RAG pipeline of LangChain through [`create_retrieval_chain`](https://python.langchain.com/docs/modules/chains/), which will help to create a chain to connect RAG components including:\n",
    "\n",
    "- [`Vector stores`](https://python.langchain.com/docs/modules/data_connection/vectorstores/)，\n",
    "- [`Retrievers`](https://python.langchain.com/docs/modules/data_connection/retrievers/)\n",
    "- [`LLM`](https://python.langchain.com/docs/integrations/llms/)\n",
    "- [`Embedding`](https://python.langchain.com/docs/integrations/text_embedding/)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0908e5e9-4dcb-4fc8-8480-3cf70fd5e934",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.prompts import PromptTemplate\n",
    "from langchain.vectorstores import Chroma\n",
    "from langchain.chains.retrieval import create_retrieval_chain\n",
    "from langchain.chains.combine_documents import create_stuff_documents_chain\n",
    "from langchain.docstore.document import Document\n",
    "from langchain.retrievers import ContextualCompressionRetriever\n",
    "from threading import Thread\n",
    "import gradio as gr\n",
    "\n",
    "stop_tokens = llm_model_configuration.get(\"stop_tokens\")\n",
    "\n",
    "\n",
    "class StopOnTokens(StoppingCriteria):\n",
    "    def __init__(self, token_ids):\n",
    "        self.token_ids = token_ids\n",
    "\n",
    "    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:\n",
    "        for stop_id in self.token_ids:\n",
    "            if input_ids[0][-1] == stop_id:\n",
    "                return True\n",
    "        return False\n",
    "\n",
    "\n",
    "if stop_tokens is not None:\n",
    "    if isinstance(stop_tokens[0], str):\n",
    "        stop_tokens = llm.pipeline.tokenizer.convert_tokens_to_ids(stop_tokens)\n",
    "\n",
    "    stop_tokens = [StopOnTokens(stop_tokens)]\n",
    "\n",
    "\n",
    "def load_single_document(file_path: str) -> List[Document]:\n",
    "    \"\"\"\n",
    "    helper for loading a single document\n",
    "\n",
    "    Params:\n",
    "      file_path: document path\n",
    "    Returns:\n",
    "      documents loaded\n",
    "\n",
    "    \"\"\"\n",
    "    ext = \".\" + file_path.rsplit(\".\", 1)[-1]\n",
    "    if ext in LOADERS:\n",
    "        loader_class, loader_args = LOADERS[ext]\n",
    "        loader = loader_class(file_path, **loader_args)\n",
    "        return loader.load()\n",
    "\n",
    "    raise ValueError(f\"File does not exist '{ext}'\")\n",
    "\n",
    "\n",
    "def default_partial_text_processor(partial_text: str, new_text: str):\n",
    "    \"\"\"\n",
    "    helper for updating partially generated answer, used by default\n",
    "\n",
    "    Params:\n",
    "      partial_text: text buffer for storing previosly generated text\n",
    "      new_text: text update for the current step\n",
    "    Returns:\n",
    "      updated text string\n",
    "\n",
    "    \"\"\"\n",
    "    partial_text += new_text\n",
    "    return partial_text\n",
    "\n",
    "\n",
    "text_processor = llm_model_configuration.get(\"partial_text_processor\", default_partial_text_processor)\n",
    "\n",
    "\n",
    "def create_vectordb(docs, spliter_name, chunk_size, chunk_overlap, vector_search_top_k, vector_search_top_n, run_rerank, search_method, score_threshold):\n",
    "    \"\"\"\n",
    "    Initialize a vector database\n",
    "\n",
    "    Params:\n",
    "      doc: orignal documents provided by user\n",
    "      chunk_size:  size of a single sentence chunk\n",
    "      chunk_overlap: overlap size between 2 chunks\n",
    "      vector_search_top_k: Vector search top k\n",
    "\n",
    "    \"\"\"\n",
    "    documents = []\n",
    "    for doc in docs:\n",
    "        documents.extend(load_single_document(doc.name))\n",
    "\n",
    "    text_splitter = TEXT_SPLITERS[spliter_name](chunk_size=chunk_size, chunk_overlap=chunk_overlap)\n",
    "\n",
    "    texts = text_splitter.split_documents(documents)\n",
    "\n",
    "    global db\n",
    "    db = Chroma.from_documents(texts, embedding)\n",
    "\n",
    "    global retriever\n",
    "    if search_method == \"similarity_score_threshold\":\n",
    "        search_kwargs = {\"k\": vector_search_top_k, \"score_threshold\": score_threshold}\n",
    "    else:\n",
    "        search_kwargs = {\"k\": vector_search_top_k}\n",
    "    retriever = db.as_retriever(search_kwargs=search_kwargs, search_type=search_method)\n",
    "    if run_rerank:\n",
    "        reranker.top_n = vector_search_top_n\n",
    "        retriever = ContextualCompressionRetriever(base_compressor=reranker, base_retriever=retriever)\n",
    "    prompt = PromptTemplate.from_template(llm_model_configuration[\"rag_prompt_template\"])\n",
    "\n",
    "    global combine_docs_chain\n",
    "    combine_docs_chain = create_stuff_documents_chain(llm, prompt)\n",
    "\n",
    "    global rag_chain\n",
    "    rag_chain = create_retrieval_chain(retriever, combine_docs_chain)\n",
    "\n",
    "    return \"Vector database is Ready\"\n",
    "\n",
    "\n",
    "def update_retriever(vector_search_top_k, vector_rerank_top_n, run_rerank, search_method, score_threshold):\n",
    "    \"\"\"\n",
    "    Update retriever\n",
    "\n",
    "    Params:\n",
    "      vector_search_top_k: size of searching results\n",
    "      vector_rerank_top_n:  size of rerank results\n",
    "      run_rerank: whether run rerank step\n",
    "      search_method: search method used by vector store\n",
    "\n",
    "    \"\"\"\n",
    "    global retriever\n",
    "    global db\n",
    "    global rag_chain\n",
    "    global combine_docs_chain\n",
    "\n",
    "    if search_method == \"similarity_score_threshold\":\n",
    "        search_kwargs = {\"k\": vector_search_top_k, \"score_threshold\": score_threshold}\n",
    "    else:\n",
    "        search_kwargs = {\"k\": vector_search_top_k}\n",
    "    retriever = db.as_retriever(search_kwargs=search_kwargs, search_type=search_method)\n",
    "    if run_rerank:\n",
    "        retriever = ContextualCompressionRetriever(base_compressor=reranker, base_retriever=retriever)\n",
    "        reranker.top_n = vector_rerank_top_n\n",
    "    rag_chain = create_retrieval_chain(retriever, combine_docs_chain)\n",
    "\n",
    "\n",
    "def user(message, history):\n",
    "    \"\"\"\n",
    "    callback function for updating user messages in interface on submit button click\n",
    "\n",
    "    Params:\n",
    "      message: current message\n",
    "      history: conversation history\n",
    "    Returns:\n",
    "      None\n",
    "    \"\"\"\n",
    "    # Append the user's message to the conversation history\n",
    "    return \"\", history + [[message, \"\"]]\n",
    "\n",
    "\n",
    "def bot(history, temperature, top_p, top_k, repetition_penalty, hide_full_prompt):\n",
    "    \"\"\"\n",
    "    callback function for running chatbot on submit button click\n",
    "\n",
    "    Params:\n",
    "      history: conversation history\n",
    "      temperature:  parameter for control the level of creativity in AI-generated text.\n",
    "                    By adjusting the `temperature`, you can influence the AI model's probability distribution, making the text more focused or diverse.\n",
    "      top_p: parameter for control the range of tokens considered by the AI model based on their cumulative probability.\n",
    "      top_k: parameter for control the range of tokens considered by the AI model based on their cumulative probability, selecting number of tokens with highest probability.\n",
    "      repetition_penalty: parameter for penalizing tokens based on how frequently they occur in the text.\n",
    "      hide_full_prompt: whether to show searching results in promopt.\n",
    "\n",
    "    \"\"\"\n",
    "    streamer = TextIteratorStreamer(\n",
    "        llm.pipeline.tokenizer,\n",
    "        timeout=60.0,\n",
    "        skip_prompt=hide_full_prompt,\n",
    "        skip_special_tokens=True,\n",
    "    )\n",
    "    llm.pipeline._forward_params = dict(\n",
    "        max_new_tokens=512,\n",
    "        temperature=temperature,\n",
    "        do_sample=temperature > 0.0,\n",
    "        top_p=top_p,\n",
    "        top_k=top_k,\n",
    "        repetition_penalty=repetition_penalty,\n",
    "        streamer=streamer,\n",
    "    )\n",
    "    if stop_tokens is not None:\n",
    "        llm.pipeline._forward_params[\"stopping_criteria\"] = StoppingCriteriaList(stop_tokens)\n",
    "\n",
    "    t1 = Thread(target=rag_chain.invoke, args=({\"input\": history[-1][0]},))\n",
    "    t1.start()\n",
    "\n",
    "    # Initialize an empty string to store the generated text\n",
    "    partial_text = \"\"\n",
    "    for new_text in streamer:\n",
    "        partial_text = text_processor(partial_text, new_text)\n",
    "        history[-1][1] = partial_text\n",
    "        yield history\n",
    "\n",
    "\n",
    "def request_cancel():\n",
    "    llm.pipeline.model.request.cancel()\n",
    "\n",
    "\n",
    "with gr.Blocks(\n",
    "    theme=gr.themes.Soft(),\n",
    "    css=\".disclaimer {font-variant-caps: all-small-caps;}\",\n",
    ") as demo:\n",
    "    gr.Markdown(\"\"\"<h1><center>QA over Document</center></h1>\"\"\")\n",
    "    gr.Markdown(f\"\"\"<center>Powered by OpenVINO and {llm_model_id.value} </center>\"\"\")\n",
    "    with gr.Row():\n",
    "        with gr.Column(scale=1):\n",
    "            docs = gr.File(\n",
    "                label=\"Step 1: Load text files\",\n",
    "                value=[text_example_path],\n",
    "                file_count=\"multiple\",\n",
    "                file_types=[\n",
    "                    \".csv\",\n",
    "                    \".doc\",\n",
    "                    \".docx\",\n",
    "                    \".enex\",\n",
    "                    \".epub\",\n",
    "                    \".html\",\n",
    "                    \".md\",\n",
    "                    \".odt\",\n",
    "                    \".pdf\",\n",
    "                    \".ppt\",\n",
    "                    \".pptx\",\n",
    "                    \".txt\",\n",
    "                ],\n",
    "            )\n",
    "            load_docs = gr.Button(\"Step 2: Build Vector Store\")\n",
    "            db_argument = gr.Accordion(\"Vector Store Configuration\", open=False)\n",
    "            with db_argument:\n",
    "                spliter = gr.Dropdown(\n",
    "                    [\"Character\", \"RecursiveCharacter\", \"Markdown\", \"Chinese\"],\n",
    "                    value=\"RecursiveCharacter\",\n",
    "                    label=\"Text Spliter\",\n",
    "                    info=\"Method used to splite the documents\",\n",
    "                    multiselect=False,\n",
    "                )\n",
    "\n",
    "                chunk_size = gr.Slider(\n",
    "                    label=\"Chunk size\",\n",
    "                    value=700,\n",
    "                    minimum=100,\n",
    "                    maximum=2000,\n",
    "                    step=50,\n",
    "                    interactive=True,\n",
    "                    info=\"Size of sentence chunk\",\n",
    "                )\n",
    "\n",
    "                chunk_overlap = gr.Slider(\n",
    "                    label=\"Chunk overlap\",\n",
    "                    value=100,\n",
    "                    minimum=0,\n",
    "                    maximum=400,\n",
    "                    step=10,\n",
    "                    interactive=True,\n",
    "                    info=(\"Overlap between 2 chunks\"),\n",
    "                )\n",
    "\n",
    "            langchain_status = gr.Textbox(\n",
    "                label=\"Vector Store Status\",\n",
    "                value=\"Vector Store is Not ready\",\n",
    "                interactive=False,\n",
    "            )\n",
    "            with gr.Accordion(\"Generation Configuration\", open=False):\n",
    "                with gr.Row():\n",
    "                    with gr.Column():\n",
    "                        with gr.Row():\n",
    "                            temperature = gr.Slider(\n",
    "                                label=\"Temperature\",\n",
    "                                value=0.1,\n",
    "                                minimum=0.0,\n",
    "                                maximum=1.0,\n",
    "                                step=0.1,\n",
    "                                interactive=True,\n",
    "                                info=\"Higher values produce more diverse outputs\",\n",
    "                            )\n",
    "                    with gr.Column():\n",
    "                        with gr.Row():\n",
    "                            top_p = gr.Slider(\n",
    "                                label=\"Top-p (nucleus sampling)\",\n",
    "                                value=1.0,\n",
    "                                minimum=0.0,\n",
    "                                maximum=1,\n",
    "                                step=0.01,\n",
    "                                interactive=True,\n",
    "                                info=(\n",
    "                                    \"Sample from the smallest possible set of tokens whose cumulative probability \"\n",
    "                                    \"exceeds top_p. Set to 1 to disable and sample from all tokens.\"\n",
    "                                ),\n",
    "                            )\n",
    "                    with gr.Column():\n",
    "                        with gr.Row():\n",
    "                            top_k = gr.Slider(\n",
    "                                label=\"Top-k\",\n",
    "                                value=50,\n",
    "                                minimum=0.0,\n",
    "                                maximum=200,\n",
    "                                step=1,\n",
    "                                interactive=True,\n",
    "                                info=\"Sample from a shortlist of top-k tokens — 0 to disable and sample from all tokens.\",\n",
    "                            )\n",
    "                    with gr.Column():\n",
    "                        with gr.Row():\n",
    "                            repetition_penalty = gr.Slider(\n",
    "                                label=\"Repetition Penalty\",\n",
    "                                value=1.1,\n",
    "                                minimum=1.0,\n",
    "                                maximum=2.0,\n",
    "                                step=0.1,\n",
    "                                interactive=True,\n",
    "                                info=\"Penalize repetition — 1.0 to disable.\",\n",
    "                            )\n",
    "        with gr.Column(scale=4):\n",
    "            chatbot = gr.Chatbot(\n",
    "                height=600,\n",
    "                label=\"Step 3: Input Query\",\n",
    "            )\n",
    "            with gr.Row():\n",
    "                with gr.Column():\n",
    "                    with gr.Row():\n",
    "                        msg = gr.Textbox(\n",
    "                            label=\"QA Message Box\",\n",
    "                            placeholder=\"Chat Message Box\",\n",
    "                            show_label=False,\n",
    "                            container=False,\n",
    "                        )\n",
    "                with gr.Column():\n",
    "                    with gr.Row():\n",
    "                        submit = gr.Button(\"Submit\")\n",
    "                        stop = gr.Button(\"Stop\")\n",
    "                        clear = gr.Button(\"Clear\")\n",
    "            gr.Examples(examples, inputs=msg, label=\"Click on any example and press the 'Submit' button\")\n",
    "            retriever_argument = gr.Accordion(\"Retriever Configuration\", open=True)\n",
    "            with retriever_argument:\n",
    "                with gr.Row():\n",
    "                    with gr.Row():\n",
    "                        do_rerank = gr.Checkbox(\n",
    "                            value=True,\n",
    "                            label=\"Rerank searching result\",\n",
    "                            interactive=True,\n",
    "                        )\n",
    "                        hide_context = gr.Checkbox(\n",
    "                            value=True,\n",
    "                            label=\"Hide searching result in prompt\",\n",
    "                            interactive=True,\n",
    "                        )\n",
    "                    with gr.Row():\n",
    "                        search_method = gr.Dropdown(\n",
    "                            [\"similarity_score_threshold\", \"similarity\", \"mmr\"],\n",
    "                            value=\"similarity_score_threshold\",\n",
    "                            label=\"Searching Method\",\n",
    "                            info=\"Method used to search vector store\",\n",
    "                            multiselect=False,\n",
    "                            interactive=True,\n",
    "                        )\n",
    "                    with gr.Row():\n",
    "                        score_threshold = gr.Slider(\n",
    "                            0.01,\n",
    "                            0.99,\n",
    "                            value=0.5,\n",
    "                            step=0.01,\n",
    "                            label=\"Similarity Threshold\",\n",
    "                            info=\"Only working for 'similarity score threshold' method\",\n",
    "                            interactive=True,\n",
    "                        )\n",
    "                    with gr.Row():\n",
    "                        vector_rerank_top_n = gr.Slider(\n",
    "                            1,\n",
    "                            10,\n",
    "                            value=2,\n",
    "                            step=1,\n",
    "                            label=\"Rerank top n\",\n",
    "                            info=\"Number of rerank results\",\n",
    "                            interactive=True,\n",
    "                        )\n",
    "                    with gr.Row():\n",
    "                        vector_search_top_k = gr.Slider(\n",
    "                            1,\n",
    "                            50,\n",
    "                            value=10,\n",
    "                            step=1,\n",
    "                            label=\"Search top k\",\n",
    "                            info=\"Number of searching results, must >= Rerank top n\",\n",
    "                            interactive=True,\n",
    "                        )\n",
    "    load_docs.click(\n",
    "        create_vectordb,\n",
    "        inputs=[docs, spliter, chunk_size, chunk_overlap, vector_search_top_k, vector_rerank_top_n, do_rerank, search_method, score_threshold],\n",
    "        outputs=[langchain_status],\n",
    "        queue=False,\n",
    "    )\n",
    "    submit_event = msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(\n",
    "        bot,\n",
    "        [chatbot, temperature, top_p, top_k, repetition_penalty, hide_context],\n",
    "        chatbot,\n",
    "        queue=True,\n",
    "    )\n",
    "    submit_click_event = submit.click(user, [msg, chatbot], [msg, chatbot], queue=False).then(\n",
    "        bot,\n",
    "        [chatbot, temperature, top_p, top_k, repetition_penalty, hide_context],\n",
    "        chatbot,\n",
    "        queue=True,\n",
    "    )\n",
    "    stop.click(\n",
    "        fn=request_cancel,\n",
    "        inputs=None,\n",
    "        outputs=None,\n",
    "        cancels=[submit_event, submit_click_event],\n",
    "        queue=False,\n",
    "    )\n",
    "    clear.click(lambda: None, None, chatbot, queue=False)\n",
    "    vector_search_top_k.release(\n",
    "        update_retriever,\n",
    "        [vector_search_top_k, vector_rerank_top_n, do_rerank, search_method, score_threshold],\n",
    "    )\n",
    "    vector_rerank_top_n.release(\n",
    "        update_retriever,\n",
    "        [vector_search_top_k, vector_rerank_top_n, do_rerank, search_method, score_threshold],\n",
    "    )\n",
    "    do_rerank.change(\n",
    "        update_retriever,\n",
    "        [vector_search_top_k, vector_rerank_top_n, do_rerank, search_method, score_threshold],\n",
    "    )\n",
    "    search_method.change(\n",
    "        update_retriever,\n",
    "        [vector_search_top_k, vector_rerank_top_n, do_rerank, search_method, score_threshold],\n",
    "    )\n",
    "    score_threshold.change(\n",
    "        update_retriever,\n",
    "        [vector_search_top_k, vector_rerank_top_n, do_rerank, search_method, score_threshold],\n",
    "    )\n",
    "\n",
    "\n",
    "demo.queue()\n",
    "# if you are launching remotely, specify server_name and server_port\n",
    "#  demo.launch(server_name='your server name', server_port='server port in int')\n",
    "# if you have any issue to launch on your platform, you can pass share=True to launch method:\n",
    "# demo.launch(share=True)\n",
    "# it creates a publicly shareable link for the interface. Read more in the docs: https://gradio.app/docs/\n",
    "demo.launch()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6f4b5a84-bebf-49b9-b2fa-5e788ed2cbac",
   "metadata": {},
   "outputs": [],
   "source": [
    "# please run this cell for stopping gradio interface\n",
    "demo.close()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.12"
  },
  "openvino_notebooks": {
   "imageUrl": "https://github.com/openvinotoolkit/openvino_notebooks/assets/29454499/304aa048-f10c-41c6-bb31-6d2bfdf49cf5",
   "tags": {
    "categories": [
     "Model Demos",
     "AI Trends"
    ],
    "libraries": [],
    "other": [
     "LLM"
    ],
    "tasks": [
     "Text Generation"
    ]
   }
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "state": {},
    "version_major": 2,
    "version_minor": 0
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}