notebooks/273-stable-zephyr-3b-chatbot/273-stable-zephyr-3b-chatbot.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "a1dafb31-bac1-4d01-9356-102114115c51",
   "metadata": {},
   "source": [
    "# LLM-powered chatbot using Stable-Zephyr-3b and OpenVINO\n",
    "\n",
    "In the rapidly evolving world of artificial intelligence (AI), chatbots have become 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",
    "`Stable Zephyr 3B` is a 3 billion parameter model that demonstrated outstanding results on many LLM evaluation benchmarks outperforming many popular models in relatively small size. Inspired by [HugginFaceH4's Zephyr 7B](https://huggingface.co/HuggingFaceH4/zephyr-7b-beta) training pipeline this model was trained on a mix of publicly available datasets, synthetic datasets using [Direct Preference Optimization (DPO)](https://arxiv.org/abs/2305.18290), evaluation for this model based on [MT Bench](https://tatsu-lab.github.io/alpaca_eval/) and [Alpaca Benchmark](https://tatsu-lab.github.io/alpaca_eval/). More details about model can be found in [model card](https://huggingface.co/stabilityai/stablelm-zephyr-3b)\n",
    "\n",
    "In this tutorial, we consider how to optimize and run this model using the OpenVINO toolkit. For the convenience of the conversion step and model performance evaluation, we will use [llm_bench](https://github.com/openvinotoolkit/openvino.genai/tree/master/llm_bench/python) tool, which provides a unified approach to estimate performance for LLM. It is based on pipelines provided by Optimum-Intel and allows to estimate performance for Pytorch and OpenVINO models using almost the same code. We also demonstrate how to apply BetterTransformer optimization and making model stateful, that provides opportunity for processing model cache state."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "93504605-33cc-4eb1-8331-dd7d7ea9a51d",
   "metadata": {},
   "source": [
    "## Prerequisites\n",
    "\n",
    "For starting work, we should install required packages first"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dd1d538e-c22b-4a6e-9f5b-4462b1a2a43f",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "import sys\n",
    "\n",
    "\n",
    "genai_llm_bench = Path(\"openvino.genai/llm_bench/python\")\n",
    "\n",
    "if not genai_llm_bench.exists():\n",
    "    !git clone  https://github.com/openvinotoolkit/openvino.genai.git\n",
    "    %cd openvino.genai\n",
    "    !git checkout e5d7861\n",
    "    %cd ..\n",
    "\n",
    "sys.path.append(str(genai_llm_bench))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "342d2a0e-08cb-4a65-9164-98cb82f25d8c",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install -q --extra-index-url https://download.pytorch.org/whl/cpu -r ./openvino.genai/llm_bench/python/requirements.txt\n",
    "%pip uninstall -q -y openvino-dev openvino openvino-nightly\n",
    "%pip install openvino-nightly"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eb0882f4-e206-4dde-a67e-67d40b73250b",
   "metadata": {},
   "source": [
    "## Convert model to OpenVINO Intermediate Representation (IR) and compress model weights to INT4 using NNCF\n",
    "\n",
    "llm_bench provides conversion script for converting LLMS into OpenVINO IR format compatible with Optimum-Intel. It also allows to compress model weights into INT8 or INT4 precision with [NNCF](https://github.com/openvinotoolkit/nncf). For enabling weights compression in INT4 we should use `--compress_weights 4BIT_DEFAULT` argument. 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. Additionally, it support way to optimize models using [BetterTransformer](https://huggingface.co/docs/optimum/bettertransformer) library, a fast path of standard PyTorch Transformer APIs to benefit from interesting speedups on CPU & GPU through sparsity and fused kernels as Flash Attention using `--bettertransformer` flag."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "aeacd53b-755a-464f-b984-d88c3625d687",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_path = Path(\"stable-zephyr-3b/pytorch/dldt/compressed_weights/OV_FP16-4BIT_DEFAULT\") \n",
    "\n",
    "convert_script = genai_llm_bench / \"convert.py\"\n",
    "\n",
    "if not (model_path / \"openvino_model.xml\").exists():\n",
    "    !python $convert_script --model_id stabilityai/stable-zephyr-3b --precision FP16 --compress_weights 4BIT_DEFAULT --bettertransformer --output stable-zephyr-3b"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "98c8167c-d17c-43e8-9ff9-a9aa2f16cf3f",
   "metadata": {},
   "source": [
    "## Estimate model performance\n",
    "\n",
    "openvino.genai / llm_bench / python / benchmark.py script allow to estimate text generation pipeline inference on specific input prompt with given number of maximum generated tokens."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "f57f0ac8-74ea-4709-a1a8-eb8c4d6da3fc",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:nncf:NNCF initialized successfully. Supported frameworks detected: torch, onnx, openvino\n",
      "No CUDA runtime is found, using CUDA_HOME='/usr/local/cuda'\n",
      "[ INFO ] ==SUCCESS FOUND==: use_case: text_gen, model_name: stable-zephyr-3b\n",
      "[ INFO ] ov_config={'PERFORMANCE_HINT': 'LATENCY', 'CACHE_DIR': '', 'NUM_STREAMS': '1'}\n",
      "OPENVINO_TORCH_BACKEND_DEVICE=CPU\n",
      "[ INFO ] model_path=stable-zephyr-3b/pytorch/dldt/compressed_weights/OV_FP16-4BIT_DEFAULT, openvino runtime version:2023.3.0-13522-0a7d1d770ff\n",
      "Compiling the model to CPU ...\n",
      "[ INFO ] From pretrained time: 5.11s\n",
      "Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n",
      "[ INFO ] num_iters=0, num_text_list=1\n",
      "[ INFO ] input_text=Tell me story about cats\n",
      "[ INFO ] Input token size:5, max_output_token_size:512\n",
      "Setting `pad_token_id` to `eos_token_id`:0 for open-end generation.\n",
      "/home/ea/work/openvino_notebooks/notebooks/273-stable-zephyr-3b-chatbot/openvino.genai/llm_bench/python/utils/ov_model_classes.py:998: 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",
      "[ INFO ] [warm-up] Input token size: 5\n",
      "[ INFO ] [warm-up] Output size: 512\n",
      "[ INFO ] [warm-up] Infer count: 512\n",
      "[ INFO ] [warm-up] Generation Time: 33.60s\n",
      "[ INFO ] [warm-up] Latency: 65.62 ms/token\n",
      "[ INFO ] [warm-up] Generated:\n",
      "Tell me story about cats and dogs\n",
      "Once upon a time, in a small village, there lived a kind and gentle old lady named Grandma.\n",
      "Grandma had a small garden behind her house, where she grew beautiful flowers and vegetables. She loved to spend her days in the garden, tending to her plants and enjoying the peaceful surroundings.\n",
      "One day, a curious little kitten wandered into the garden. It was playful and curious, and it seemed to be friendly towards Grandma. The kitten was so cute that Grandma couldn't resist it and decided to keep it. From that day on, the garden was filled with the joyful sounds of the kitten playing and exploring.\n",
      "One day, while Grandma was watering the plants, she heard a loud barking coming from the other side of the garden. It was a big, friendly dog who had wandered into the village. The dog was looking for someone to help him find his owner, as he had been separated from his family during a move.\n",
      "Grandma felt sorry for the dog and decided to take care of him. She gave him some food and water and made him a cozy bed in the shed. The dog was grateful for the kind treatment and soon became a regular visitor in the garden.\n",
      "The kitten and the dog soon became good friends, playing and exploring together. They would often take long walks through the village, looking for new adventures to share. The kitten would chase after the dog, and the dog would bring the kitten back little treasures he found along the way.\n",
      "One day, while they were walking, they heard a loud noise coming from the other side of the garden. It was a group of stray cats who had come to visit the village. The cats were sneaky and mischievous, and they loved to play tricks on the villagers.\n",
      "The kitten and the dog were both scared at first, but they soon realized that the cats were not a threat to them. The kitten and the dog decided to make friends with the cats and show them that they were friendly too.\n",
      "The cats were amazed by the kitten's playful nature and the dog's friendly demeanor. They soon became good friends with the kitten and the dog, and the garden became a lively and colorful place once again.\n",
      "From that day on, the kitten, the dog, and the cats would often come to visit Grandma in the garden, bringing with them new adventures and memories to cherish. And Grandma would always be happy to see\n",
      "[ INFO ] [warm-up] Result MD5:['ff3143d4365f5fd1a308832647e3e565']\n",
      "[ INFO ] [warm-up] First token latency: 821.26 ms/token, other tokens latency: 64.11 ms/token, len of tokens: 512\n",
      "[ INFO ] [warm-up] First token infer latency: 820.47 ms/token, other tokens infer latency: 63.59 ms/token, len of tokens: 512\n"
     ]
    }
   ],
   "source": [
    "benchmark_script = genai_llm_bench / \"benchmark.py\"\n",
    "\n",
    "!python $benchmark_script -m $model_path -ic 512 -p \"Tell me story about cats\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1ac23ffe-a3ef-4c81-b193-a51bf3e392b7",
   "metadata": {},
   "source": [
    "## Apply stateful transformation for automatic handling model state\n",
    "\n",
    "Stable Zephyr is a decoder-only transformer model and generates text token by token in an autoregressive fashion. 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.  To optimize the generation process and use memory more efficiently, HuggingFace transformers API provides a mechanism for caching model state externally using `use_cache=True` parameter and `past_key_values` argument in inputs and outputs. 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. 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",
    "With increasing model size like in modern LLMs, we also can note an increase in the number of attention blocks and size past key values tensors respectively. The strategy for handling cache state as model inputs and outputs in the inference cycle may become a bottleneck for memory-bounded systems, especially with processing long input sequences, for example in a chatbot scenario. OpenVINO suggests a transformation that removes inputs and corresponding outputs with cache tensors from the model keeping cache handling logic inside the model. Hiding the cache enables storing and updating the cache values in a more device-friendly representation. It helps to reduce memory consumption and additionally optimize model performance.\n",
    "\n",
    "You can estimate the model performance by adding stateful transformation using `--make_stateful` flag."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "a586d0be-bb8e-449f-8dd5-61bc4f0a2e20",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:nncf:NNCF initialized successfully. Supported frameworks detected: torch, onnx, openvino\n",
      "No CUDA runtime is found, using CUDA_HOME='/usr/local/cuda'\n",
      "[ INFO ] ==SUCCESS FOUND==: use_case: text_gen, model_name: stable-zephyr-3b\n",
      "[ INFO ] ov_config={'PERFORMANCE_HINT': 'LATENCY', 'CACHE_DIR': '', 'NUM_STREAMS': '1'}\n",
      "OPENVINO_TORCH_BACKEND_DEVICE=CPU\n",
      "[ INFO ] model_path=stable-zephyr-3b/pytorch/dldt/compressed_weights/OV_FP16-4BIT_DEFAULT, openvino runtime version:2023.3.0-13522-0a7d1d770ff\n",
      "Compiling the model to CPU ...\n",
      "[ INFO ] From pretrained time: 7.12s\n",
      "Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n",
      "[ INFO ] num_iters=0, num_text_list=1\n",
      "[ INFO ] input_text=Tell me story about cats\n",
      "[ INFO ] Input token size:5, max_output_token_size:512\n",
      "Setting `pad_token_id` to `eos_token_id`:0 for open-end generation.\n",
      "[ INFO ] [warm-up] Input token size: 5\n",
      "[ INFO ] [warm-up] Output size: 512\n",
      "[ INFO ] [warm-up] Infer count: 512\n",
      "[ INFO ] [warm-up] Generation Time: 27.29s\n",
      "[ INFO ] [warm-up] Latency: 53.30 ms/token\n",
      "[ INFO ] [warm-up] Generated:\n",
      "Tell me story about cats and dogs\n",
      "Once upon a time, in a small village, there lived a kind and gentle old lady named Grandma.\n",
      "Grandma had a small garden behind her house, where she grew beautiful flowers and vegetables. She loved to spend her days in the garden, tending to her plants and enjoying the peaceful surroundings.\n",
      "One day, a curious little kitten wandered into the garden. It was playful and curious, and it seemed to be friendly towards Grandma. The kitten was so cute that Grandma couldn't resist it and decided to keep it. From that day on, the garden was filled with the joyful sounds of the kitten playing and exploring.\n",
      "One day, while Grandma was watering the plants, she heard a loud barking coming from the other side of the garden. It was a big, friendly dog who had wandered into the village. The dog was looking for someone to help him find his owner, as he had been separated from his family during a move.\n",
      "Grandma felt sorry for the dog and decided to take care of him. She gave him some food and water and made him a cozy bed in the shed. The dog was grateful for the kind treatment and soon became a regular visitor in the garden.\n",
      "The kitten and the dog soon became good friends, playing and exploring together. They would often take long walks through the village, looking for new adventures to share. The kitten would chase after the dog, and the dog would bring the kitten back little treasures he found along the way.\n",
      "One day, while they were walking, they heard a loud noise coming from the other side of the garden. It was a group of stray cats who had come to visit the village. The cats were sneaky and mischievous, and they loved to play tricks on the villagers.\n",
      "The kitten and the dog were both scared at first, but they soon realized that the cats were not a threat to them. The kitten and the dog decided to make friends with the cats and show them that they were friendly too.\n",
      "The cats were amazed by the kitten's playful nature and the dog's friendly demeanor. They soon became good friends with the kitten and the dog, and the garden became a lively and colorful place once again.\n",
      "From that day on, the kitten, the dog, and the cats would often come to visit Grandma in the garden, bringing with them new adventures and memories to cherish. And Grandma would always be happy to see\n",
      "[ INFO ] [warm-up] Result MD5:['ff3143d4365f5fd1a308832647e3e565']\n",
      "[ INFO ] [warm-up] First token latency: 699.40 ms/token, other tokens latency: 51.99 ms/token, len of tokens: 512\n",
      "[ INFO ] [warm-up] First token infer latency: 698.34 ms/token, other tokens infer latency: 51.41 ms/token, len of tokens: 512\n"
     ]
    }
   ],
   "source": [
    "!python $benchmark_script -m $model_path -ic 512 -p \"Tell me story about cats\" --make_stateful "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b32d0beb-b564-47f0-99f0-4e11e4a68678",
   "metadata": {},
   "source": [
    "## Using model with Optimum Intel\n",
    "\n",
    "Running model with Optimum-Intel API required following steps:\n",
    "1. register normalized config for model\n",
    "2. create instance of `OVModelForCausalLM` class using `from_pretrained` method and providing path to the model.\n",
    "3. patch inter processing for applying `make_stateful=True`\n",
    "\n",
    "The model text generation interface remains without changes, the text generation process started with running `ov_model.generate` method and passing text encoded by the tokenizer as input. This method returns a sequence of generated token ids that should be decoded using a tokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "433988ee-bda7-4224-9bf5-b013de4fcd65",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:nncf:NNCF initialized successfully. Supported frameworks detected: torch, onnx, openvino\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "No CUDA runtime is found, using CUDA_HOME='/usr/local/cuda'\n",
      "Compiling the model to CPU ...\n",
      "Setting OpenVINO CACHE_DIR to stable-zephyr-3b/pytorch/dldt/compressed_weights/OV_FP16-4BIT_DEFAULT/model_cache\n"
     ]
    }
   ],
   "source": [
    "from utils.ov_model_classes import register_normalized_configs\n",
    "from utils.ov_utils import patch_inter_processing\n",
    "from optimum.intel.openvino import OVModelForCausalLM\n",
    "from transformers import AutoConfig\n",
    "\n",
    "# Load model into Optimum Interface\n",
    "register_normalized_configs()\n",
    "\n",
    "ov_model = OVModelForCausalLM.from_pretrained(model_path, compile=False, config=AutoConfig.from_pretrained(model_path, trust_remote_code=True))\n",
    "patch_inter_processing(ov_model, fuse_decoding_strategy=False, fuse_cache_reorder=False, make_stateful=True, save_prepared_model=None)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ef1e67fa-422c-4b7b-b3dd-cad81318b9b2",
   "metadata": {},
   "source": [
    "## Interactive chatbot demo\n",
    "\n",
    "Now, our model ready to use. Let's see it in action. We will use Gradio interface for interaction with model. Put text message into `Chat message box` and click `Submit` button for starting conversation.\n",
    "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.\n",
    "\n",
    "You can modify them in `Advanced generation options` section."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b7f014a4-2b3b-41fc-aae1-b4b0729b978d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from threading import Event, Thread\n",
    "from uuid import uuid4\n",
    "from typing import List, Tuple\n",
    "import gradio as gr\n",
    "from transformers import (\n",
    "    AutoTokenizer,\n",
    "    StoppingCriteria,\n",
    "    StoppingCriteriaList,\n",
    "    TextIteratorStreamer,\n",
    ")\n",
    "\n",
    "\n",
    "model_name = \"stable-zephyr-3b\"\n",
    "\n",
    "tok = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)\n",
    "\n",
    "DEFAULT_SYSTEM_PROMPT = \"\"\"\\\n",
    "You are a helpful, respectful and honest assistant. Always answer as helpfully as possible, while being safe.  Your answers should not include any harmful, unethical, racist, sexist, toxic, dangerous, or illegal content. Please ensure that your responses are socially unbiased and positive in nature.\n",
    "If a question does not make any sense or is not factually coherent, explain why instead of answering something not correct. If you don't know the answer to a question, please don't share false information.\\\n",
    "\"\"\"\n",
    "\n",
    "model_configuration = {\n",
    "    \"start_message\": f\"<|system|>\\n {DEFAULT_SYSTEM_PROMPT }<|endoftext|>\",\n",
    "    \"history_template\": \"<|user|>\\n{user}<|endoftext|><|assistant|>\\n{assistant}<|endoftext|>\",\n",
    "    \"current_message_template\": '<|user|>\\n{user}<|endoftext|><|assistant|>\\n{assistant}',\n",
    "}\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",
    "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",
    "    [\n",
    "        \"Write a 100-word blog post on “Benefits of Artificial Intelligence and OpenVINO“\"\n",
    "    ],\n",
    "]\n",
    "\n",
    "max_new_tokens = 256\n",
    "\n",
    "\n",
    "class StopOnTokens(StoppingCriteria):\n",
    "    def __init__(self, token_ids):\n",
    "        self.token_ids = token_ids\n",
    "\n",
    "    def __call__(\n",
    "        self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs\n",
    "    ) -> bool:\n",
    "        for stop_id in self.token_ids:\n",
    "            if input_ids[0][-1] == stop_id:\n",
    "                return True\n",
    "        return False\n",
    "\n",
    "\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",
    "\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",
    "\n",
    "text_processor = model_configuration.get(\n",
    "    \"partial_text_processor\", default_partial_text_processor\n",
    ")\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",
    "                [history_template.format(num=round, user=item[0], assistant=item[1])]\n",
    "            )\n",
    "            for round, item in enumerate(history[:-1])\n",
    "        ]\n",
    "    )\n",
    "    text += \"\".join(\n",
    "        [\n",
    "            \"\".join(\n",
    "                [\n",
    "                    current_message_template.format(\n",
    "                        num=len(history) + 1,\n",
    "                        user=history[-1][0],\n",
    "                        assistant=history[-1][1],\n",
    "                    )\n",
    "                ]\n",
    "            )\n",
    "        ]\n",
    "    )\n",
    "    return text\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(\n",
    "        tok, timeout=30.0, skip_prompt=True, skip_special_tokens=True\n",
    "    )\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(f\"\"\"<h1><center>OpenVINO {model_name} Chatbot</center></h1>\"\"\")\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",
    "        examples, inputs=msg, 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()"
   ]
  }
 ],
 "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
}