notebooks/llm-question-answering/llm-question-answering.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "f6aa8c52-9bde-41a6-a5f7-4fa93d5c2a6c",
   "metadata": {},
   "source": [
    "# LLM Instruction-following pipeline with OpenVINO \n",
    "\n",
    "LLM stands for “Large Language Model,” which refers to a type of artificial intelligence model that is designed to understand and generate human-like text based on the input it receives. LLMs are trained on large datasets of text to learn patterns, grammar, and semantic relationships, allowing them to generate coherent and contextually relevant responses. One core capability of Large Language Models (LLMs) is to follow natural language instructions. Instruction-following models are capable of generating text in response to prompts and are often used for tasks like writing assistance, chatbots, and content generation.\n",
    "\n",
    "In this tutorial, we consider how to run an instruction-following text generation pipeline using popular LLMs and OpenVINO. We will use pre-trained models from the [Hugging Face Transformers](https://huggingface.co/docs/transformers/index) library. The [Hugging Face Optimum Intel](https://huggingface.co/docs/optimum/intel/index) library converts the models to OpenVINO™ IR format. To simplify the user experience, we will use [OpenVINO Generate API](https://github.com/openvinotoolkit/openvino.genai) for generation of instruction-following inference pipeline.  \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 INT8 and INT4 with [OpenVINO NNCF](https://github.com/openvinotoolkit/nncf)\n",
    "- Create an instruction-following inference pipeline with [Generate API](https://github.com/openvinotoolkit/openvino.genai)\n",
    "- Run instruction-following pipeline\n",
    "\n",
    "\n",
    "#### Table of contents:\n",
    "\n",
    "- [Prerequisites](#Prerequisites)\n",
    "- [Select model for inference](#Select-model-for-inference)\n",
    "- [Download and convert model to OpenVINO IR via Optimum Intel CLI](#Download-and-convert-model-to-OpenVINO-IR-via-Optimum-Intel-CLI)\n",
    "- [Compress model weights](#Compress-model-weights)\n",
    "    - [Weights Compression using Optimum Intel CLI](#Weights-Compression-using-Optimum-Intel-CLI)\n",
    "    - [Weights Compression using NNCF](#Weights-Compression-using-NNCF)\n",
    "- [Select device for inference and model variant](#Select-device-for-inference-and-model-variant)\n",
    "- [Create an instruction-following inference pipeline](#Create-an-instruction-following-inference-pipeline)\n",
    "    - [Setup imports](#Setup-imports)\n",
    "    - [Prepare text streamer to get results runtime](#Prepare-text-streamer-to-get-results-runtime)\n",
    "    - [Main generation function](#Main-generation-function)\n",
    "    - [Helpers for application](#Helpers-for-application)\n",
    "- [Run instruction-following pipeline](#Run-instruction-following-pipeline)\n",
    "\n",
    "\n",
    "### Installation Instructions\n",
    "\n",
    "This is a self-contained example that relies solely on its own code.\n",
    "\n",
    "We recommend  running the notebook in a virtual environment. You only need a Jupyter server to start.\n",
    "For details, please refer to [Installation Guide](https://github.com/openvinotoolkit/openvino_notebooks/blob/latest/README.md#-installation-guide).\n",
    "\n",
    "<img referrerpolicy=\"no-referrer-when-downgrade\" src=\"https://static.scarf.sh/a.png?x-pxid=5b5a4db0-7875-4bfb-bdbd-01698b5b1a77&file=notebooks/llm-question-answering/llm-question-answering.ipynb\" />\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "027108c2-1fbe-4be5-9e23-3fc359185a42",
   "metadata": {},
   "source": [
    "## Prerequisites\n",
    "[back to top ⬆️](#Table-of-contents:)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5d0473c6-3734-422d-a370-2e39d576be0e",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip uninstall -q -y optimum optimum-intel\n",
    "%pip install  -Uq \"openvino>=2024.3.0\" \"openvino-genai\"\n",
    "%pip install -q \"torch>=2.1\" \"nncf>=2.7\" \"transformers>=4.40.0\" \"huggingface-hub>=0.26.5\" \"onnx<1.16.2\" \"optimum>=1.16.1\" \"accelerate\" \"datasets>=2.14.6\" \"gradio>=4.19\" \"git+https://github.com/huggingface/optimum-intel.git\" --extra-index-url https://download.pytorch.org/whl/cpu"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "611cc777-d5bc-4c7b-92e4-a4befa13b2ce",
   "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",
    "* **tiny-llama-1b-chat** - This is the chat model finetuned on top of [TinyLlama/TinyLlama-1.1B-intermediate-step-1431k-3T](https://huggingface.co/TinyLlama/TinyLlama-1.1B-intermediate-step-1431k-3T). 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-v1.0)\n",
    "* **phi-2** - Phi-2 is a Transformer with 2.7 billion parameters. It was trained using the same data sources as [Phi-1.5](https://huggingface.co/microsoft/phi-1_5), augmented with a new data source that consists of various NLP synthetic texts and filtered websites (for safety and educational value). When assessed against benchmarks testing common sense, language understanding, and logical reasoning, Phi-2 showcased a nearly state-of-the-art performance among models with less than 13 billion parameters. More details about model can be found in [model card](https://huggingface.co/microsoft/phi-2#limitations-of-phi-2).\n",
    "* **dolly-v2-3b** - Dolly 2.0 is an instruction-following large language model trained on the Databricks machine-learning platform that is licensed for commercial use. It is based on [Pythia](https://github.com/EleutherAI/pythia) and is trained on ~15k instruction/response fine-tuning records generated by Databricks employees in various capability domains, including brainstorming, classification, closed QA, generation, information extraction, open QA, and summarization. Dolly 2.0 works by processing natural language instructions and generating responses that follow the given instructions. It can be used for a wide range of applications, including closed question-answering, summarization, and generation. More details about model can be found in [model card](https://huggingface.co/databricks/dolly-v2-3b).\n",
    "* **red-pajama-3b-instruct** -  A 2.8B parameter pre-trained language model based on GPT-NEOX architecture. The model was fine-tuned for few-shot applications on the data of [GPT-JT](https://huggingface.co/togethercomputer/GPT-JT-6B-v1), with exclusion of tasks that overlap with the HELM core scenarios.More details about model can be found in [model card](https://huggingface.co/togethercomputer/RedPajama-INCITE-Instruct-3B-v1).\n",
    "* **mistral-7b** - The Mistral-7B-v0.2 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-Instruct-v0.2), [paper](https://arxiv.org/abs/2310.06825) and [release blog post](https://mistral.ai/news/announcing-mistral-7b/).\n",
    "* **llama-3-8b-instruct** - Llama 3 is an auto-regressive language model that uses an optimized transformer architecture. The tuned versions use supervised fine-tuning (SFT) and reinforcement learning with human feedback (RLHF) to align with human preferences for helpfulness and safety. The Llama 3 instruction tuned models are optimized for dialogue use cases and outperform many of the available open source chat models on common industry benchmarks. More details about model can be found in [Meta blog post](https://ai.meta.com/blog/meta-llama-3/), [model website](https://llama.meta.com/llama3) and [model card](https://huggingface.co/meta-llama/Meta-Llama-3-8B-Instruct).\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",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "51920510-effc-49ef-8c6f-81c951d96a9b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "import requests\n",
    "\n",
    "if not Path(\"notebook_utils.py\").exists():\n",
    "    # Fetch `notebook_utils` module\n",
    "    r = requests.get(\n",
    "        url=\"https://raw.githubusercontent.com/openvinotoolkit/openvino_notebooks/latest/utils/notebook_utils.py\",\n",
    "    )\n",
    "    open(\"notebook_utils.py\", \"w\").write(r.text)\n",
    "from notebook_utils import download_file, device_widget\n",
    "\n",
    "if not Path(\"./config.py\").exists():\n",
    "    download_file(url=\"https://raw.githubusercontent.com/openvinotoolkit/openvino_notebooks/latest/notebooks/llm-question-answering/config.py\")\n",
    "from config import SUPPORTED_LLM_MODELS\n",
    "import ipywidgets as widgets\n",
    "\n",
    "# Read more about telemetry collection at https://github.com/openvinotoolkit/openvino_notebooks?tab=readme-ov-file#-telemetry\n",
    "from notebook_utils import collect_telemetry\n",
    "\n",
    "collect_telemetry(\"llm-question-answering.ipynb\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "27b42290-a9b5-4453-9a4c-ffa44bbd966d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "aed83fe4a63f44a8af7508cd115fcf4b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Dropdown(description='Model:', index=1, options=('tiny-llama-1b', 'phi-2', 'dolly-v2-3b', 'red-pajama-instruct…"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model_ids = list(SUPPORTED_LLM_MODELS)\n",
    "\n",
    "model_id = widgets.Dropdown(\n",
    "    options=model_ids,\n",
    "    value=model_ids[1],\n",
    "    description=\"Model:\",\n",
    "    disabled=False,\n",
    ")\n",
    "\n",
    "model_id"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "37e9634f-4fc7-4d9c-9ade-b3e8684a0828",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Selected model dolly-v2-3b\n"
     ]
    }
   ],
   "source": [
    "model_configuration = SUPPORTED_LLM_MODELS[model_id.value]\n",
    "print(f\"Selected model {model_id.value}\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "4e4fd394-b4fb-4eef-8bdc-d116572aa8f0",
   "metadata": {},
   "source": [
    "## Download and convert model to OpenVINO IR via Optimum Intel CLI\n",
    "[back to top ⬆️](#Table-of-contents:)\n",
    "\n",
    "Listed model are available for downloading via the [HuggingFace hub](https://huggingface.co/models). We will use optimum-cli interface for exporting it into OpenVINO Intermediate Representation (IR) format.\n",
    "\n",
    "Optimum CLI interface for converting models supports export to OpenVINO (supported starting optimum-intel 1.12 version).\n",
    "General command format:\n",
    "\n",
    "```bash\n",
    "optimum-cli export openvino --model <model_id_or_path> --task <task> <output_dir>\n",
    "```\n",
    "\n",
    "where `--model` argument is model id from HuggingFace Hub or local directory with model (saved using `.save_pretrained` method), `--task ` is one of [supported task](https://huggingface.co/docs/optimum/exporters/task_manager) that exported model should solve. For LLMs it will be `text-generation-with-past`. If model initialization requires to use remote code, `--trust-remote-code` flag additionally should be passed. Full list of supported arguments available via `--help` For more details and examples of usage, please check [optimum documentation](https://huggingface.co/docs/optimum/intel/inference#export).\n",
    "\n",
    "## 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 CLI\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 `--weight-format int8` to `optimum-cli` command line. For 4 bit compression we provide `--weight-format int4` and some other options containing number of bits and other compression parameters. An example of this approach usage you can find in [llm-chatbot notebook](../llm-chatbot)\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 the OpenVINO model instance and compresses its weights for Linear and Embedding layers. We will consider this variant in this notebook for both int4 and int8 compression.\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.\n",
    ">**Note**: There may be no speedup for INT4/INT8 compressed models on dGPU"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "f81602ca-4674-4b61-b2c8-ca11631428b1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "18f82da95c7048fe9377f0546e4e885a",
       "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": "204a990e71c84893aab873b0c3fc2f0a",
       "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": "faca0d176aa04e2f927b1c1152934593",
       "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, Markdown\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": 6,
   "id": "0430babb-f8c3-4b7e-86e0-1ea23de68477",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:nncf:NNCF initialized successfully. Supported frameworks detected: torch, onnx, openvino\n"
     ]
    }
   ],
   "source": [
    "from pathlib import Path\n",
    "import logging\n",
    "import openvino as ov\n",
    "import nncf\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",
    "core = ov.Core()\n",
    "\n",
    "\n",
    "def convert_to_fp16():\n",
    "    if (fp16_model_dir / \"openvino_model.xml\").exists():\n",
    "        return\n",
    "    export_command_base = \"optimum-cli export openvino --model {} --task text-generation-with-past --weight-format fp16\".format(pt_model_id)\n",
    "    export_command = export_command_base + \" \" + str(fp16_model_dir)\n",
    "    display(Markdown(\"**Export command:**\"))\n",
    "    display(Markdown(f\"`{export_command}`\"))\n",
    "    ! $export_command\n",
    "\n",
    "\n",
    "def convert_to_int8():\n",
    "    if (int8_model_dir / \"openvino_model.xml\").exists():\n",
    "        return\n",
    "    int8_model_dir.mkdir(parents=True, exist_ok=True)\n",
    "    export_command_base = \"optimum-cli export openvino --model {} --task text-generation-with-past --weight-format int8\".format(pt_model_id)\n",
    "    export_command = export_command_base + \" \" + str(int8_model_dir)\n",
    "    display(Markdown(\"**Export command:**\"))\n",
    "    display(Markdown(f\"`{export_command}`\"))\n",
    "    ! $export_command\n",
    "\n",
    "\n",
    "def convert_to_int4():\n",
    "    compression_configs = {\n",
    "        \"mistral-7b\": {\n",
    "            \"sym\": True,\n",
    "            \"group_size\": 64,\n",
    "            \"ratio\": 0.6,\n",
    "        },\n",
    "        \"red-pajama-3b-instruct\": {\n",
    "            \"sym\": False,\n",
    "            \"group_size\": 128,\n",
    "            \"ratio\": 0.5,\n",
    "        },\n",
    "        \"dolly-v2-3b\": {\"sym\": False, \"group_size\": 32, \"ratio\": 0.5},\n",
    "        \"llama-3-8b-instruct\": {\"sym\": True, \"group_size\": 128, \"ratio\": 1.0},\n",
    "        \"default\": {\n",
    "            \"sym\": False,\n",
    "            \"group_size\": 128,\n",
    "            \"ratio\": 0.8,\n",
    "        },\n",
    "    }\n",
    "\n",
    "    model_compression_params = compression_configs.get(model_id.value, compression_configs[\"default\"])\n",
    "    if (int4_model_dir / \"openvino_model.xml\").exists():\n",
    "        return\n",
    "    export_command_base = \"optimum-cli export openvino --model {} --task text-generation-with-past --weight-format int4\".format(pt_model_id)\n",
    "    int4_compression_args = \" --group-size {} --ratio {}\".format(model_compression_params[\"group_size\"], model_compression_params[\"ratio\"])\n",
    "    if model_compression_params[\"sym\"]:\n",
    "        int4_compression_args += \" --sym\"\n",
    "    export_command_base += int4_compression_args\n",
    "    export_command = export_command_base + \" \" + str(int4_model_dir)\n",
    "    display(Markdown(\"**Export command:**\"))\n",
    "    display(Markdown(f\"`{export_command}`\"))\n",
    "    ! $export_command\n",
    "\n",
    "\n",
    "if prepare_fp16_model.value:\n",
    "    convert_to_fp16()\n",
    "if prepare_int8_model.value:\n",
    "    convert_to_int8()\n",
    "if prepare_int4_model.value:\n",
    "    convert_to_int4()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "60355f86-2250-4ebe-82ac-950f2d4fb01b",
   "metadata": {},
   "source": [
    "Let's compare model size for different compression types"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "42c7b254-1ce4-4f23-813a-9bdc23aed327",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Size of FP16 model is 5297.21 MB\n",
      "Size of model with INT8 compressed weights is 2656.29 MB\n",
      "Compression rate for INT8 model: 1.994\n",
      "Size of model with INT4 compressed weights is 2154.54 MB\n",
      "Compression rate for INT4 model: 2.459\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}\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "3df73379-bccc-41b1-9c94-c3040819805b",
   "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."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "d2d7bf5b-8a05-4c3b-a36b-631af5c197e9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "76c6c52c8af04d9084b1e3b56686f563",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Dropdown(description='Device:', options=('CPU', 'AUTO'), value='CPU')"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "core = ov.Core()\n",
    "\n",
    "device = device_widget(\"CPU\", exclude=[\"NPU\"])\n",
    "\n",
    "device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "24532480-80a5-4953-9cd6-78ac51a1cd8f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "44219c88d2da477b887a479227591704",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Dropdown(description='Model to run:', options=('INT4', 'INT8', 'FP16'), value='INT4')"
      ]
     },
     "execution_count": 9,
     "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": 10,
   "id": "5259c1c5-4128-4210-9ad2-faf33ee40e86",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading model from dolly-v2-3b/INT8_compressed_weights\n"
     ]
    }
   ],
   "source": [
    "from transformers import AutoTokenizer\n",
    "from openvino_tokenizers import convert_tokenizer\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",
    "# optionally convert tokenizer if used cached model without it\n",
    "if not (model_dir / \"openvino_tokenizer.xml\").exists() or not (model_dir / \"openvino_detokenizer.xml\").exists():\n",
    "    hf_tokenizer = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=True)\n",
    "    ov_tokenizer, ov_detokenizer = convert_tokenizer(hf_tokenizer, with_detokenizer=True)\n",
    "    ov.save_model(ov_tokenizer, model_dir / \"openvino_tokenizer.xml\")\n",
    "    ov.save_model(ov_tokenizer, model_dir / \"openvino_detokenizer.xml\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "bf13f6c3-6671-408e-ae0e-aaa3d8a6eaac",
   "metadata": {},
   "source": [
    "## Create an instruction-following inference pipeline\n",
    "[back to top ⬆️](#Table-of-contents:)\n",
    " \n",
    "The `run_generation` function accepts user-provided text input, tokenizes it, and runs the generation process. Text generation is an iterative process, where each next token depends on previously generated until a maximum number of tokens or stop generation condition is not reached. \n",
    " \n",
    "The diagram below illustrates how the instruction-following pipeline works\n",
    "\n",
    "![generation pipeline)](https://github.com/openvinotoolkit/openvino_notebooks/assets/29454499/e881f4a4-fcc8-427a-afe1-7dd80aebd66e)\n",
    "\n",
    "As can be seen, on the first iteration, the user provided instructions. Instructions is 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).\n",
    "\n",
    "To simplify user experience we will use [OpenVINO Generate API](https://github.com/openvinotoolkit/openvino.genai/blob/master/src/README.md).\n",
    "Firstly we will create pipeline with `LLMPipeline`. `LLMPipeline` is the main object used for decoding. You can construct it straight away from the folder with the converted model. It will automatically load the `main model`, `tokenizer`, `detokenizer` and default `generation configuration`. \n",
    "After that we will configure parameters for decoding. We can get default config with `get_generation_config()`, setup parameters and apply the updated version with `set_generation_config(config)` or put config directly to `generate()`. It's also possible to specify the needed options just as inputs in the `generate()` method, as shown below.\n",
    "Then we just run `generate` method and get the output in text format. We do not need to encode input prompt according to model expected template or write post-processing code for logits decoder, it will be done easily with LLMPipeline. \n",
    "\n",
    "To obtain intermediate generation results without waiting until when generation is finished, we will write class-iterator based on `StreamerBase` class of `openvino_genai`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "f1f295d4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " of the presence of chlorophyll\n",
      "in its leaves. Chlorophyll absorbs all\n",
      "visible sunlight and this causes it to\n",
      "turn from a green to yellow colour.\n",
      "The Sun is yellow bacause of the presence of chlorophyll in its leaves. Chlorophyll absorbs all\n",
      "visible sunlight and this causes it to\n",
      "turn from a green to yellow colour.\n",
      "The yellow colour of the Sun is the\n",
      "colour we perceive as the colour of the\n",
      "sun. It also causes us to perceive the\n",
      "sun as yellow. This property is called\n",
      "the yellow colouration of the Sun and it\n",
      "is caused by the presence of chlorophyll\n",
      "in the leaves of plants.\n",
      "Chlorophyll is also responsible for the green colour of plants\n"
     ]
    }
   ],
   "source": [
    "import openvino_genai as ov_genai\n",
    "\n",
    "pipe = ov_genai.LLMPipeline(model_dir.as_posix(), device.value)\n",
    "print(pipe.generate(\"The Sun is yellow bacause\", temperature=1.2, top_k=4, do_sample=True, max_new_tokens=150))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "abcd2f8e",
   "metadata": {},
   "source": [
    "There are several parameters that can control text generation quality:\n",
    "\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 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",
    "\n",
    "The generation cycle repeats until the end of the sequence token is reached or it also can be interrupted when maximum tokens will be generated. As already mentioned before, we can enable printing current generated tokens without waiting until when the whole generation is finished using Streaming API, it adds a new token to the output queue and then prints them when they are ready."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "eb7af692-0a15-4493-86ef-a80cda21551c",
   "metadata": {},
   "source": [
    "### Setup imports\n",
    "[back to top ⬆️](#Table-of-contents:)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "19d23a99-3284-42db-b7dc-5805d219f70d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from threading import Thread\n",
    "from time import perf_counter\n",
    "from typing import List\n",
    "import numpy as np\n",
    "from queue import Queue\n",
    "import re"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "b341a9e6-290d-4780-90da-1ee64cee436d",
   "metadata": {},
   "source": [
    "### Prepare text streamer to get results runtime\n",
    "[back to top ⬆️](#Table-of-contents:)\n",
    "\n",
    "Load the `detokenizer`, use it to convert token_id to string output format. We will collect print-ready text in a queue and give the text when it is needed. It will help estimate performance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "e2638c5b-47ad-4213-80da-8cfc2659b3aa",
   "metadata": {},
   "outputs": [],
   "source": [
    "core = ov.Core()\n",
    "\n",
    "detokinizer_dir = Path(model_dir, \"openvino_detokenizer.xml\")\n",
    "\n",
    "\n",
    "class TextIteratorStreamer(ov_genai.StreamerBase):\n",
    "    def __init__(self, tokenizer):\n",
    "        super().__init__()\n",
    "        self.tokenizer = tokenizer\n",
    "        self.compiled_detokenizer = core.compile_model(detokinizer_dir.as_posix())\n",
    "        self.text_queue = Queue()\n",
    "        self.stop_signal = None\n",
    "\n",
    "    def __iter__(self):\n",
    "        return self\n",
    "\n",
    "    def __next__(self):\n",
    "        value = self.text_queue.get()\n",
    "        if value == self.stop_signal:\n",
    "            raise StopIteration()\n",
    "        else:\n",
    "            return value\n",
    "\n",
    "    def put(self, token_id):\n",
    "        openvino_output = self.compiled_detokenizer(np.array([[token_id]], dtype=int))\n",
    "        text = str(openvino_output[\"string_output\"][0])\n",
    "        # remove labels/special symbols\n",
    "        text = text.lstrip(\"!\")\n",
    "        text = re.sub(\"<.*>\", \"\", text)\n",
    "        self.text_queue.put(text)\n",
    "\n",
    "    def end(self):\n",
    "        self.text_queue.put(self.stop_signal)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "55dbc4ae-da28-4be8-b928-3dd68c197937",
   "metadata": {},
   "source": [
    "### Main generation function\n",
    "[back to top ⬆️](#Table-of-contents:)\n",
    "\n",
    "As it was discussed above, `run_generation` function is the entry point for starting generation. It gets provided input instruction as parameter and returns model response."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "27802e81-9d42-4d71-99ea-5f76db5237f1",
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_generation(\n",
    "    user_text: str,\n",
    "    top_p: float,\n",
    "    temperature: float,\n",
    "    top_k: int,\n",
    "    max_new_tokens: int,\n",
    "    perf_text: str,\n",
    "):\n",
    "    \"\"\"\n",
    "    Text generation function\n",
    "\n",
    "    Parameters:\n",
    "      user_text (str): User-provided instruction for a generation.\n",
    "      top_p (float):  Nucleus sampling. If set to < 1, only the smallest set of most probable tokens with probabilities that add up to top_p or higher are kept for a generation.\n",
    "      temperature (float): The value used to module the logits distribution.\n",
    "      top_k (int): The number of highest probability vocabulary tokens to keep for top-k-filtering.\n",
    "      max_new_tokens (int): Maximum length of generated sequence.\n",
    "      perf_text (str): Content of text field for printing performance results.\n",
    "    Returns:\n",
    "      model_output (str) - model-generated text\n",
    "      perf_text (str) - updated perf text filed content\n",
    "    \"\"\"\n",
    "\n",
    "    # setup config for decoding stage\n",
    "    config = pipe.get_generation_config()\n",
    "    config.temperature = temperature\n",
    "    if top_k > 0:\n",
    "        config.top_k = top_k\n",
    "    config.top_p = top_p\n",
    "    config.do_sample = True\n",
    "    config.max_new_tokens = max_new_tokens\n",
    "\n",
    "    # Start generation on a separate thread, so that we don't block the UI. The text is pulled from the streamer\n",
    "    # in the main thread.\n",
    "    streamer = TextIteratorStreamer(pipe.get_tokenizer())\n",
    "    t = Thread(target=pipe.generate, args=(user_text, config, streamer))\n",
    "    t.start()\n",
    "\n",
    "    model_output = \"\"\n",
    "    per_token_time = []\n",
    "    num_tokens = 0\n",
    "    start = perf_counter()\n",
    "    for new_text in streamer:\n",
    "        current_time = perf_counter() - start\n",
    "        model_output += new_text\n",
    "        perf_text, num_tokens = estimate_latency(current_time, perf_text, per_token_time, num_tokens)\n",
    "        yield model_output, perf_text\n",
    "        start = perf_counter()\n",
    "    return model_output, perf_text"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "6d7b7182-1d9b-485b-81f2-1a9fab6d5cd4",
   "metadata": {},
   "source": [
    "### Helpers for application\n",
    "[back to top ⬆️](#Table-of-contents:)\n",
    "\n",
    "For making interactive user interface we will use Gradio library. The code bellow provides useful functions used for communication with UI elements."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "c9b2a2b7-5a91-470f-98d7-7355664bc1be",
   "metadata": {},
   "outputs": [],
   "source": [
    "def estimate_latency(\n",
    "    current_time: float,\n",
    "    current_perf_text: str,\n",
    "    per_token_time: List[float],\n",
    "    num_tokens: int,\n",
    "):\n",
    "    \"\"\"\n",
    "    Helper function for performance estimation\n",
    "\n",
    "    Parameters:\n",
    "      current_time (float): This step time in seconds.\n",
    "      current_perf_text (str): Current content of performance UI field.\n",
    "      per_token_time (List[float]): history of performance from previous steps.\n",
    "      num_tokens (int): Total number of generated tokens.\n",
    "\n",
    "    Returns:\n",
    "      update for performance text field\n",
    "      update for a total number of tokens\n",
    "    \"\"\"\n",
    "    num_tokens += 1\n",
    "    per_token_time.append(1 / current_time)\n",
    "    if len(per_token_time) > 10 and len(per_token_time) % 4 == 0:\n",
    "        current_bucket = per_token_time[:-10]\n",
    "        return (\n",
    "            f\"Average generation speed: {np.mean(current_bucket):.2f} tokens/s. Total generated tokens: {num_tokens}\",\n",
    "            num_tokens,\n",
    "        )\n",
    "    return current_perf_text, num_tokens"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "31ebb167-0e55-4271-aedd-13814c2356d2",
   "metadata": {},
   "source": [
    "## Run instruction-following pipeline\n",
    "[back to top ⬆️](#Table-of-contents:)\n",
    "\n",
    "Now, we are ready to explore model capabilities. This demo provides a simple interface that allows communication with a model using text instruction. Type your instruction into the `User instruction` field or select one from predefined examples and click on the `Submit` button to start generation. Additionally, you can modify advanced generation parameters:\n",
    "\n",
    "* `Device` - allows switching inference device. Please note, every time when new device is selected, model will be recompiled and this takes some time.\n",
    "* `Max New Tokens` - maximum size of generated text.\n",
    "* `Top-p (nucleus sampling)` -  if set to < 1, only the smallest set of most probable tokens with probabilities that add up to top_p or higher are kept for a generation.\n",
    "* `Top-k` - the number of highest probability vocabulary tokens to keep for top-k-filtering.\n",
    "* `Temperature` - the value used to module the logits distribution."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9f222d02-847a-490f-8d66-02608a53259b",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "if not Path(\"gradio_helper.py\").exists():\n",
    "    r = requests.get(url=\"https://raw.githubusercontent.com/openvinotoolkit/openvino_notebooks/latest/notebooks/llm-question-answering/gradio_helper.py\")\n",
    "    open(\"gradio_helper.py\", \"w\").write(r.text)\n",
    "\n",
    "from gradio_helper import make_demo\n",
    "\n",
    "demo = make_demo(run_fn=run_generation, title=f\"Question Answering with {model_id.value} and OpenVINO\")\n",
    "\n",
    "try:\n",
    "    demo.queue().launch(height=800)\n",
    "except Exception:\n",
    "    demo.queue().launch(share=True, height=800)\n",
    "# If you are launching remotely, specify server_name and server_port\n",
    "# EXAMPLE: `demo.launch(server_name='your server name', server_port='server port in int')`\n",
    "# To learn more please refer to the Gradio docs: https://gradio.app/docs/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "59038a29",
   "metadata": {},
   "outputs": [],
   "source": [
    "# please uncomment and 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"
  },
  "openvino_notebooks": {
   "imageUrl": "https://github.com/openvinotoolkit/openvino_notebooks/assets/29454499/daafd702-5a42-4f54-ae72-2e4480d73501",
   "tags": {
    "categories": [
     "Model Demos",
     "AI Trends"
    ],
    "libraries": [],
    "other": [
     "LLM"
    ],
    "tasks": [
     "Text Generation"
    ]
   }
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "state": {},
    "version_major": 2,
    "version_minor": 0
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}