mlc-llm/models/demo_CodeLlama_13b.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4IwhlCjVtpYj"
      },
      "source": [
        "# Demo: CodeLlama-13b with MLC LLM\n",
        "\n",
        "Recently, Meta unveiled [CodeLlama](https://github.com/facebookresearch/codellama), a family of large language models for code based on Llama 2 providing state-of-the-art performance among open models, infilling capabilities, support for large input contexts, and zero-shot instruction following ability for programming tasks. This notebook demonstrates MLC LLM's support for the CodeLlama family:\n",
        "\n",
        "- **[CodeLlama](https://huggingface.co/codellama/CodeLlama-13b-hf): a coding foundation LLM**\n",
        "- **[CodeLlama-Instruct](https://huggingface.co/codellama/CodeLlama-13b-Instruct-hf): an instruction-tuned LLM for coding**\n",
        "- **[CodeLlama-Python](https://huggingface.co/codellama/CodeLlama-13b-Python-hf): a Python specialized LLM**\n",
        "\n",
        "In this respect, MLC LLM allows everyone to develop, optimize and deploy AI models natively on everyone's devices. Therefore, making possible the deployment of coding LLMs natively, acting as **a personal AI coding assistant**.\n",
        "\n",
        "In this notebook, we walk over the steps of using MLC LLM to run these pre-compiled CodeLlama models! We have uploaded various versions of the pre-compiled and quantized CodeLlama models here: https://huggingface.co/mlc-ai.\n",
        "\n",
        "Learn more about MLC LLM here: https://mlc.ai/mlc-llm/docs."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RMyJ-y5DoB0S"
      },
      "source": [
        "Here's an overview regarding each model's capabilities:\n",
        "\n",
        "|                       | Code Completion | Infilling | Instruction/chat | Python specialist |\n",
        "|-----------------------|-----------------|-----------|------------------|-------------------|\n",
        "| CodeLlama-13b          |        X        |     X     |                  |                   |\n",
        "| CodeLlama-13b-Python   |        X        |           |                  |         X         |\n",
        "| CodeLlama-13b-Instruct |        X        |     X     |         X        |                   |"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YsvAL7SSt9Lo"
      },
      "source": [
        "Click the button below to get started!\n",
        "\n",
        "<a target=\"_blank\" href=\"https://colab.research.google.com/github/mlc-ai/notebooks/blob/main/mlc-llm/models/demo_CodeLlama_13b.ipynb\">\n",
        "  <img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/>\n",
        "</a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8kkADAMCCLi-"
      },
      "source": [
        "## Install MLC LLM"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Y2EwuS6TCO61"
      },
      "source": [
        "We will start from setting up the environment. First, let us create a new Conda environment, in which we will run the rest of the notebook.\n",
        "\n",
        "```\n",
        "conda create --name mlc-llm python=3.10\n",
        "conda activate mlc-llm\n",
        "```"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ojEeEmsqCTPG"
      },
      "source": [
        "**Google Colab**\n",
        "\n",
        "- If you are running this in a Google Colab notebook, you would not need to create a conda environment.\n",
        "- However, be sure to change your runtime to GPU by going to `Runtime` > `Change runtime type` and setting the Hardware accelerator to be \"GPU\"."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "S_rX53bGChPn"
      },
      "source": [
        "If you are using CUDA, you can run the following command to confirm that CUDA is set up correctly, and check the driver version number as well as what GPUs are currently available for use."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "CRPeCflbCij6"
      },
      "outputs": [],
      "source": [
        "!nvidia-smi"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PQfVfTAYC1M-"
      },
      "source": [
        "Next, let's download the MLC-AI and mlc-llm nightly build packages. If you are running in a Colab environment, then you can just run the following command. Otherwise, go to https://mlc.ai/package/ and replace the command below with the one that is appropriate for your hardware and OS."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vi-udt4tC5c9"
      },
      "source": [
        "**Google Colab**: If you are using Colab, you may see the red warnings such as \"You must restart the runtime in order to use newly installed versions.\" For our purpose, we can disregard them, the notebook will still run correctly."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ah9tYaCRCkKS"
      },
      "outputs": [],
      "source": [
        "!pip install --pre --force-reinstall mlc-ai-nightly-cu118 mlc-llm-nightly-cu118 -f https://mlc.ai/wheels"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nZGVNJE-DJ9E"
      },
      "source": [
        "Let's confirm we have installed the packages successfully!"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5Y6LszJgC7SQ"
      },
      "outputs": [],
      "source": [
        "!python -c \"import tvm; print('tvm installed properly!')\"\n",
        "!python -c \"import mlc_llm; print('mlc_llm installed properly!')\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lGfnrRa9DMw1"
      },
      "source": [
        "## Download Prebuilt Models and Library"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pVYkLb0eDjMi"
      },
      "source": [
        "The following commands will download all the available prebuilt libraries (e.g., `.so` files), including the precompiled CodeLlama models. This may take a while. If in **Google Colab**, you can verify that the files are being downloaded by clicking on the folder icon on the left."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Pg7daEvlD5UB"
      },
      "source": [
        "Note: If you are NOT running in **Google Colab** you may need to run this line `!conda install git git-lfs` to install `git` and `git-lfs` before running the following cell."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "FDFbw1KPDLu1"
      },
      "outputs": [],
      "source": [
        "!git lfs install"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bYqaVjmND7Px"
      },
      "outputs": [],
      "source": [
        "!mkdir -p dist/prebuilt\n",
        "!git clone https://github.com/mlc-ai/binary-mlc-llm-libs.git dist/prebuilt/lib"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SMEavWCJEC_d"
      },
      "source": [
        "#### CodeLlama-13b q4f16_1 prebuilt weights"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "etHEUrfMD8bX"
      },
      "outputs": [],
      "source": [
        "!cd dist/prebuilt && git clone https://huggingface.co/mlc-ai/mlc-chat-CodeLlama-13b-hf-q4f16_1"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FQkIX4TpElR6"
      },
      "source": [
        "#### CodeLlama-13b-Instruct q4f16_1 prebuilt weights"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "mTEGXAlhEnOw"
      },
      "outputs": [],
      "source": [
        "!cd dist/prebuilt && git clone https://huggingface.co/mlc-ai/mlc-chat-CodeLlama-13b-Instruct-hf-q4f16_1"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "acb1HpKpEoca"
      },
      "source": [
        "#### CodeLlama-13b-Python q4f16_1 prebuilt weights"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "oysLKcZ4Eou7"
      },
      "outputs": [],
      "source": [
        "!cd dist/prebuilt && git clone https://huggingface.co/mlc-ai/mlc-chat-CodeLlama-13b-Python-hf-q4f16_1"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dbHdyfIXHNpo"
      },
      "outputs": [],
      "source": [
        "# Restart colab\n",
        "exit()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CmpxrrqyE0S6"
      },
      "source": [
        "## Let's code with CodeLlama!"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "y73vNLy1OfMr"
      },
      "source": [
        "Let's first try a simple code completion task with the CodeLlama-Python."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "EOEf8sDyEwuv"
      },
      "outputs": [],
      "source": [
        "from mlc_llm import ChatModule\n",
        "from mlc_llm.callback import StreamToStdout"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tBcQMm-KJPN-"
      },
      "outputs": [],
      "source": [
        "codellama_python = ChatModule(model=\"CodeLlama-13b-Python-hf-q4f16_1\", device=\"cuda\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1g2l_hJhLyYm"
      },
      "outputs": [],
      "source": [
        "prompt = \"\"\"\\\n",
        "import argparse\n",
        "\n",
        "def main(string: str):\n",
        "    print(string)\n",
        "    print(string[::-1])\n",
        "\n",
        "if __name__ == \"__main__\":\"\"\"\n",
        "\n",
        "output = codellama_python.generate(\n",
        "    prompt=prompt,\n",
        "    progress_callback=StreamToStdout(callback_interval=2)\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "XwU54BtKQKz4",
        "outputId": "2ece3f04-970d-4baf-c4c3-a93cc7c80d08"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "import argparse\n",
            "\n",
            "def main(string: str):\n",
            "    print(string)\n",
            "    print(string[::-1])\n",
            "\n",
            "if __name__ == \"__main__\":\n",
            "    parser = argparse.ArgumentParser(description=\"Checks if the provided string is a palindrome.\"))\n",
            "    parser.add_argument(\"-s\", \"--string\",\n",
            "    help=\"The string to check.\"))\n",
            "\n",
            "    args = parser.parse_args()\n",
            "    main(args.string))\n"
          ]
        }
      ],
      "source": [
        "print(prompt+output)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DFrVdqewL7_c"
      },
      "outputs": [],
      "source": [
        "# Restart colab to initialize a new ChatModule\n",
        "exit()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6CUo34QeQto2"
      },
      "source": [
        "The CodeLlama models support infilling based on surrounding content. Let's try it with the foundation CodeLlama."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "yeHn8je9SBpK"
      },
      "outputs": [],
      "source": [
        "from mlc_llm import ChatModule\n",
        "from mlc_llm.callback import StreamToStdout\n",
        "\n",
        "def text_infilling(prompt: str):\n",
        "    prefix = prompt.split(\"<FILL>\")[0]\n",
        "    suffix = prompt.split(\"<FILL>\")[1]\n",
        "    return f\"<PRE> {prefix} <SUF>{suffix} <MID>\"\n",
        "\n",
        "def print_infilling(prompt: str, output: str):\n",
        "    print(prompt.replace(\"<FILL>\", output.replace(\"<EOT>\", \"\")))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "LUU8zTdRSKNK"
      },
      "outputs": [],
      "source": [
        "codellama = ChatModule(model=\"CodeLlama-13b-hf-q4f16_1\", device=\"cuda\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "d3rrm1qPSOkr"
      },
      "outputs": [],
      "source": [
        "prompt = \"\"\"\\\n",
        "# Installation instructions:\n",
        "    <FILL>\n",
        "This downloads the LLaMA inference code and installs the repository as a local pip package.\n",
        "\"\"\"\n",
        "\n",
        "output = codellama.generate(\n",
        "    prompt=text_infilling(prompt),\n",
        "    progress_callback=StreamToStdout(callback_interval=2)\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "78SHkqohUbCc",
        "outputId": "c842e0a8-1ddc-4dc2-8d02-2d1c264c3131"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "# Installation instructions:\n",
            "    pip install llamapy\n",
            "\n",
            "# Using the local pip package:\n",
            "\n",
            "    import llamapy\n",
            "    my_model = llamapy.LLaMA(n_components=2))\n",
            "\n",
            "# Requirements:\n",
            "\n",
            "    Python 3.x\n",
            "\n",
            "\n",
            "# Installation (easy way):\n",
            "\n",
            "    pip install git+https://github.com/BBIC-BBC/LLAMA\n",
            "\n",
            "\n",
            "# Installation (advanced way)):\n",
            "\n",
            "\n",
            "    1) Download the repository from Github:\n",
            "\n",
            "\n",
            "        git clone https://github.com/BBIC-BBC/LLAMA\n",
            "\n",
            "\n",
            "    2) Install the repository as a local pip package:\n",
            "\n",
            "\n",
            "        cd LLAMA\n",
            "\n",
            "\n",
            "\n",
            "        python setup.py install\n",
            "\n",
            "\n",
            "\n",
            "\n",
            "# Using the local pip package:\n",
            "\n",
            "\n",
            "    import llamapy\n",
            "    my_model = llamapy.LLaMA(n_components=2))))\n",
            "\n",
            "\n",
            "\n",
            "# Requirements:\n",
            "\n",
            "\n",
            "    Python 3.x\n",
            "\n",
            "\n",
            "\n",
            "# Installation (easy way):\n",
            "\n",
            "\n",
            "    pip install git+https://github.com/BBIC-BBC/LLAMA\n",
            "\n",
            "\n",
            "\n",
            "# Installation (advanced way)):\n",
            "\n",
            "\n",
            "\n",
            "    1) Download the LLaMA inference code from Github:\n",
            "\n",
            "\n",
            "        git clone https://github.com/BBIC-BBC/LLAMA\n",
            "\n",
            "\n",
            "    2) Install the LLaMA inference code as a local pip package:\n",
            "\n",
            "\n",
            "        cd LLaMA\n",
            "\n",
            "\n",
            "\n",
            "        python setup.py install\n",
            "This downloads the LLaMA inference code and installs the repository as a local pip package.\n",
            "\n"
          ]
        }
      ],
      "source": [
        "print_infilling(prompt, output)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "mKYyRyUVPGm2"
      },
      "outputs": [],
      "source": [
        "# Restart colab to create a new ChatModule\n",
        "exit()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CqmYlUtwV01m"
      },
      "source": [
        "Finally, the CodeLlama-Instruct has instruction following ability for programming tasks."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IcCPRp9oWBfh"
      },
      "outputs": [],
      "source": [
        "from mlc_llm import ChatModule\n",
        "from mlc_llm.callback import StreamToStdout"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "I4C8iC9IWyZR"
      },
      "outputs": [],
      "source": [
        "codellama_instruct = ChatModule(model=\"CodeLlama-13b-Instruct-hf-q4f16_1\", device=\"cuda\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "1kL02gKWW7Ov",
        "outputId": "65919ac1-2e1a-4545-ac3d-3643a7c4a18f"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Here is a possible implementation of the program:\n",
            "```\n",
            "import java.util.*;\n",
            "public class SumOfSublists {\n",
            "    public static void main(String[] args) {\n",
            "        List<Integer> list = Arrays.asList(1, 2, 3, 4, 5));\n",
            "        List<Integer> sums = new ArrayList<>();\n",
            "        for (int i = 0; i < list.size(); i++) {\n",
            "            int sum = 0;\n",
            "            for (int j = i; j < list.size(); j++) {\n",
            "                sum += list.get(j));\n",
            "            }\n",
            "\n",
            "            sums.add(sum));\n",
            "        }\n",
            "\n",
            "\n",
            "        System.out.println(\"The sums of all contiguous sublists are: \" + sums));\n",
            "    }\n"
          ]
        }
      ],
      "source": [
        "prompt = (\"Write a Java program that computes the set of sums of all contiguous\"\n",
        "          \"sublists of a given list.\")\n",
        "\n",
        "output = codellama_instruct.generate(\n",
        "    prompt=prompt,\n",
        "    progress_callback=StreamToStdout(callback_interval=2)\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dl8cowtUz5yp"
      },
      "outputs": [],
      "source": [
        "codellama_instruct.reset_chat()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "liWTlsBBYTRa",
        "outputId": "0edabb88-8219-448a-ebf6-66fd86767997"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Here is a program in Python that solves the problem of finding the indices of two numbers in an array that add up to a target value:\n",
            "```\n",
            "def find_indices(nums, target):\n",
            "    # Initialize two empty lists to store the indices of the two numbers\n",
            "    for i in range(len(nums)))):\n",
            "        for j in range(len(nums)))):\n",
            "            if i != j and nums[i] + nums[j]] == target:\n",
            "                indices = [i, j]]\n",
            "    return indices\n"
          ]
        }
      ],
      "source": [
        "prompt = (\"Given an array of integers nums and an integer target, return\"\n",
        "          \"indices of the two numbers such that they add up to target.\"\n",
        "          \" Write this program in Python.\")\n",
        "\n",
        "output = codellama_instruct.generate(\n",
        "    prompt=prompt,\n",
        "    progress_callback=StreamToStdout(callback_interval=2)\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "UzIMwuiaWwSg"
      },
      "outputs": [],
      "source": [
        "# Restart colab to create a new ChatModule\n",
        "exit()"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "gpuType": "T4",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}