notebooks/254-llm-chatbot/254-rag-chatbot.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "02a561f4",
   "metadata": {},
   "source": [
    "# Create an LLM-powered RAG system using OpenVINO\n",
    "\n",
    "> **Note:** This notebook has been moved to a new branch named \"latest\". [Click here](https://github.com/openvinotoolkit/openvino_notebooks/blob/latest/notebooks/llm-chatbot/rag-chatbot.ipynb) to get the most updated version of the notebook. This branch is deprecated.\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 Markdown or CSV 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 chat pipeline\n",
    "\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](#Convert-model)\n",
    "    - [Convert LLM model](#Convert-LLM-model)\n",
    "- [Compress model weights](#Compress-model-weights)\n",
    "    - [Weights Compression using Optimum Intel](#Weights-Compression-using-Optimum-Intel)\n",
    "    - [Weights Compression using NNCF](#Weights-Compression-using-NNCF)\n",
    "    - [Convert embedding model](#Convert-embedding-model)\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 LLM model inference](#Select-device-for-LLM-model-inference)\n",
    "- [Load model](#Load-model)\n",
    "    - [Load embedding model](#Load-embedding-model)\n",
    "    - [Load LLM model](#Load-LLM-model)\n",
    "- [Run QA over Document](#Run-QA-over-Document)\n",
    "\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "7c09cb8f",
   "metadata": {},
   "source": [
    "## Prerequisites\n",
    "[back to top ⬆️](#Table-of-contents:)\n",
    "\n",
    "Install required dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "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",
      "\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m A new release of pip is available: \u001b[0m\u001b[31;49m23.3.2\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m24.0\u001b[0m\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m To update, run: \u001b[0m\u001b[32;49mpip install --upgrade pip\u001b[0m\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.7\" \"langchainhub\" \"unstructured\" \"scikit-learn\" \"python-docx\" \"pdfminer.six\""
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "c8e7965f",
   "metadata": {},
   "source": [
    "## Select model for inference\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",
    ">**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",
    "* **all-mpnet-base-v2(All)** - This is a [sentence-transformers](https://huggingface.co/sentence-transformers) model: It maps sentences & paragraphs to a 768 dimensional dense vector space and can be used for tasks like clustering or semantic search. More details about model can be found in [model card](https://huggingface.co/sentence-transformers/all-mpnet-base-v2)\n",
    "* **text2vec-large-chinese(Chinese)** - This is a [CoSENT](https://github.com/bojone/CoSENT) model. It can be used for tasks like sentence embeddings, text matching or semantic search. More details about model can be found in [model card](https://huggingface.co/GanymedeNil/text2vec-base-chinese)\n",
    "\n",
    "The available LLM model options are:\n",
    "\n",
    "* **tiny-llama-1b-chat** - This is the chat model finetuned on top of [TinyLlama/TinyLlama-1.1B-intermediate-step-955k-2T](https://huggingface.co/TinyLlama/TinyLlama-1.1B-intermediate-step-955k-token-2T). The TinyLlama project aims to pretrain a 1.1B Llama model on 3 trillion tokens with the adoption of the same architecture and tokenizer as Llama 2. This means TinyLlama can be plugged and played in many open-source projects built upon Llama. Besides, TinyLlama is compact with only 1.1B parameters. This compactness allows it to cater to a multitude of applications demanding a restricted computation and memory footprint. More details about model can be found in [model card](https://huggingface.co/TinyLlama/TinyLlama-1.1B-Chat-v0.6)\n",
    "*  **gemma-2b-it** - Gemma is a family of lightweight, state-of-the-art open models from Google, built from the same research and technology used to create the Gemini models. They are text-to-text, decoder-only large language models, available in English, with open weights, pre-trained variants, and instruction-tuned variants. Gemma models are well-suited for a variety of text generation tasks, including question answering, summarization, and reasoning. This model is instruction-tuned version of 2B parameters model. More details about model can be found in [model card](https://huggingface.co/google/gemma-2b-it).\n",
    ">**Note**: run model with demo, you will need to accept license agreement. \n",
    ">You must be a registered user in 🤗 Hugging Face Hub. Please visit [HuggingFace model card](https://huggingface.co/google/gemma-2b-it), carefully read terms of usage and click accept button.  You will need to use an access token for the code below to run. For more information on access tokens, refer to [this section of the documentation](https://huggingface.co/docs/hub/security-tokens).\n",
    ">You can login on Hugging Face Hub in notebook environment, using following code:\n",
    " \n",
    "```python\n",
    "    ## login to huggingfacehub to get access to pretrained model \n",
    "\n",
    "    from huggingface_hub import notebook_login, whoami\n",
    "\n",
    "    try:\n",
    "        whoami()\n",
    "        print('Authorization token already provided')\n",
    "    except OSError:\n",
    "        notebook_login()\n",
    "```\n",
    "* **red-pajama-3b-chat** - A 2.8B parameter pre-trained language model based on GPT-NEOX architecture. It was developed by Together Computer and leaders from the open-source AI community. The model is fine-tuned on OASST1 and Dolly2 datasets to enhance chatting ability. More details about model can be found in [HuggingFace model card](https://huggingface.co/togethercomputer/RedPajama-INCITE-Chat-3B-v1).\n",
    "*  **gemma-7b-it** - Gemma is a family of lightweight, state-of-the-art open models from Google, built from the same research and technology used to create the Gemini models. They are text-to-text, decoder-only large language models, available in English, with open weights, pre-trained variants, and instruction-tuned variants. Gemma models are well-suited for a variety of text generation tasks, including question answering, summarization, and reasoning. This model is instruction-tuned version of 7B parameters model. More details about model can be found in [model card](https://huggingface.co/google/gemma-7b-it).\n",
    ">**Note**: run model with demo, you will need to accept license agreement. \n",
    ">You must be a registered user in 🤗 Hugging Face Hub. Please visit [HuggingFace model card](https://huggingface.co/google/gemma-7b-it), carefully read terms of usage and click accept button.  You will need to use an access token for the code below to run. For more information on access tokens, refer to [this section of the documentation](https://huggingface.co/docs/hub/security-tokens).\n",
    ">You can login on Hugging Face Hub in notebook environment, using following code:\n",
    " \n",
    "```python\n",
    "    ## login to huggingfacehub to get access to pretrained model \n",
    "\n",
    "    from huggingface_hub import notebook_login, whoami\n",
    "\n",
    "    try:\n",
    "        whoami()\n",
    "        print('Authorization token already provided')\n",
    "    except OSError:\n",
    "        notebook_login()\n",
    "```\n",
    "* **llama-2-7b-chat** - LLama 2 is the second generation of LLama models developed by Meta. Llama 2 is a collection of pre-trained and fine-tuned generative text models ranging in scale from 7 billion to 70 billion parameters. llama-2-7b-chat is 7 billions parameters version of LLama 2 finetuned and optimized for dialogue use case. More details about model can be found in the [paper](https://ai.meta.com/research/publications/llama-2-open-foundation-and-fine-tuned-chat-models/), [repository](https://github.com/facebookresearch/llama) and [HuggingFace model card](https://huggingface.co/meta-llama/Llama-2-7b-chat-hf).\n",
    ">**Note**: run model with demo, you will need to accept license agreement. \n",
    ">You must be a registered user in 🤗 Hugging Face Hub. Please visit [HuggingFace model card](https://huggingface.co/meta-llama/Llama-2-7b-chat-hf), carefully read terms of usage and click accept button.  You will need to use an access token for the code below to run. For more information on access tokens, refer to [this section of the documentation](https://huggingface.co/docs/hub/security-tokens).\n",
    ">You can login on Hugging Face Hub in notebook environment, using following code:\n",
    " \n",
    "```python\n",
    "    ## login to huggingfacehub to get access to pretrained model \n",
    "\n",
    "    from huggingface_hub import notebook_login, whoami\n",
    "\n",
    "    try:\n",
    "        whoami()\n",
    "        print('Authorization token already provided')\n",
    "    except OSError:\n",
    "        notebook_login()\n",
    "```\n",
    "* **qwen1.5-1.8b-chat/qwen1.5-7b-chat** - Qwen1.5 is the beta version of Qwen2, a transformer-based decoder-only language model pretrained on a large amount of data. Qwen1.5 is a language model series including decoder language models of different model sizes. It is based on the Transformer architecture with SwiGLU activation, attention QKV bias, group query attention, mixture of sliding window attention and full attention. You can find more details about model in the [model repository](https://huggingface.co/Qwen).\n",
    "* **qwen-7b-chat** - Qwen-7B is the 7B-parameter version of the large language model series, Qwen (abbr. Tongyi Qianwen), proposed by Alibaba Cloud. Qwen-7B is a Transformer-based large language model, which is pretrained on a large volume of data, including web texts, books, codes, etc. For more details about Qwen, please refer to the [GitHub](https://github.com/QwenLM/Qwen) code repository.\n",
    "* **mpt-7b-chat** - MPT-7B is part of the family of MosaicPretrainedTransformer (MPT) models, which use a modified transformer architecture optimized for efficient training and inference. These architectural changes include performance-optimized layer implementations and the elimination of context length limits by replacing positional embeddings with Attention with Linear Biases ([ALiBi](https://arxiv.org/abs/2108.12409)). Thanks to these modifications, MPT models can be trained with high throughput efficiency and stable convergence. MPT-7B-chat is a chatbot-like model for dialogue generation. It was built by finetuning MPT-7B on the [ShareGPT-Vicuna](https://huggingface.co/datasets/jeffwan/sharegpt_vicuna), [HC3](https://huggingface.co/datasets/Hello-SimpleAI/HC3), [Alpaca](https://huggingface.co/datasets/tatsu-lab/alpaca), [HH-RLHF](https://huggingface.co/datasets/Anthropic/hh-rlhf), and [Evol-Instruct](https://huggingface.co/datasets/victor123/evol_instruct_70k) datasets. More details about the model can be found in [blog post](https://www.mosaicml.com/blog/mpt-7b), [repository](https://github.com/mosaicml/llm-foundry/) and [HuggingFace model card](https://huggingface.co/mosaicml/mpt-7b-chat).\n",
    "* **chatglm3-6b** - ChatGLM3-6B is the latest open-source model in the ChatGLM series. While retaining many excellent features such as smooth dialogue and low deployment threshold from the previous two generations, ChatGLM3-6B employs a more diverse training dataset, more sufficient training steps, and a more reasonable training strategy. ChatGLM3-6B adopts a newly designed [Prompt format](https://github.com/THUDM/ChatGLM3/blob/main/PROMPT_en.md), in addition to the normal multi-turn dialogue. You can find more details about model in the [model card](https://huggingface.co/THUDM/chatglm3-6b)\n",
    "* **mistral-7b** - The Mistral-7B-v0.1 Large Language Model (LLM) is a pretrained generative text model with 7 billion parameters. You can find more details about model in the [model card](https://huggingface.co/mistralai/Mistral-7B-v0.1), [paper](https://arxiv.org/abs/2310.06825) and [release blog post](https://mistral.ai/news/announcing-mistral-7b/).\n",
    "* **zephyr-7b-beta** - Zephyr is a series of language models that are trained to act as helpful assistants. Zephyr-7B-beta is the second model in the series, and is a fine-tuned version of [mistralai/Mistral-7B-v0.1](https://huggingface.co/mistralai/Mistral-7B-v0.1) that was trained on on a mix of publicly available, synthetic datasets using [Direct Preference Optimization (DPO)](https://arxiv.org/abs/2305.18290). You can find more details about model in [technical report](https://arxiv.org/abs/2310.16944) and [HuggingFace model card](https://huggingface.co/HuggingFaceH4/zephyr-7b-beta).\n",
    "* **neural-chat-7b-v3-1** - Mistral-7b model fine-tuned using Intel Gaudi. The model fine-tuned on the open source dataset [Open-Orca/SlimOrca](https://huggingface.co/datasets/Open-Orca/SlimOrca) and aligned with [Direct Preference Optimization (DPO) algorithm](https://arxiv.org/abs/2305.18290). More details can be found in [model card](https://huggingface.co/Intel/neural-chat-7b-v3-1) and [blog post](https://medium.com/@NeuralCompressor/the-practice-of-supervised-finetuning-and-direct-preference-optimization-on-habana-gaudi2-a1197d8a3cd3).\n",
    "* **notus-7b-v1** - Notus is a collection of fine-tuned models using [Direct Preference Optimization (DPO)](https://arxiv.org/abs/2305.18290). and related [RLHF](https://huggingface.co/blog/rlhf) techniques. This model is the first version, fine-tuned with DPO over zephyr-7b-sft. Following a data-first approach, the only difference between Notus-7B-v1 and Zephyr-7B-beta is the preference dataset used for dDPO. Proposed approach for dataset creation helps to effectively fine-tune Notus-7b that surpasses Zephyr-7B-beta and Claude 2 on [AlpacaEval](https://tatsu-lab.github.io/alpaca_eval/). More details about model can be found in [model card](https://huggingface.co/argilla/notus-7b-v1).\n",
    "* **youri-7b-chat** - Youri-7b-chat is a Llama2 based model. [Rinna Co., Ltd.](https://rinna.co.jp/) conducted further pre-training for the Llama2 model with a mixture of English and Japanese datasets to improve Japanese task capability. The model is publicly released on Hugging Face hub. You can find detailed information at the [rinna/youri-7b-chat project page](https://huggingface.co/rinna/youri-7b). \n",
    "* **baichuan2-7b-chat** - Baichuan 2 is the new generation of large-scale open-source language models launched by [Baichuan Intelligence inc](https://www.baichuan-ai.com/home). It is trained on a high-quality corpus with 2.6 trillion tokens and has achieved the best performance in authoritative Chinese and English benchmarks of the same size."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "d3b57cfb-e727-43a5-b2c9-8f1b1ba72061",
   "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-03-07 23:06:47.788169: 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-03-07 23:06:47.791855: I tensorflow/tsl/cuda/cudart_stub.cc:28] Could not find cuda drivers on your machine, GPU will not be used.\n",
      "2024-03-07 23:06:47.834258: 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-03-07 23:06:47.834288: 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-03-07 23:06:47.834330: 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-03-07 23:06:47.842773: I tensorflow/tsl/cuda/cudart_stub.cc:28] Could not find cuda drivers on your machine, GPU will not be used.\n",
      "2024-03-07 23:06:47.844036: 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-03-07 23:06:48.759435: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT\n"
     ]
    }
   ],
   "source": [
    "from pathlib import Path\n",
    "from optimum.intel.openvino import OVModelForCausalLM\n",
    "import openvino as ov\n",
    "import torch\n",
    "import nncf\n",
    "import logging\n",
    "import shutil\n",
    "import gc\n",
    "import ipywidgets as widgets\n",
    "from transformers import (\n",
    "    AutoModelForCausalLM,\n",
    "    AutoModel,\n",
    "    AutoTokenizer,\n",
    "    AutoConfig,\n",
    "    TextIteratorStreamer,\n",
    "    pipeline,\n",
    "    StoppingCriteria,\n",
    "    StoppingCriteriaList,\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "b51ff5a9",
   "metadata": {},
   "source": [
    "## Convert model\n",
    "[back to top ⬆️](#Table-of-contents:)\n",
    "\n",
    "### Convert LLM model\n",
    "[back to top ⬆️](#Table-of-contents:)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "37bf49d7",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8f481cf6f5af459495323c305c9f2b14",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Dropdown(description='Model Language:', options=('English', 'Chinese', 'Japanese'), value='English')"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from config import SUPPORTED_EMBEDDING_MODELS, SUPPORTED_LLM_MODELS\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": 31,
   "id": "184d1678-0e73-4f35-8af5-1a7d291c2e6e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "815948450b6b441cb33407d6e888ec84",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Dropdown(description='Model:', index=2, options=('qwen1.5-1.8b-chat', 'qwen1.5-7b-chat', 'qwen-7b-chat', 'chat…"
      ]
     },
     "execution_count": 31,
     "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": 11,
   "id": "49ea95f8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Selected LLM model internlm2-chat-1.8b\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}\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "5d370f13",
   "metadata": {},
   "source": [
    "Optimum Intel can be used to load optimized models from the [Hugging Face Hub](https://huggingface.co/docs/optimum/intel/hf.co/models) and create pipelines to run an inference with OpenVINO Runtime using Hugging Face APIs. The Optimum Inference models are API compatible with Hugging Face Transformers models.  This means we just need to replace `AutoModelForXxx` class with the corresponding `OVModelForXxx` class.\n",
    "\n",
    "Below is an example of the RedPajama model\n",
    "\n",
    "```diff\n",
    "-from transformers import AutoModelForCausalLM\n",
    "+from optimum.intel.openvino import OVModelForCausalLM\n",
    "from transformers import AutoTokenizer, pipeline\n",
    "\n",
    "model_id = \"togethercomputer/RedPajama-INCITE-Chat-3B-v1\"\n",
    "-model = AutoModelForCausalLM.from_pretrained(model_id)\n",
    "+model = OVModelForCausalLM.from_pretrained(model_id, export=True)\n",
    "```\n",
    "\n",
    "Model class initialization starts with calling `from_pretrained` method. When downloading and converting Transformers model, the parameter `export=True` should be added. We can save the converted model for the next usage with the `save_pretrained` method.\n",
    "Tokenizer class and pipelines API are compatible with Optimum models.\n",
    "\n",
    "To optimize the generation process and use memory more efficiently, the `use_cache=True` option is enabled. Since the output side is auto-regressive, an output token hidden state remains the same once computed for every further generation step. Therefore, recomputing it every time you want to generate a new token seems wasteful. With the cache, the model saves the hidden state once it has been computed. The model only computes the one for the most recently generated output token at each time step, re-using the saved ones for hidden tokens. This reduces the generation complexity from $O(n^3)$ to $O(n^2)$ for a transformer model. More details about how it works can be found in this [article](https://scale.com/blog/pytorch-improvements#Text%20Translation). With this option, the model gets the previous step's hidden states (cached attention keys and values) as input and additionally provides hidden states for the current step as output. It means for all next iterations, it is enough to provide only a new token obtained from the previous step and cached key values to get the next token prediction. \n",
    "\n",
    "In our case, MPT, Qwen and ChatGLM model currently is not covered by Optimum Intel, we will convert it manually and create wrapper compatible with Optimum Intel."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "6337eab8",
   "metadata": {},
   "source": [
    "## Compress model weights\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.\n",
    "\n",
    "\n",
    "### Weights Compression using Optimum Intel\n",
    "[back to top ⬆️](#Table-of-contents:)\n",
    "\n",
    "Optimum Intel supports weight compression via NNCF out of the box. For 8-bit compression we pass `load_in_8bit=True` to `from_pretrained()` method of `OVModelForCausalLM`. For 4 bit compression we provide `quantization_config=OVWeightQuantizationConfig(bits=4, ...)` argument containing number of bits and other compression parameters. We will consider how to do it on RedPajama, LLAMA and Zephyr examples.\n",
    "\n",
    ">**Note**: There may be no speedup for INT4/INT8 compressed models on dGPU.\n",
    "\n",
    "### Weights Compression using NNCF\n",
    "[back to top ⬆️](#Table-of-contents:)\n",
    "\n",
    "You also can perform weights compression for OpenVINO models using NNCF directly. `nncf.compress_weights` function accepts OpenVINO model instance and compresses its weights for Linear and Embedding layers. We will consider this variant based on MPT model.\n",
    "\n",
    "\n",
    ">**Note**: This tutorial involves conversion model for FP16 and INT4/INT8 weights compression scenarios. It may be memory and time-consuming in the first run. You can manually control the compression precision below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "c6a38153",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4d6fd33052574569bb607d72148933d2",
       "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": "52cae9ef3fca42029949c79ce046b52e",
       "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": "2dfd7e4fd16f4e0785929d8f3fceb088",
       "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 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": 14,
   "id": "2020d522",
   "metadata": {},
   "outputs": [],
   "source": [
    "from optimum.intel import OVWeightQuantizationConfig\n",
    "from converter import converters, register_configs\n",
    "\n",
    "register_configs()\n",
    "\n",
    "nncf.set_log_level(logging.ERROR)\n",
    "\n",
    "pt_model_id = llm_model_configuration[\"model_id\"]\n",
    "pt_model_name = llm_model_id.value.split(\"-\")[0]\n",
    "model_type = AutoConfig.from_pretrained(pt_model_id, trust_remote_code=True).model_type\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",
    "    if not llm_model_configuration[\"remote\"]:\n",
    "        remote_code = llm_model_configuration.get(\"remote_code\", False)\n",
    "        model_kwargs = {}\n",
    "        if remote_code:\n",
    "            model_kwargs = {\n",
    "                \"trust_remote_code\": True,\n",
    "                \"config\": AutoConfig.from_pretrained(pt_model_id, trust_remote_code=True)\n",
    "            }\n",
    "        ov_model = OVModelForCausalLM.from_pretrained(\n",
    "            pt_model_id, export=True, compile=False, load_in_8bit=False, **model_kwargs\n",
    "        )\n",
    "        ov_model.half()\n",
    "        ov_model.save_pretrained(fp16_model_dir)\n",
    "        del ov_model\n",
    "    else:\n",
    "        model_kwargs = {}\n",
    "        if \"revision\" in llm_model_configuration:\n",
    "            model_kwargs[\"revision\"] = llm_model_configuration[\"revision\"]\n",
    "        model = AutoModelForCausalLM.from_pretrained(\n",
    "            llm_model_configuration[\"model_id\"],\n",
    "            torch_dtype=torch.float32,\n",
    "            trust_remote_code=True,\n",
    "            **model_kwargs\n",
    "        )\n",
    "        converters[pt_model_name](model, fp16_model_dir)\n",
    "        del model\n",
    "    gc.collect()\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",
    "    if not llm_model_configuration[\"remote\"]:\n",
    "        remote_code = llm_model_configuration.get(\"remote_code\", False)\n",
    "        model_kwargs = {}\n",
    "        if remote_code:\n",
    "            model_kwargs = {\n",
    "                \"trust_remote_code\": True,\n",
    "                \"config\": AutoConfig.from_pretrained(pt_model_id, trust_remote_code=True)\n",
    "            }\n",
    "        ov_model = OVModelForCausalLM.from_pretrained(\n",
    "            pt_model_id, export=True, compile=False, load_in_8bit=True, **model_kwargs\n",
    "        )\n",
    "        ov_model.save_pretrained(int8_model_dir)\n",
    "        del ov_model\n",
    "    else:\n",
    "        convert_to_fp16()\n",
    "        ov_model = ov.Core().read_model(fp16_model_dir / \"openvino_model.xml\")\n",
    "        shutil.copy(fp16_model_dir / \"config.json\", int8_model_dir / \"config.json\")\n",
    "        configuration_file = fp16_model_dir / f\"configuration_{model_type}.py\"\n",
    "        if configuration_file.exists():\n",
    "            shutil.copy(\n",
    "                configuration_file, int8_model_dir / f\"configuration_{model_type}.py\"\n",
    "            )\n",
    "        compressed_model = nncf.compress_weights(ov_model)\n",
    "        ov.save_model(compressed_model, int8_model_dir / \"openvino_model.xml\")\n",
    "        del ov_model\n",
    "        del compressed_model\n",
    "    gc.collect()\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",
    "        \"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",
    "        \"chatglm2-6b\": {\n",
    "            \"sym\": True,\n",
    "            \"group_size\": 128,\n",
    "            \"ratio\": 0.72,\n",
    "        },\n",
    "        \"qwen-7b-chat\": {\n",
    "            \"sym\": True, \n",
    "            \"group_size\": 128, \n",
    "            \"ratio\": 0.6\n",
    "        },\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(\n",
    "        llm_model_id.value, compression_configs[\"default\"]\n",
    "    )\n",
    "    if (int4_model_dir / \"openvino_model.xml\").exists():\n",
    "        return\n",
    "    int4_model_dir.mkdir(parents=True, exist_ok=True)\n",
    "    if not llm_model_configuration[\"remote\"]:\n",
    "        remote_code = llm_model_configuration.get(\"remote_code\", False)\n",
    "        model_kwargs = {}\n",
    "        if remote_code:\n",
    "            model_kwargs = {\n",
    "                \"trust_remote_code\" : True,\n",
    "                \"config\": AutoConfig.from_pretrained(pt_model_id, trust_remote_code=True)\n",
    "            }\n",
    "        ov_model = OVModelForCausalLM.from_pretrained(\n",
    "            pt_model_id, export=True, compile=False,\n",
    "            quantization_config=OVWeightQuantizationConfig(bits=4, **model_compression_params),\n",
    "            **model_kwargs\n",
    "        )\n",
    "        ov_model.save_pretrained(int4_model_dir)\n",
    "        del ov_model\n",
    "\n",
    "    else:\n",
    "        convert_to_fp16()\n",
    "        ov_model = ov.Core().read_model(fp16_model_dir / \"openvino_model.xml\")\n",
    "        shutil.copy(fp16_model_dir / \"config.json\", int4_model_dir / \"config.json\")\n",
    "        configuration_file = fp16_model_dir / f\"configuration_{model_type}.py\"\n",
    "        if configuration_file.exists():\n",
    "            shutil.copy(\n",
    "                configuration_file, int4_model_dir / f\"configuration_{model_type}.py\"\n",
    "            )\n",
    "        mode = nncf.CompressWeightsMode.INT4_SYM if model_compression_params[\"sym\"] else \\\n",
    "            nncf.CompressWeightsMode.INT4_ASYM\n",
    "        del model_compression_params[\"sym\"]\n",
    "        compressed_model = nncf.compress_weights(ov_model, mode=mode, **model_compression_params)\n",
    "        ov.save_model(compressed_model, int4_model_dir / \"openvino_model.xml\")\n",
    "        del ov_model\n",
    "        del compressed_model\n",
    "    gc.collect()\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()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "f1d1d1a2",
   "metadata": {},
   "source": [
    "Let's compare model size for different compression types"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "8e127215",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Size of FP16 model is 1819.91 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(\n",
    "            f\"Size of model with INT{precision} compressed weights is {compressed_weights.stat().st_size / 1024 / 1024:.2f} MB\"\n",
    "        )\n",
    "    if compressed_weights.exists() and fp16_weights.exists():\n",
    "        print(\n",
    "            f\"Compression rate for INT{precision} model: {fp16_weights.stat().st_size / compressed_weights.stat().st_size:.3f}\"\n",
    "        )"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "4f943465",
   "metadata": {},
   "source": [
    "### Convert embedding model\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."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "ff80e6eb-7923-40ef-93d8-5e6c56e50667",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "10bc2ae1b14d4cb69f0db31fe6133643",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Dropdown(description='Embedding Model:', options=('all-mpnet-base-v2', 'text2vec-large-chinese'), value='all-m…"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "embedding_model_id = list(SUPPORTED_EMBEDDING_MODELS)\n",
    "\n",
    "if model_language.value != \"Chinese\":\n",
    "    embedding_model_id = [x for x in embedding_model_id if \"chinese\" not in x]\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": 17,
   "id": "790afcf8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Selected all-mpnet-base-v2 model\n"
     ]
    }
   ],
   "source": [
    "embedding_model_configuration = SUPPORTED_EMBEDDING_MODELS[embedding_model_id.value]\n",
    "print(f\"Selected {embedding_model_id.value} model\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "58d75dad-2eeb-4edd-8d12-d77a365f8eda",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "embedding_model_dir = Path(embedding_model_id.value)\n",
    "\n",
    "if not (embedding_model_dir / \"openvino_model.xml\").exists():\n",
    "    model = AutoModel.from_pretrained(embedding_model_configuration[\"model_id\"])\n",
    "    converters[embedding_model_id.value](model, embedding_model_dir)\n",
    "    tokenizer = AutoTokenizer.from_pretrained(embedding_model_configuration[\"model_id\"])\n",
    "    tokenizer.save_pretrained(embedding_model_dir)\n",
    "    del model"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "749b5bbd",
   "metadata": {},
   "source": [
    "## Select device for inference and model variant\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",
    "[back to top ⬆️](#Table-of-contents:)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "e11e73cf",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d92f5255fdad4d3cbd056e52e99f5ee9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Dropdown(description='Device:', options=('CPU', 'AUTO'), value='CPU')"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "core = ov.Core()\n",
    "embedding_device = widgets.Dropdown(\n",
    "    options=core.available_devices + [\"AUTO\"],\n",
    "    value=\"CPU\",\n",
    "    description=\"Device:\",\n",
    "    disabled=False,\n",
    ")\n",
    "\n",
    "embedding_device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "9ab29b85",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Embedding model will be loaded to CPU device for response generation\n"
     ]
    }
   ],
   "source": [
    "print(f\"Embedding model will be loaded to {embedding_device.value} device for response generation\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "ef31656a",
   "metadata": {},
   "source": [
    "### Select device for LLM model inference\n",
    "[back to top ⬆️](#Table-of-contents:)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "6d044d01",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "21accfb73c1a45b19b331ef2159dcd81",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Dropdown(description='Device:', options=('CPU', 'AUTO'), value='CPU')"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "llm_device = widgets.Dropdown(\n",
    "    options=core.available_devices + [\"AUTO\"],\n",
    "    value=\"CPU\",\n",
    "    description=\"Device:\",\n",
    "    disabled=False,\n",
    ")\n",
    "\n",
    "llm_device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "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\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "bc225391",
   "metadata": {},
   "source": [
    "## Load model\n",
    "[back to top ⬆️](#Table-of-contents:)\n",
    "\n",
    "### Load embedding model\n",
    "[back to top ⬆️](#Table-of-contents:)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "e85aeffa",
   "metadata": {},
   "source": [
    "Wrapper around a text embedding model for LangChain, used for converting text to embeddings."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "df3e8fd1-d4c1-4e33-b46e-7840e392f8ee",
   "metadata": {},
   "outputs": [],
   "source": [
    "from ov_embedding_model import OVEmbeddings\n",
    "\n",
    "embedding = OVEmbeddings.from_model_id(\n",
    "    embedding_model_dir,\n",
    "    do_norm=embedding_model_configuration[\"do_norm\"],\n",
    "    ov_config={\n",
    "        \"device_name\": embedding_device.value,\n",
    "        \"config\": {\"PERFORMANCE_HINT\": \"THROUGHPUT\"},\n",
    "    },\n",
    "    model_kwargs={\n",
    "        \"model_max_length\": 512,\n",
    "    },\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "79fe990a",
   "metadata": {},
   "source": [
    "### Load LLM model\n",
    "[back to top ⬆️](#Table-of-contents:)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "e2610f4b",
   "metadata": {},
   "source": [
    "The cell below create `OVMPTModel` and `OVCHATGLM2Model` wrapper based on `OVModelForCausalLM` model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "efe29701",
   "metadata": {},
   "outputs": [],
   "source": [
    "from ov_llm_model import model_classes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "8b014f24-aa5b-4d40-924d-d579ad7fcec6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3138dedc2102456db478482b1260c704",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Dropdown(description='Model to run:', options=('FP16',), value='FP16')"
      ]
     },
     "execution_count": 25,
     "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": "code",
   "execution_count": 26,
   "id": "f7f708db-8de1-4efd-94b2-fcabc48d52f4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading model from internlm2-chat-1.8b/FP16\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "The argument `trust_remote_code` is to be used along with export=True. It will be ignored.\n",
      "Compiling the model to CPU ...\n"
     ]
    }
   ],
   "source": [
    "from langchain.llms 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",
    "model_name = llm_model_configuration[\"model_id\"]\n",
    "stop_tokens = llm_model_configuration.get(\"stop_tokens\")\n",
    "class_key = llm_model_id.value.split(\"-\")[0]\n",
    "tok = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)\n",
    "\n",
    "class StopOnTokens(StoppingCriteria):\n",
    "    def __init__(self, token_ids):\n",
    "        self.token_ids = token_ids\n",
    "\n",
    "    def __call__(\n",
    "        self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs\n",
    "    ) -> 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",
    "if stop_tokens is not None:\n",
    "    if isinstance(stop_tokens[0], str):\n",
    "        stop_tokens = tok.convert_tokens_to_ids(stop_tokens)\n",
    "\n",
    "    stop_tokens = [StopOnTokens(stop_tokens)]\n",
    "\n",
    "model_class = (\n",
    "    OVModelForCausalLM\n",
    "    if not llm_model_configuration[\"remote\"]\n",
    "    else model_classes[class_key]\n",
    ")\n",
    "ov_model = model_class.from_pretrained(\n",
    "    model_dir,\n",
    "    device=llm_device.value,\n",
    "    ov_config=ov_config,\n",
    "    config=AutoConfig.from_pretrained(model_dir, trust_remote_code=True),\n",
    "    trust_remote_code=True,\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "4fb8b0e4",
   "metadata": {},
   "source": [
    "## Run QA over Document\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",
    "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",
    "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": 27,
   "id": "5b97eeeb",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "from langchain.text_splitter import CharacterTextSplitter, RecursiveCharacterTextSplitter, MarkdownTextSplitter\n",
    "from langchain.document_loaders import (\n",
    "    CSVLoader,\n",
    "    EverNoteLoader,\n",
    "    PDFMinerLoader,\n",
    "    TextLoader,\n",
    "    UnstructuredEPubLoader,\n",
    "    UnstructuredHTMLLoader,\n",
    "    UnstructuredMarkdownLoader,\n",
    "    UnstructuredODTLoader,\n",
    "    UnstructuredPowerPointLoader,\n",
    "    UnstructuredWordDocumentLoader, )\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(\n",
    "            '([﹒﹔﹖﹗．。！？][\"’”」』]{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\": (PDFMinerLoader, {}),\n",
    "    \".ppt\": (UnstructuredPowerPointLoader, {}),\n",
    "    \".pptx\": (UnstructuredPowerPointLoader, {}),\n",
    "    \".txt\": (TextLoader, {\"encoding\": \"utf8\"}),\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "602f8ebd-789c-4eb2-b54d-b23d8f1d8e7b",
   "metadata": {},
   "source": [
    "To reuse RAG pipeline of LangChain, we can wrapper around a LLM/chat model for LangChain, used for generating the response text. An OpenVINO compiled model can be run locally through the `HuggingFacePipeline` class."
   ]
  },
  {
   "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 import RetrievalQA\n",
    "from langchain.docstore.document import Document\n",
    "from threading import Event, Thread\n",
    "import gradio as gr\n",
    "import re\n",
    "from uuid import uuid4\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(\n",
    "    \"partial_text_processor\", default_partial_text_processor\n",
    ")\n",
    "\n",
    "\n",
    "def build_retriever(docs, spliter_name, chunk_size, chunk_overlap, vector_search_top_k):\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](\n",
    "        chunk_size=chunk_size, chunk_overlap=chunk_overlap\n",
    "    )\n",
    "\n",
    "    texts = text_splitter.split_documents(documents)\n",
    "\n",
    "    db = Chroma.from_documents(texts, embedding)\n",
    "\n",
    "    global retriever\n",
    "    retriever = db.as_retriever(search_kwargs={\"k\": vector_search_top_k})\n",
    "\n",
    "    return \"Retriever is Ready\"\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, conversation_id):\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",
    "      conversation_id: unique conversation identifier.\n",
    "\n",
    "    \"\"\"\n",
    "    streamer = TextIteratorStreamer(\n",
    "        tok, timeout=60.0, skip_prompt=True, skip_special_tokens=True\n",
    "    )\n",
    "    generate_kwargs = dict(\n",
    "        model=ov_model,\n",
    "        tokenizer=tok,\n",
    "        max_new_tokens=256,\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",
    "        generate_kwargs[\"stopping_criteria\"] = StoppingCriteriaList(stop_tokens)\n",
    "        \n",
    "    pipe = pipeline(\"text-generation\", **generate_kwargs)\n",
    "    llm = HuggingFacePipeline(pipeline=pipe)\n",
    "    \n",
    "    prompt = PromptTemplate.from_template(llm_model_configuration[\"rag_prompt_template\"])\n",
    "    chain_type_kwargs = {\"prompt\": prompt}\n",
    "    rag_chain = RetrievalQA.from_chain_type(\n",
    "        llm=llm,\n",
    "        chain_type=\"stuff\",\n",
    "        retriever=retriever,\n",
    "        chain_type_kwargs=chain_type_kwargs,\n",
    "    )\n",
    "    \n",
    "    stream_complete = Event()\n",
    "\n",
    "    def infer(question):\n",
    "        rag_chain.invoke(question)\n",
    "        stream_complete.set()\n",
    "\n",
    "    t1 = Thread(target=infer, args=(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",
    "    ov_model.request.cancel()\n",
    "\n",
    "\n",
    "def get_uuid():\n",
    "    \"\"\"\n",
    "    universal unique identifier for thread\n",
    "    \"\"\"\n",
    "    return str(uuid4())\n",
    "\n",
    "\n",
    "with gr.Blocks(\n",
    "    theme=gr.themes.Soft(),\n",
    "    css=\".disclaimer {font-variant-caps: all-small-caps;}\",\n",
    ") as demo:\n",
    "    conversation_id = gr.State(get_uuid)\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=\"Load text files\",\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(\"Build Retriever\")\n",
    "            retriever_argument = gr.Accordion(\"Retriever Configuration\", open=False)\n",
    "            with retriever_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=1000,\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=200,\n",
    "                    minimum=0,\n",
    "                    maximum=400,\n",
    "                    step=10,\n",
    "                    interactive=True,\n",
    "                    info=(\"Overlap between 2 chunks\"),\n",
    "                )\n",
    "\n",
    "                vector_search_top_k = gr.Slider(\n",
    "                    1,\n",
    "                    10,\n",
    "                    value=4,\n",
    "                    step=1,\n",
    "                    label=\"Vector search top k\",\n",
    "                    interactive=True,\n",
    "                )\n",
    "            langchain_status = gr.Textbox(\n",
    "                label=\"Status\", value=\"Retriever is Not ready\", 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(height=600)\n",
    "            with gr.Row():\n",
    "                with gr.Column():\n",
    "                    msg = gr.Textbox(\n",
    "                        label=\"Chat 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",
    "    load_docs.click(\n",
    "        build_retriever,\n",
    "        inputs=[docs, spliter, chunk_size, chunk_overlap, vector_search_top_k],\n",
    "        outputs=[langchain_status],\n",
    "        queue=False,\n",
    "    )\n",
    "    submit_event = msg.submit(\n",
    "        user, [msg, chatbot], [msg, chatbot], queue=False\n",
    "    ).then(bot, [chatbot, temperature, top_p, top_k, repetition_penalty, conversation_id], chatbot, queue=True)\n",
    "    submit_click_event = submit.click(\n",
    "        user, [msg, chatbot], [msg, chatbot], queue=False\n",
    "    ).then(bot, [chatbot, temperature, top_p, top_k, repetition_penalty, conversation_id], chatbot, queue=True)\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",
    "\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": 29,
   "id": "6f4b5a84-bebf-49b9-b2fa-5e788ed2cbac",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Closing server running on port: 4545\n"
     ]
    }
   ],
   "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": [],
    "tasks": [
     "Text Generation"
    ]
   }
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "state": {},
    "version_major": 2,
    "version_minor": 0
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}