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

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "50e1d1d5-3bdd-4224-9f93-bf5d9a83f424",
   "metadata": {},
   "source": [
    "# Create an LLM-powered Chatbot using OpenVINO\n",
    "\n",
    "In the rapidly evolving world of artificial intelligence (AI), chatbots have emerged as powerful tools for businesses to enhance customer interactions and streamline operations. \n",
    "Large Language Models (LLMs) are artificial intelligence systems that can understand and generate human language. They use deep learning algorithms and massive amounts of data to learn the nuances of language and produce coherent and relevant responses.\n",
    "While a decent intent-based chatbot can answer basic, one-touch inquiries like order management, FAQs, and policy questions, LLM chatbots can tackle more complex, multi-touch questions. LLM enables chatbots to provide support in a conversational manner, similar to how humans do, through contextual memory. Leveraging the capabilities of Language Models, chatbots are becoming increasingly intelligent, capable of understanding and responding to human language with remarkable accuracy.\n",
    "\n",
    "Previously, we already discussed how to build an instruction-following pipeline using OpenVINO and Optimum Intel, please check out [Dolly example](../240-dolly-2-instruction-following) for reference.\n",
    "In this tutorial, we consider how to use the power of OpenVINO for running Large Language Models for chat. We will use a pre-trained model from the [Hugging Face Transformers](https://huggingface.co/docs/transformers/index) library. To simplify the user experience, the [Hugging Face Optimum Intel](https://huggingface.co/docs/optimum/intel/index) library is used to convert the models to OpenVINO™ IR format.\n",
    "\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 chat inference pipeline\n",
    "- Run chat pipeline\n",
    "\n",
    "#### Table of contents:\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",
    "- [Instantiate Model using Optimum Intel](#Instantiate-Model-using-Optimum-Intel)\n",
    "- [Compress model weights $\\uparrow$(#Table-of-content:)](#Compress-model-weights-\\uparrow(#Table-of-content:))\n",
    "    - [Weights Compression using Optimum Intel $\\uparrow$(#Table-of-content:)](#Weights-Compression-using-Optimum-Intel-\\uparrow(#Table-of-content:))\n",
    "    - [Weights Compression using NNCF $\\uparrow$(#Table-of-content:)](#Weights-Compression-using-NNCF-\\uparrow(#Table-of-content:))\n",
    "- [Select device for inference and model variant $\\uparrow$(#Table-of-content:)](#Select-device-for-inference-and-model-variant-\\uparrow(#Table-of-content:))\n",
    "- [Run Chatbot](#Run-Chatbot)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5df233b0-0369-4fff-9952-7957a90394a5",
   "metadata": {},
   "source": [
    "## Prerequisites\n",
    "[back to top ⬆️](#Table-of-contents:)\n",
    "\n",
    "Install required dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "563ecf9f-346b-4f14-85ef-c66ff0c95f65",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%pip uninstall -q -y openvino-dev openvino openvino-nightly\n",
    "%pip install -q openvino-nightly\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@release_v270\"\\\n",
    "\"gradio\"\\\n",
    "\"onnx\" \"einops\" \"transformers>=4.33.0\"\\"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "81983176-e571-4652-ba21-4bd608c35146",
   "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 options are:\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",
    "* **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",
    "[back to top ⬆️](#Table-of-contents:)\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",
    "* **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)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "f93282b6-f1f1-4153-84af-31aac79c3ef4",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from config import SUPPORTED_MODELS\n",
    "import ipywidgets as widgets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "8d22fedb-d1f6-4306-b910-efac5b849c7c",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "01dc75bdf73d43bfbc5b878cc3efcdfa",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Dropdown(description='Model:', index=2, options=('red-pajama-3b-chat', 'llama-2-chat-7b', 'mpt-7b-chat'), valu…"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model_ids = list(SUPPORTED_MODELS)\n",
    "\n",
    "model_id = widgets.Dropdown(\n",
    "    options=model_ids,\n",
    "    value=model_ids[2],\n",
    "    description='Model:',\n",
    "    disabled=False,\n",
    ")\n",
    "\n",
    "model_id"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "906022ec-96bf-41a9-9447-789d2e875250",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Selected model mpt-7b-chat\n"
     ]
    }
   ],
   "source": [
    "model_configuration = SUPPORTED_MODELS[model_id.value]\n",
    "print(f\"Selected model {model_id.value}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "62af3e8a-915a-49b4-8007-803777ba9eaf",
   "metadata": {},
   "source": [
    "## Instantiate Model using Optimum Intel\n",
    "[back to top ⬆️](#Table-of-contents:)\n",
    "\n",
    "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 model currently is not covered by Optimum Intel, we will convert it manually and create wrapper compatible with Optimum Intel.\n",
    "\n",
    "Below is some code required for MPT conversion."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "d36e20d2-70c3-43d0-b0e1-cb2d1e06bdaa",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:nncf:NNCF initialized successfully. Supported frameworks detected: torch, onnx, openvino\n"
     ]
    }
   ],
   "source": [
    "from functools import wraps\n",
    "import torch\n",
    "from transformers import AutoModelForCausalLM\n",
    "from nncf import compress_weights\n",
    "import openvino as ov\n",
    "from pathlib import Path\n",
    "from typing import Optional, Union, Dict, Tuple, List\n",
    "\n",
    "def flattenize_inputs(inputs):\n",
    "    \"\"\"\n",
    "    Helper function for making nested inputs flattens\n",
    "    \"\"\"\n",
    "    flatten_inputs = []\n",
    "    for input_data in inputs:\n",
    "        if input_data is None:\n",
    "            continue\n",
    "        if isinstance(input_data, (list, tuple)):\n",
    "            flatten_inputs.extend(flattenize_inputs(input_data))\n",
    "        else:\n",
    "            flatten_inputs.append(input_data)\n",
    "    return flatten_inputs\n",
    "\n",
    "def cleanup_torchscript_cache():\n",
    "    \"\"\"\n",
    "    Helper for removing cached model representation\n",
    "    \"\"\"\n",
    "    torch._C._jit_clear_class_registry()\n",
    "    torch.jit._recursive.concrete_type_store = torch.jit._recursive.ConcreteTypeStore()\n",
    "    torch.jit._state._clear_class_state()\n",
    "\n",
    "def convert_mpt(pt_model:torch.nn.Module, model_path:Path):\n",
    "    \"\"\"\n",
    "    MPT model conversion function\n",
    "    \n",
    "    Params:\n",
    "      pt_model: PyTorch model\n",
    "      model_path: path for saving model\n",
    "    Returns:\n",
    "      None\n",
    "    \"\"\"\n",
    "    ov_out_path = Path(model_path) / \"openvino_model.xml\"\n",
    "    pt_model.config.save_pretrained(ov_out_path.parent)\n",
    "    pt_model.config.use_cache = True\n",
    "    outs = pt_model(input_ids=torch.ones((1, 10), dtype=torch.long), attention_mask=torch.ones((1, 10), dtype=torch.long))\n",
    "    inputs = [\"input_ids\"]\n",
    "    outputs = [\"logits\"]\n",
    "\n",
    "    dynamic_shapes = {\"input_ids\": {1: \"seq_len\"}, \"attention_mask\": {1: \"seq_len\"}}\n",
    "    for idx in range(len(outs.past_key_values)):\n",
    "        inputs.extend([f\"past_key_values.{idx}.key\", f\"past_key_values.{idx}.value\"])\n",
    "        dynamic_shapes[inputs[-1]] = {2: \"past_sequence + sequence\"}\n",
    "        dynamic_shapes[inputs[-2]] = {3: \"past_sequence + sequence\"}\n",
    "        outputs.extend([f\"present.{idx}.key\", f\"present.{idx}.value\"])\n",
    "            \n",
    "    inputs.append(\"attention_mask\")\n",
    "    dummy_inputs = {\"input_ids\": torch.ones((1,2), dtype=torch.long), \"past_key_values\": outs.past_key_values, \"attention_mask\": torch.ones((1,12), dtype=torch.long)}\n",
    "    pt_model.config.torchscript = True\n",
    "    orig_forward = pt_model.forward\n",
    "    @wraps(orig_forward)\n",
    "    def ts_patched_forward(input_ids: torch.Tensor, past_key_values: Tuple[Tuple[torch.Tensor]], attention_mask: torch.Tensor):\n",
    "        pkv_list = list(past_key_values)\n",
    "        outs = orig_forward(input_ids=input_ids, past_key_values=pkv_list, attention_mask=attention_mask)\n",
    "        return (outs.logits, tuple(outs.past_key_values))\n",
    "    pt_model.forward = ts_patched_forward\n",
    "    ov_model = ov.convert_model(pt_model, example_input=dummy_inputs)\n",
    "    pt_model.forward = orig_forward\n",
    "    for inp_name, m_input, input_data in zip(inputs, ov_model.inputs, flattenize_inputs(dummy_inputs.values())):\n",
    "        input_node = m_input.get_node()\n",
    "        if input_node.element_type == ov.Type.dynamic:\n",
    "            m_input.get_node().set_element_type(ov.Type.f32)\n",
    "        shape = list(input_data.shape)\n",
    "        if inp_name in dynamic_shapes:\n",
    "            for k in dynamic_shapes[inp_name]:\n",
    "                shape[k] = -1\n",
    "        input_node.set_partial_shape(ov.PartialShape(shape))\n",
    "        m_input.get_tensor().set_names({inp_name})\n",
    "        \n",
    "    for out, out_name in zip(ov_model.outputs, outputs):\n",
    "        out.get_tensor().set_names({out_name})\n",
    "\n",
    "    ov_model.validate_nodes_and_infer_types()\n",
    "    ov.save_model(ov_model, ov_out_path)\n",
    "    del ov_model\n",
    "    cleanup_torchscript_cache()\n",
    "    del pt_model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "13694bf8-ee7b-4186-a3e0-a8705be9733c",
   "metadata": {},
   "source": [
    "## Compress model weights [$\\uparrow$](#Table-of-content:)\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 [$\\uparrow$](#Table-of-content:)\n",
    "[back to top ⬆️](#Table-of-contents:)\n",
    "",
    "\n",
    "To enable weights compression via NNCF for models supported by Optimum Intel `OVQuantizer` class should be used instantiated by PyTorch model using `from_pretrained` method. `OVQuantizer.quantize(save_directory=save_dir, weights_only=True)` enables weights compression and model conversion to OpenVINO Intermediate Representation format. We will consider how to do it on RedPajama and LLAMA examples. \n",
    "\n",
    ">**Note**: Weights Compression using Optimum Intel currently supports only INT8 compression. We will apply INT4 compression for these model using NNCF API described below.\n",
    "\n",
    ">**Note**: There may be no speedup for INT4/INT8 compressed models on dGPU.\n",
    "\n",
    "### Weights Compression using NNCF [$\\uparrow$](#Table-of-content:)\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 maybe memory and time-consuming in the first run. You can manually control compression precision below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "91eb2ccf",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2a1838d4ac0e482f8cfddeab8a522ef4",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Dropdown(description='Compression mode:', options=('INT4', 'INT8', 'FP16'), value='INT4')"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from IPython.display import display\n",
    "\n",
    "# TODO: red-pajama-3b-chat currently can't be compiled in INT4 or FP16 due to ticket 123973\n",
    "is_pajama_model = model_id.value == 'red-pajama-3b-chat'\n",
    "prepare_int4_model = widgets.Checkbox(\n",
    "    value=True and not is_pajama_model,\n",
    "    description='Prepare INT4 model',\n",
    "    disabled=is_pajama_model,\n",
    ")\n",
    "prepare_int8_model = widgets.Checkbox(\n",
    "    value=False or is_pajama_model,\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=is_pajama_model,\n",
    ")\n",
    "\n",
    "display(prepare_int4_model)\n",
    "display(prepare_int8_model)\n",
    "display(prepare_fp16_model)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "130a037a-7d98-4152-81ea-92ffb01da5a2",
   "metadata": {},
   "source": [
    "We can now save floating point and compressed model variants"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "c4ef9112",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "No CUDA runtime is found, using CUDA_HOME='/usr/local/cuda-11.7'\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0d0e425ace2d40c3809a597fc5cd7fed",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Output()"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"></pre>\n"
      ],
      "text/plain": []
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4ec5e866626348c38275478e0e8227eb",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Output()"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"></pre>\n"
      ],
      "text/plain": []
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from pathlib import Path\n",
    "from optimum.intel import OVQuantizer\n",
    "from optimum.intel.openvino import OVModelForCausalLM\n",
    "import shutil\n",
    "import logging\n",
    "import nncf\n",
    "import gc\n",
    "\n",
    "nncf.set_log_level(logging.ERROR)\n",
    "\n",
    "pt_model_id = model_configuration[\"model_id\"]\n",
    "fp16_model_dir = Path(model_id.value) / \"FP16\"\n",
    "int8_model_dir = Path(model_id.value) / \"INT8_compressed_weights\"\n",
    "int4_model_dir = Path(model_id.value) / \"INT4_compressed_weights\"\n",
    "\n",
    "def convert_to_fp16():\n",
    "    if fp16_model_dir.exists():\n",
    "        return\n",
    "    if \"mpt\" not in model_id.value:\n",
    "        ov_model = OVModelForCausalLM.from_pretrained(pt_model_id, export=True, compile=False)\n",
    "        ov_model.half()\n",
    "        ov_model.save_pretrained(fp16_model_dir)\n",
    "        del ov_model\n",
    "    else:\n",
    "        model = AutoModelForCausalLM.from_pretrained(model_configuration[\"model_id\"], torch_dtype=torch.float32, trust_remote_code=True)\n",
    "        convert_mpt(model, fp16_model_dir)\n",
    "        del model\n",
    "    gc.collect()\n",
    "\n",
    "def convert_to_int8():\n",
    "    if int8_model_dir.exists():\n",
    "        return\n",
    "    if \"mpt\" not in model_id.value:\n",
    "        pt_model = AutoModelForCausalLM.from_pretrained(pt_model_id)\n",
    "        quantizer = OVQuantizer.from_pretrained(pt_model)\n",
    "        quantizer.quantize(save_directory=int8_model_dir, weights_only=True)\n",
    "        del quantizer\n",
    "        del pt_model\n",
    "    else:\n",
    "        model = AutoModelForCausalLM.from_pretrained(model_configuration[\"model_id\"], torch_dtype=torch.float32, trust_remote_code=True)\n",
    "        compressed_model = compress_weights(model)\n",
    "        convert_mpt(compressed_model, int8_model_dir)\n",
    "        del model\n",
    "        del compressed_model\n",
    "    gc.collect()\n",
    "\n",
    "\n",
    "def convert_to_int4(group_size, ratio):\n",
    "    if int4_model_dir.exists():\n",
    "        return\n",
    "    int4_model_dir.mkdir(parents=True)\n",
    "    if \"mpt\" not in model_id.value:\n",
    "        # TODO: remove compression via NNCF for non-MPT models when INT4 weight compression is added to optimum-intel\n",
    "        model = OVModelForCausalLM.from_pretrained(pt_model_id, export=True, compile=False)\n",
    "        model.config.save_pretrained(int4_model_dir)\n",
    "        ov_model = model.half().model\n",
    "        del 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', int4_model_dir / 'config.json')\n",
    "    compressed_model = nncf.compress_weights(ov_model, mode=nncf.CompressWeightsMode.INT4_ASYM, group_size=group_size, ratio=ratio)\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",
    "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(group_size=128, ratio=0.8)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "671a17d4",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "Let's compare model size for different compression types"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "281f1d07-998e-4e13-ba95-0264564ede82",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Size of FP16 model is 12682.63 MB\n",
      "Size of model with INT8 compressed weights is 6542.01 MB\n",
      "Compression rate for INT8 model: 1.939\n",
      "Size of model with INT4 compressed weights is 4036.34 MB\n",
      "Compression rate for INT4 model: 3.142\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": "6d62f9f4-5434-4550-b372-c86b5a5089d5",
   "metadata": {},
   "source": [
    "## Select device for inference and model variant [$\\uparrow$](#Table-of-content:)\n",
    "[back to top ⬆️](#Table-of-contents:)\n",
    "",
    "\n",
    ">**Note**: There may be no speedup for INT4/INT8 compressed models on dGPU."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "837b4a3b-ccc3-4004-9577-2b2c7b802dea",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "core = ov.Core()\n",
    "device = widgets.Dropdown(\n",
    "    options=core.available_devices + [\"AUTO\"],\n",
    "    value='CPU',\n",
    "    description='Device:',\n",
    "    disabled=False,\n",
    ")\n",
    "\n",
    "device"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bd55ade7-0445-47e1-90f0-72b82da016ca",
   "metadata": {},
   "source": [
    "The cell below create `OVMPTModel` model wrapper based on `OVModelForCausalLM` model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "5333ab9b-ff5d-4a7f-bcdc-9cca5d56dc0a",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from transformers import AutoConfig, PretrainedConfig\n",
    "import torch\n",
    "\n",
    "from optimum.utils import NormalizedTextConfig, NormalizedConfigManager\n",
    "from transformers.modeling_outputs import CausalLMOutputWithPast\n",
    "from optimum.intel.openvino.utils import OV_XML_FILE_NAME\n",
    "import numpy as np\n",
    "from pathlib import Path\n",
    "\n",
    "    \n",
    "class OVMPTModel(OVModelForCausalLM):\n",
    "    \"\"\"\n",
    "    Optimum intel compatible model wrapper for MPT\n",
    "    \"\"\"\n",
    "    def __init__(\n",
    "        self,\n",
    "        model: \"Model\",\n",
    "        config: \"PretrainedConfig\" = None,\n",
    "        device: str = \"CPU\",\n",
    "        dynamic_shapes: bool = True,\n",
    "        ov_config: Optional[Dict[str, str]] = None,\n",
    "        model_save_dir: Optional[Union[str, Path]] = None,\n",
    "        **kwargs,\n",
    "    ):\n",
    "        NormalizedConfigManager._conf[\"mpt\"] = NormalizedTextConfig.with_args(num_layers=\"n_layers\", num_attention_heads=\"n_heads\")\n",
    "        super().__init__(model, config, device, dynamic_shapes, ov_config, model_save_dir, **kwargs)\n",
    "\n",
    "    def _reshape(\n",
    "        self,\n",
    "        model: \"Model\",\n",
    "        *args,\n",
    "        **kwargs\n",
    "    ):\n",
    "        shapes = {}\n",
    "        for inputs in model.inputs:\n",
    "            shapes[inputs] = inputs.get_partial_shape()\n",
    "            if shapes[inputs].rank.get_length() in [2, 3]:\n",
    "                shapes[inputs][1] = -1\n",
    "            else:\n",
    "                if \".key\" in inputs.get_any_name():\n",
    "                    shapes[inputs][3] = -1\n",
    "                else:\n",
    "                    shapes[inputs][2] = -1\n",
    "                \n",
    "        model.reshape(shapes)\n",
    "        return model\n",
    "\n",
    "    def forward(\n",
    "        self,\n",
    "        input_ids: torch.LongTensor,\n",
    "        attention_mask: Optional[torch.LongTensor] = None,\n",
    "        past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,\n",
    "        **kwargs,\n",
    "    ) -> CausalLMOutputWithPast:\n",
    "        self.compile()\n",
    "\n",
    "        if self.use_cache and past_key_values is not None:\n",
    "            input_ids = input_ids[:, -1:]\n",
    "\n",
    "        inputs = {}\n",
    "        if past_key_values is not None:\n",
    "            # Flatten the past_key_values\n",
    "            past_key_values = tuple(\n",
    "                past_key_value for pkv_per_layer in past_key_values for past_key_value in pkv_per_layer\n",
    "            )\n",
    "            # Add the past_key_values to the decoder inputs\n",
    "            inputs = dict(zip(self.key_value_input_names, past_key_values))\n",
    "\n",
    "        # Create empty past_key_values for decoder_with_past first generation step\n",
    "        elif self.use_cache:\n",
    "            shape_input_ids = input_ids.shape\n",
    "            num_attention_heads = (\n",
    "                self.normalized_config.num_attention_heads if self.config.model_type == \"bloom\" else 1\n",
    "            )\n",
    "            for input_name in self.key_value_input_names:\n",
    "                model_inputs = self.model.input(input_name)\n",
    "                shape = model_inputs.get_partial_shape()\n",
    "                shape[0] = shape_input_ids[0] * num_attention_heads\n",
    "                if shape[2].is_dynamic:\n",
    "                    shape[2] = 0\n",
    "                if shape[1].is_dynamic:\n",
    "                    shape[1] = 0\n",
    "                if shape.rank.get_length() == 4 and shape[3].is_dynamic:\n",
    "                    shape[3] = 0\n",
    "                inputs[input_name] = ov.Tensor(model_inputs.get_element_type(), shape.get_shape())\n",
    "\n",
    "        inputs[\"input_ids\"] = np.array(input_ids)\n",
    "\n",
    "        # Add the attention_mask inputs when needed\n",
    "        if \"attention_mask\" in self.input_names and attention_mask is not None:\n",
    "            inputs[\"attention_mask\"] = np.array(attention_mask)\n",
    "\n",
    "        # Run inference\n",
    "        self.request.start_async(inputs, shared_memory=True)\n",
    "        self.request.wait()\n",
    "\n",
    "        logits = torch.from_numpy(self.request.get_tensor(\"logits\").data).to(self.device)\n",
    "\n",
    "        if self.use_cache:\n",
    "            # Tuple of length equal to : number of layer * number of past_key_value per decoder layer (2 corresponds to the self-attention layer)\n",
    "            past_key_values = tuple(self.request.get_tensor(key).data for key in self.key_value_output_names)\n",
    "            # Tuple of tuple of length `n_layers`, with each tuple of length equal to 2 (k/v of self-attention)\n",
    "            past_key_values = tuple(\n",
    "                past_key_values[i : i + self.num_pkv] for i in range(0, len(past_key_values), self.num_pkv)\n",
    "            )\n",
    "        else:\n",
    "            past_key_values = None\n",
    "\n",
    "        return CausalLMOutputWithPast(logits=logits, past_key_values=past_key_values)\n",
    "\n",
    "    @classmethod\n",
    "    def _from_pretrained(\n",
    "        cls,\n",
    "        model_id: Union[str, Path],\n",
    "        config: PretrainedConfig,\n",
    "        use_auth_token: Optional[Union[bool, str, None]] = None,\n",
    "        revision: Optional[Union[str, None]] = None,\n",
    "        force_download: bool = False,\n",
    "        cache_dir: Optional[str] = None,\n",
    "        file_name: Optional[str] = None,\n",
    "        subfolder: str = \"\",\n",
    "        from_onnx: bool = False,\n",
    "        local_files_only: bool = False,\n",
    "        load_in_8bit: bool = False,\n",
    "        **kwargs,\n",
    "    ):\n",
    "        model_path = Path(model_id)\n",
    "        default_file_name = OV_XML_FILE_NAME\n",
    "        file_name = file_name or default_file_name\n",
    "\n",
    "        model_cache_path = cls._cached_file(\n",
    "            model_path=model_path,\n",
    "            use_auth_token=use_auth_token,\n",
    "            revision=revision,\n",
    "            force_download=force_download,\n",
    "            cache_dir=cache_dir,\n",
    "            file_name=file_name,\n",
    "            subfolder=subfolder,\n",
    "            local_files_only=local_files_only,\n",
    "        )\n",
    "\n",
    "        model = cls.load_model(model_cache_path, load_in_8bit=load_in_8bit)\n",
    "        init_cls = OVMPTModel\n",
    "\n",
    "        return init_cls(model=model, config=config, model_save_dir=model_cache_path.parent, **kwargs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c53001e7-615f-4eb5-b831-4e2b2ff32826",
   "metadata": {
    "tags": []
   },
   "source": [
    "The cell below demonstrates how to instantiate model based on selected variant of model weights and inference device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3536a1a7",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "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",
    "model_to_run"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "7a041101-7336-40fd-96c9-cd298015a0f3",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/nsavel/.cache/huggingface/modules/transformers_modules/mosaicml/mpt-7b-chat/df5a0d74bf7a93f8f76f64ae6b45e2b996ca4764/configuration_mpt.py:97: UserWarning: alibi is turned on, setting `learned_pos_emb` to `False.`\n",
      "  warnings.warn(f'alibi is turned on, setting `learned_pos_emb` to `False.`')\n",
      "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 pathlib import Path\n",
    "from optimum.intel.openvino import OVModelForCausalLM\n",
    "from transformers import AutoTokenizer\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",
    "model_name = model_configuration[\"model_id\"]\n",
    "\n",
    "ov_config = {'PERFORMANCE_HINT': 'LATENCY', 'NUM_STREAMS': '1', \"CACHE_DIR\": \"\"}\n",
    "\n",
    "tok = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)\n",
    "\n",
    "model_class = OVModelForCausalLM if \"mpt\" not in model_id.value else OVMPTModel\n",
    "ov_model = model_class.from_pretrained(model_dir, device=device.value, ov_config=ov_config, config=AutoConfig.from_pretrained(model_dir, trust_remote_code=True), trust_remote_code=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "8f6f7596-5677-4931-875b-aaabfa23cabc",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_50832/2282042463.py:95: FutureWarning: `shared_memory` is deprecated and will be removed in 2024.0. Value of `shared_memory` is going to override `share_inputs` value. Please use only `share_inputs` explicitly.\n",
      "  self.request.start_async(inputs, shared_memory=True)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2 + 2 = 4:\n"
     ]
    }
   ],
   "source": [
    "tokenizer_kwargs = model_configuration.get(\"tokenizer_kwargs\", {})\n",
    "test_string = \"2 + 2 =\"\n",
    "input_tokens = tok(test_string, return_tensors=\"pt\", **tokenizer_kwargs)\n",
    "answer = ov_model.generate(**input_tokens, max_new_tokens=2)\n",
    "print(tok.batch_decode(answer)[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "24d622d0-be46-47c0-a762-88cb50ab15a9",
   "metadata": {},
   "source": [
    "## Run Chatbot\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",
    "The diagram below illustrates how the chatbot pipeline works\n",
    "\n",
    "![generation pipeline](https://user-images.githubusercontent.com/29454499/255523209-d9336491-c7ba-4dc1-98f0-07f23743ce89.png)\n",
    "\n",
    "As can be seen, the pipeline very similar to instruction-following with only changes that previous conversation history additionally passed as input with next user question for getting wider input context. On the first iteration, the user provided instructions joined to conversation history (if exists) converted to token ids using a tokenizer, then prepared input provided to the model. The model generates probabilities for all tokens in logits format  The way the next token will be selected over predicted probabilities is driven by the selected decoding methodology. You can find more information about the most popular decoding methods in this [blog](https://huggingface.co/blog/how-to-generate). The result generation updates conversation history for next conversation step. it makes stronger connection of next question with previously provided and allows user to make clarifications regarding previously provided answers."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "725544ea-05ec-40d7-bbbc-1dc87cf57d04",
   "metadata": {},
   "source": [
    "There are several parameters that can control text generation quality: \n",
    "  * `Temperature` is a parameter used to control the level of creativity in AI-generated text. By adjusting the `temperature`, you can influence the AI model's probability distribution, making the text more focused or diverse.  \n",
    "  Consider the following example: The AI model has to complete the sentence \"The cat is ____.\" with the following token probabilities:  \n",
    "\n",
    "    playing: 0.5  \n",
    "    sleeping: 0.25  \n",
    "    eating: 0.15  \n",
    "    driving: 0.05  \n",
    "    flying: 0.05  \n",
    "\n",
    "    - **Low temperature** (e.g., 0.2): The AI model becomes more focused and deterministic, choosing tokens with the highest probability, such as \"playing.\"  \n",
    "    - **Medium temperature** (e.g., 1.0): The AI model maintains a balance between creativity and focus, selecting tokens based on their probabilities without significant bias, such as \"playing,\" \"sleeping,\" or \"eating.\"  \n",
    "    - **High temperature** (e.g., 2.0): The AI model becomes more adventurous, increasing the chances of selecting less likely tokens, such as \"driving\" and \"flying.\"\n",
    "  * `Top-p`, also known as nucleus sampling, is a parameter used to control the range of tokens considered by the AI model based on their cumulative probability. By adjusting the `top-p` value, you can influence the AI model's token selection, making it more focused or diverse.\n",
    "  Using the same example with the cat, consider the following top_p settings:  \n",
    "    - **Low top_p** (e.g., 0.5): The AI model considers only tokens with the highest cumulative probability, such as \"playing.\"  \n",
    "    - **Medium top_p** (e.g., 0.8): The AI model considers tokens with a higher cumulative probability, such as \"playing,\" \"sleeping,\" and \"eating.\"  \n",
    "    - **High top_p** (e.g., 1.0): The AI model considers all tokens, including those with lower probabilities, such as \"driving\" and \"flying.\" \n",
    "  * `Top-k` is an another popular sampling strategy. In comparison with Top-P, which chooses from the smallest possible set of words whose cumulative probability exceeds the probability P, in Top-K sampling K most likely next words are filtered and the probability mass is redistributed among only those K next words. In our example with cat, if k=3, then only \"playing\", \"sleeping\" and \"eating\" will be taken into account as possible next word.\n",
    "  * `Repetition Penalty` This parameter can help penalize tokens based on how frequently they occur in the text, including the input prompt. A token that has already appeared five times is penalized more heavily than a token that has appeared only one time. A value of 1 means that there is no penalty and values larger than 1 discourage repeated tokens."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "01f8f7f8-072e-45dc-b7c9-18d8c3c47754",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/nsavel/venvs/ov_notebooks_tmp/lib/python3.8/site-packages/gradio/blocks.py:889: UserWarning: api_name user already exists, using user_1\n",
      "  warnings.warn(f\"api_name {api_name} already exists, using {api_name_}\")\n",
      "/home/nsavel/venvs/ov_notebooks_tmp/lib/python3.8/site-packages/gradio/blocks.py:889: UserWarning: api_name bot already exists, using bot_1\n",
      "  warnings.warn(f\"api_name {api_name} already exists, using {api_name_}\")\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Running on local URL:  http://127.0.0.1:7860\n",
      "\n",
      "To create a public link, set `share=True` in `launch()`.\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div><iframe src=\"http://127.0.0.1:7860/\" width=\"100%\" height=\"500\" allow=\"autoplay; camera; microphone; clipboard-read; clipboard-write;\" frameborder=\"0\" allowfullscreen></iframe></div>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": []
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from threading import Event, Thread\n",
    "from uuid import uuid4\n",
    "\n",
    "import gradio as gr\n",
    "import torch\n",
    "from transformers import (\n",
    "    AutoTokenizer,\n",
    "    StoppingCriteria,\n",
    "    StoppingCriteriaList,\n",
    "    TextIteratorStreamer,\n",
    ")\n",
    "\n",
    "\n",
    "model_name = model_configuration[\"model_id\"]\n",
    "history_template = model_configuration[\"history_template\"]\n",
    "current_message_template = model_configuration[\"current_message_template\"]\n",
    "start_message = model_configuration[\"start_message\"]\n",
    "stop_tokens = model_configuration.get(\"stop_tokens\")\n",
    "tokenizer_kwargs = model_configuration.get(\"tokenizer_kwargs\", {})\n",
    "\n",
    "\n",
    "\n",
    "max_new_tokens = 256\n",
    "\n",
    "class StopOnTokens(StoppingCriteria):\n",
    "    def __init__(self, token_ids):\n",
    "        self.token_ids = token_ids\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",
    "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",
    "def default_partial_text_processor(partial_text:str, new_text:str):\n",
    "    \"\"\"\n",
    "    helper for updating partially generated answer, used by de\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",
    "text_processor = model_configuration.get(\"partial_text_processor\", default_partial_text_processor)\n",
    "    \n",
    "def convert_history_to_text(history:List[Tuple[str, str]]):\n",
    "    \"\"\"\n",
    "    function for conversion history stored as list pairs of user and assistant messages to string according to model expected conversation template\n",
    "    Params:\n",
    "      history: dialogue history\n",
    "    Returns:\n",
    "      history in text format\n",
    "    \"\"\"\n",
    "    text = start_message + \"\".join(\n",
    "        [\n",
    "            \"\".join(\n",
    "                [\n",
    "                    history_template.format(user=item[0], assistant=item[1])\n",
    "                ]\n",
    "            )\n",
    "            for item in history[:-1]\n",
    "        ]\n",
    "    )\n",
    "    text += \"\".join(\n",
    "        [\n",
    "            \"\".join(\n",
    "                [\n",
    "                    current_message_template.format(user=history[-1][0], assistant=history[-1][1])\n",
    "                ]\n",
    "            )\n",
    "        ]\n",
    "    )\n",
    "    return text\n",
    "\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",
    "\n",
    "    # Construct the input message string for the model by concatenating the current system message and conversation history\n",
    "    messages = convert_history_to_text(history)\n",
    "\n",
    "    # Tokenize the messages string\n",
    "    input_ids = tok(messages, return_tensors=\"pt\", **tokenizer_kwargs).input_ids\n",
    "    if input_ids.shape[1] > 2000:\n",
    "        history = [history[-1]]\n",
    "        messages = convert_history_to_text(history)\n",
    "        input_ids = tok(messages, return_tensors=\"pt\", **tokenizer_kwargs).input_ids\n",
    "    streamer = TextIteratorStreamer(tok, timeout=30.0, skip_prompt=True, skip_special_tokens=True)\n",
    "    generate_kwargs = dict(\n",
    "        input_ids=input_ids,\n",
    "        max_new_tokens=max_new_tokens,\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",
    "    stream_complete = Event()\n",
    "\n",
    "    def generate_and_signal_complete():\n",
    "        \"\"\"\n",
    "        genration function for single thread\n",
    "        \"\"\"\n",
    "        global start_time\n",
    "        ov_model.generate(**generate_kwargs)\n",
    "        stream_complete.set()\n",
    "\n",
    "    t1 = Thread(target=generate_and_signal_complete)\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 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(\n",
    "        f\"\"\"<h1><center>OpenVINO {model_id.value} Chatbot</center></h1>\"\"\"\n",
    "    )\n",
    "    chatbot = gr.Chatbot(height=500)\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",
    "    with gr.Row():\n",
    "        with gr.Accordion(\"Advanced Options:\", 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",
    "    gr.Examples([\n",
    "        [\"Hello there! How are you doing?\"],\n",
    "        [\"What is OpenVINO?\"],\n",
    "        [\"Who are you?\"],\n",
    "        [\"Can you explain to me briefly what is Python programming language?\"],\n",
    "        [\"Explain the plot of Cinderella in a sentence.\"],\n",
    "        [\"What are some common mistakes to avoid when writing code?\"],\n",
    "        [\"Write a 100-word blog post on “Benefits of Artificial Intelligence and OpenVINO“\"]\n",
    "    ], \n",
    "        inputs=msg, \n",
    "        label=\"Click on any example and press the 'Submit' button\"\n",
    "    )\n",
    "\n",
    "    submit_event = msg.submit(\n",
    "        fn=user,\n",
    "        inputs=[msg, chatbot],\n",
    "        outputs=[msg, chatbot],\n",
    "        queue=False,\n",
    "    ).then(\n",
    "        fn=bot,\n",
    "        inputs=[\n",
    "            chatbot,\n",
    "            temperature,\n",
    "            top_p,\n",
    "            top_k,\n",
    "            repetition_penalty,\n",
    "            conversation_id,\n",
    "        ],\n",
    "        outputs=chatbot,\n",
    "        queue=True,\n",
    "    )\n",
    "    submit_click_event = submit.click(\n",
    "        fn=user,\n",
    "        inputs=[msg, chatbot],\n",
    "        outputs=[msg, chatbot],\n",
    "        queue=False,\n",
    "    ).then(\n",
    "        fn=bot,\n",
    "        inputs=[\n",
    "            chatbot,\n",
    "            temperature,\n",
    "            top_p,\n",
    "            top_k,\n",
    "            repetition_penalty,\n",
    "            conversation_id,\n",
    "        ],\n",
    "        outputs=chatbot,\n",
    "        queue=True,\n",
    "    )\n",
    "    stop.click(\n",
    "        fn=None,\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(max_size=2)\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": 14,
   "id": "7b837f9e-4152-4a5c-880a-ed874aa64a74",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Closing server running on port: 7860\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.8.10"
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "state": {},
    "version_major": 2,
    "version_minor": 0
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}