{ "cells": [ { "cell_type": "markdown", "id": "d2b361bf", "metadata": {}, "source": [ "Copyright (c) 2023 Graphcore Ltd. All rights reserved." ] }, { "cell_type": "markdown", "id": "af14ce5a", "metadata": {}, "source": [ "# A Chatbot on IPUs using OpenAssistant Pythia 12B - Inference\n", "\n", "| Domain | Tasks | Model | Datasets | Workflow | Number of IPUs | Execution time |\n", "|---------|-------|-------|----------|----------|--------------|--------------|\n", "| NLP | Instruction Fine-tuned Chatbot | oasst-sft-4-pythia-12b-epoch-3.5 | N/A | Inference | recommended: 16 (minimum 4) | 30mn |\n", "\n", "The OpenAssistant `oasst-sft-4-pythia-12b-epoch-3.5` model is a chatbot trained on the OpenAssistant Conversations Dataset (OASST1).\n", "It is one of the best open-source and commercially usable chatbots currently available.\n", "\n", "Because it uses the same underlying Pythia 12B model as Dolly, we run it using the Dolly pipeline.\n", "\n", "In this notebook you will:\n", "- Create and configure an OpenAssistant inference pipeline.\n", "- Run the pipeline and use it to build up a chat conversation turn by turn.\n", "\n", "This notebook requires a minimum of 4 IPUs to run. This notebook also supports running on a POD16, resulting in a faster pipeline inference speed.\n", "\n", "[![Join our Slack Community](https://img.shields.io/badge/Slack-Join%20Graphcore's%20Community-blue?style=flat-square&logo=slack)](https://www.graphcore.ai/join-community)" ] }, { "cell_type": "markdown", "id": "db890f52", "metadata": {}, "source": [ "## Environment setup\n", "\n", "\n", "The best way to run this demo is on Paperspace Gradient's cloud IPUs because everything is already set up for you.\n", "\n", "\n", "To run the demo using other IPU hardware, you need to have the Poplar SDK enabled. Refer to the [Getting Started guide](https://docs.graphcore.ai/en/latest/getting-started.html#getting-started) for your system for details on how to enable the Poplar SDK. Also refer to the [Jupyter Quick Start guide](https://docs.graphcore.ai/projects/jupyter-notebook-quick-start/en/latest/index.html) for how to set up Jupyter to be able to run this notebook on a remote IPU machine.\n", "\n", "Run the next cell to install extra requirements for this notebook." ] }, { "cell_type": "code", "execution_count": null, "id": "5c72d24a", "metadata": {}, "outputs": [], "source": [ "%pip install -r requirements.txt\n", "%load_ext graphcore_cloud_tools.notebook_logging.gc_logger" ] }, { "cell_type": "code", "execution_count": null, "id": "965ca0a6", "metadata": {}, "outputs": [], "source": [ "import os\n", "\n", "number_of_ipus = int(os.getenv(\"NUM_AVAILABLE_IPU\", 16))\n", "number_of_ipus" ] }, { "cell_type": "markdown", "id": "7aa108da", "metadata": {}, "source": [ "## OpenAssistant inference pipeline\n", "\n", "Let's begin by loading the inference config. We use the same configuration file as Dolly and manually modify the vocab size\n", "which is the only difference between the model graphs. A configuration suitable for your instance will automatically be selected." ] }, { "cell_type": "code", "execution_count": null, "id": "a5aa39b7", "metadata": {}, "outputs": [], "source": [ "from utils.setup import dolly_config_setup\n", "\n", "config_name = \"dolly_pod4\" if number_of_ipus == 4 else \"dolly_pod16\"\n", "config, *_ = dolly_config_setup(\"config/inference.yml\", \"release\", config_name)\n", "\n", "# Update vocab size for oasst-sft-4-pythia-12b-epoch-3.5 - 50288 rather than 50280\n", "config.model.embedding.vocab_size = 50288\n", "\n", "config" ] }, { "cell_type": "markdown", "id": "7e65adf1", "metadata": {}, "source": [ "Next, we want to create our inference pipeline. Here we define the maximum\n", "sequence length and maximum micro batch size. Before executing a model on IPUs\n", "it needs to be turned into an executable format by compiling it. This will\n", "happen when the pipeline is created. All input shapes must be known before\n", "compiling, so if the maximum sequence length or micro batch size is changed, the\n", "pipeline will need to be recompiled.\n", "\n", "Selecting a longer sequence length or larger batch size will use more IPU\n", "memory. This means that increasing one may require you to decrease the other.\n", "\n", "*This cell will take approximately 18 minutes to complete, which includes downloading the model weights.*" ] }, { "cell_type": "code", "execution_count": null, "id": "21144e73", "metadata": {}, "outputs": [], "source": [ "import api\n", "\n", "# changing these parameters will trigger a recompile.\n", "sequence_length = 512 # max 2048\n", "micro_batch_size = 1\n", "\n", "# The pipeline is set to load the OpenAssistant checkpoint rather than the default Dolly one\n", "\n", "# We override the Dolly prompt templating by specifying a prompt_format. Setting the format to\n", "# just echo the instruction means that the pipeline does no formatting and it is up to the\n", "# application to provide correctly templated prompts\n", "\n", "# We set the text string that OpenAssistant uses to mark that it has finished generation, which\n", "# is different from the Dolly one\n", "\n", "oasst_pipeline = api.DollyPipeline(\n", " config,\n", " sequence_length=sequence_length,\n", " micro_batch_size=micro_batch_size,\n", " hf_dolly_checkpoint=\"OpenAssistant/oasst-sft-4-pythia-12b-epoch-3.5\",\n", " prompt_format=\"{instruction}\",\n", " end_key=\"<|endoftext|>\",\n", ")" ] }, { "cell_type": "markdown", "id": "15a35913-ffcf-4433-999e-824136caf421", "metadata": {}, "source": [ "Call the `oasst_pipeline` object you have just created to generate text from a prompt." ] }, { "cell_type": "markdown", "id": "872d03f6-dd5a-4510-8205-8b4a990e8df9", "metadata": {}, "source": [ "To make a chatbot, we will take user input and feed it to the model. So that the notebook can be tested\n", "automatically, we create a function similar to the Python built-in function `input()` that collects input\n", "from the user and returns it, but which will return canned input from a test environment variable\n", "`EXAMPLE_PROMPTS` instead if that variable is set. The variable should be set to a JSON list of strings -\n", "for example:\n", "\n", "```\n", "export EXAMPLE_PROMPTS='[\"Are you related to Dolly?\", \"Ah OK. How many islands are there in Scotland?\"]'\n", "```" ] }, { "cell_type": "code", "execution_count": null, "id": "b66046cd-b6b2-4826-b883-6e6b000c4564", "metadata": {}, "outputs": [], "source": [ "import os\n", "import json\n", "\n", "example_prompts = os.getenv(\"EXAMPLE_PROMPTS\")\n", "if example_prompts is not None:\n", " example_prompts = json.loads(example_prompts)\n", "\n", "# Get input from user like input() but with possibility of automation via EXAMPLE_PROMPTS environment variable\n", "def auto_input(prompt: str) -> str:\n", " if example_prompts is None:\n", " return input(prompt)\n", " auto_in = example_prompts.pop(0) if example_prompts != [] else \"\"\n", " print(prompt, \"AUTO\", auto_in)\n", " return auto_in" ] }, { "cell_type": "markdown", "id": "b5ae263d", "metadata": {}, "source": [ "A chatbot conversation is built up from a number of turns of user input and the model writing a reply.\n", "As a conversation develops, the prompt should be extended turn by turn, so the model has access to the full\n", "context.\n", "\n", "The model has been trained on a specific prompt template to represent the conversation as it is built up:\n", "\n", "\n", "> *`<|prompter|>`**user1**`<|endoftext|>` `<|assistant|>`**reply1**`<|endoftext|>` `<|prompter|>`**user2**`<|endoftext|>` `<|assistant|>`...*\n", "\n", "\n", "There are a some optional parameters to the pipeline call you can use to control the generation behaviour:\n", "- `temperature` – Indicates whether you want more or less creative output. A value of 1.0 corresponds to the model's default behaviour. Smaller values than this accentuate the next token distribution and make the model more likely to pick a highly probable next token. A value of 0.0 means the model will always pick the most probable token. Temperatures greater than 1.0 flatten the next token distribution making more unusual next tokens more likely. Temperature must be zero or positive.\n", "- `k` – Indicates that only among the highest `k` probable tokens can be sampled. This is known as \"top k\" sampling. Set to 0 to disable top k sampling and sample from all possible tokens. The value for `k` must be between a minimum of 0 and a maximum of `config.model.embedding.vocab_size` which is 50,288. The default is 5.\n", "- `output_length` - Sets a maximum output length in tokens. Generation normally stops when the model generates its end_key text, but can be made to stop before that by specifying this option. A value of `None` disables the limit. The default is 'None'.\n", "\n", "You can start with any user input. For instance \"**What other animals are similar to Alpacas?**\"" ] }, { "cell_type": "code", "execution_count": null, "id": "f9e87e5f", "metadata": {}, "outputs": [], "source": [ "import logging\n", "\n", "# Conduct a complete conversation - with ability to set pipeline optional parameters\n", "def chat(temperature=None, top_k=None, output_length=None):\n", "\n", " options = {}\n", " if temperature is not None:\n", " options[\"temperature\"] = temperature\n", " if top_k is not None:\n", " options[\"k\"] = top_k\n", " if output_length is not None:\n", " options[\"output_length\"] = output_length\n", "\n", " # Suppress INFO logging to make a better interactive chat display\n", " logging.disable(logging.INFO)\n", "\n", " print(\"To complete the chat, enter an empty prompt\")\n", "\n", " prompt = \"\"\n", " while True:\n", " user_input = auto_input(\"Prompter:\")\n", " if user_input == \"\":\n", " break\n", "\n", " prompt += f\"<|prompter|>{user_input}<|endoftext|><|assistant|>\"\n", " chat_output = oasst_pipeline(prompt, **options)[0]\n", " prompt += chat_output + \"<|endoftext|>\"\n", "\n", " # Restore logging to what it was before\n", " logging.disable(logging.NOTSET)\n", "\n", "\n", "chat(temperature=0.0)\n", "\n", "# Note the first time you run this cell and enter a prompt, there will be a delay of ~1 minute where\n", "# nothing appears to happen. This is where the server is attaching to the IPUs" ] }, { "cell_type": "markdown", "id": "f2ca473c", "metadata": {}, "source": [ "Remember to detach your pipeline when you are finished to free up resources:" ] }, { "cell_type": "code", "execution_count": null, "id": "d9e91b9a", "metadata": {}, "outputs": [], "source": [ "oasst_pipeline.detach()" ] }, { "cell_type": "markdown", "id": "7ece4041", "metadata": {}, "source": [ "## Conclusion\n", "\n", "In this notebook, we have demonstrated how you can easily run an OpenAssistant chatbot on Graphcore IPUs. Instruction fine-tuning is\n", "a powerful method of turning a base LLM into one more suited for human interactivity, such as a question-answer model or a chatbot.\n", "\n", "Although larger instruction fine-tuned LLMs exist with more world knowledge such as ChatGPT, they are closed-source or are subject to\n", "non-commercial licensing. OpenAssistant is currently one of the best open-source and commercially usable chatbot models. Have fun! " ] }, { "cell_type": "markdown", "id": "73e9821a", "metadata": {}, "source": [ "## Next steps\n", "\n", "Check out the full list of [IPU-powered Jupyter Notebooks](https://www.graphcore.ai/ipu-jupyter-notebooks) to see how IPUs perform on other tasks." ] } ], "metadata": { "language_info": { "name": "python" } }, "nbformat": 4, "nbformat_minor": 5 }