llm_deployment/RLChina24_4_LLM_Deployment.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0u4Y_5mZ5IQJ"
      },
      "source": [
        "# Full Stack Practice of LLM Training - LLM Depolyment @ RLChina 2024\n",
        "\n",
        "- Author: [Cheng Deng](https://www.cdeng.net/)[✉️]((davendw49@gmail.com), [Jun Wang](http://www0.cs.ucl.ac.uk/staff/jun.wang/)\n",
        "\n",
        "---\n",
        "## Main Task\n",
        "\n",
        "Given the high cost of pretraining, we'll load the pre-trained weights into our custom-built architecture. This section will cover how to load pretrained weights from models like Llama, Phi, Gemma, and Mistral or customer LLM using related framework like Transformers, vLLM, and llama.c library. Before advanced knowledge in LLM deployment, we will go throught the main methods in LLM evaluation.\n",
        "\n",
        "![](https://www.cdeng.net/resources/imgs/RLChina24/d.png)\n",
        "\n",
        "Here is the prerequisite knowledge required:\n",
        "\n",
        "- `Transformers` model\n",
        "- Basic OS knowledge\n",
        "\n",
        "Now, let's start with the preparation like package intallation, utilities import and package import."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "igKgsWXeErDE"
      },
      "outputs": [],
      "source": [
        "!pip install vllm"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iCi6LYcUEwRz"
      },
      "source": [
        "And then import the essential packages."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6PUdaAAP1sco"
      },
      "outputs": [],
      "source": [
        "from transformers import AutoModelForCausalLM, AutoTokenizer"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fSVPbbaNAbHH"
      },
      "source": [
        "## 0. Warm-up\n",
        "\n",
        "First, let's warm-up a normal model loading using Hugging Face `Transformers` library."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5rW0au7UAZpq"
      },
      "outputs": [],
      "source": [
        "checkpoint = \"HuggingFaceTB/SmolLM-135M\"\n",
        "device = \"cuda\" # for GPU usage or \"cpu\" for CPU usage\n",
        "tokenizer = AutoTokenizer.from_pretrained(checkpoint)\n",
        "# for multiple GPUs install accelerate and do `model = AutoModelForCausalLM.from_pretrained(checkpoint, device_map=\"auto\")`\n",
        "model = AutoModelForCausalLM.from_pretrained(checkpoint).to(device)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_0mgnFYFRcG2"
      },
      "outputs": [],
      "source": [
        "del model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "oC7X2oefA8Lw",
        "outputId": "5dfd4ec5-5e98-4ead-a964-2ceff0188ced"
      },
      "outputs": [],
      "source": [
        "inputs = tokenizer.encode(\"def print_hello_world():\", return_tensors=\"pt\").to(device)\n",
        "outputs = model.generate(inputs)\n",
        "print(tokenizer.decode(outputs[0]))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BDCbcU6tYWsY"
      },
      "source": [
        "## 1. Model Evaluation\n",
        "\n",
        "### Introduction to LLM Evaluation\n",
        "\n",
        "#### Dataset Preparation\n",
        "   - **Multiple-Choice Format**: Create a dataset where each question has a set of possible answers (usually labeled as A, B, C, D, etc.).\n",
        "   - **Correct Answer Annotation**: Ensure each question has the correct answer clearly marked.\n",
        "   - Example format:\n",
        "     ```json\n",
        "     {\n",
        "       \"question\": \"What is the capital of France?\",\n",
        "       \"choices\": [\"A. Paris\", \"B. London\", \"C. Rome\", \"D. Madrid\"],\n",
        "       \"correct_answer\": \"A\"\n",
        "     }\n",
        "     ```\n",
        "\n",
        "#### Prompting the Model\n",
        "   - You can prompt the LLM to select the correct answer by presenting the question and the available choices. Structure the input in a way that the LLM understands it is a multiple-choice task.\n",
        "   - Example prompt for GPT-style models:\n",
        "     ```\n",
        "     Question: What is the capital of France?\n",
        "     A) Paris\n",
        "     B) London\n",
        "     C) Rome\n",
        "     D) Madrid\n",
        "     Choose the correct answer:\n",
        "     ```\n",
        "\n",
        "#### Model Inference\n",
        "   - Run the model on each question. Depending on the design of the LLM, you can:\n",
        "     - **Direct Answer Selection**: The model directly outputs one of the choices (e.g., \"A\" for Paris).\n",
        "     - **Text Completion**: The model completes the sentence with the correct answer (e.g., \"The capital of France is Paris.\") and then maps the output back to the correct choice.\n",
        "     - **Logits Comparison**: For models that output probability distributions over tokens, you can compare the logits or probabilities of each choice and select the one with the highest score.\n",
        "\n",
        "#### Evaluation Metrics\n",
        "   \n",
        "- **Accuracy**: The most straightforward metric, simply the ratio of correctly answered questions to the total number of questions.\n",
        "\n",
        "$$Accuracy = \\frac{\\text{Number of Correct Answers}}{\\text{Total Number of Questions}}$$\n",
        "\n",
        "- **Top-N Accuracy**: If evaluating with more complex answer spaces or ambiguous questions, you might use top-N accuracy, where the model is considered correct if the correct answer is within its top N predicted choices.\n",
        "- **Log-Loss**: In cases where the model provides probabilities for each answer, log-loss (cross-entropy) can be used to evaluate how confident the model is about its predictions.\n",
        "- **F1 Score**: If the questions involve multiple correct answers or require ranking, you can compute precision, recall, and the F1 score.\n",
        "\n",
        "So, let's go through these three methods.\n",
        "\n",
        "### Direct Answer Selection (aka Exactly Match)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Sulv26LKyINd"
      },
      "outputs": [],
      "source": [
        "import json\n",
        "from datasets import load_dataset\n",
        "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
        "import torch\n",
        "\n",
        "# Load MMLU dataset (example using \"high_school_geography\" topic from MMLU)\n",
        "mmlu_dataset = load_dataset(\"openai_humaneval\", \"mmlu/high_school_geography\")\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CvAGhkAnyJTz"
      },
      "source": [
        "- For API-based model like ChatGPT, Yiyan, etc."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wcpH3kHPyc_3"
      },
      "outputs": [],
      "source": [
        "# Define your OpenAI API key here\n",
        "openai.api_key = \"YOUR_OPENAI_API_KEY\"\n",
        "\n",
        "def get_model_prediction(prompt, model=\"gpt-3.5-turbo\"):\n",
        "    response = openai.ChatCompletion.create(\n",
        "        model=model,\n",
        "        messages=[\n",
        "            {\"role\": \"system\", \"content\": \"You are a helpful assistant.\"},\n",
        "            {\"role\": \"user\", \"content\": prompt}\n",
        "        ]\n",
        "    )\n",
        "    return response['choices'][0]['message']['content'].strip()\n",
        "\n",
        "def evaluate_exact_match(dataset, model=\"gpt-3.5-turbo\"):\n",
        "    correct_predictions = 0\n",
        "    total_questions = len(dataset)\n",
        "\n",
        "    for idx, example in enumerate(dataset):\n",
        "        # Construct the prompt for the model\n",
        "        prompt = example['question']\n",
        "        for i, choice in enumerate(example['choices']):\n",
        "            prompt += f\"\\n{i + 1}: {choice}\"\n",
        "        prompt += \"\\nAnswer with the number corresponding to the correct choice.\"\n",
        "\n",
        "        # Get model prediction\n",
        "        prediction = get_model_prediction(prompt, model)\n",
        "\n",
        "        # Check if the prediction matches the answer (convert to string)\n",
        "        correct_answer = str(example['answer'])\n",
        "        if prediction == correct_answer:\n",
        "            correct_predictions += 1\n",
        "\n",
        "        print(f\"Question {idx + 1}/{total_questions} | Prediction: {prediction} | Correct Answer: {correct_answer}\")\n",
        "\n",
        "    # Calculate and return exact match accuracy\n",
        "    accuracy = correct_predictions / total_questions\n",
        "    return accuracy\n",
        "\n",
        "# Evaluate the model on MMLU\n",
        "dataset_split = mmlu_dataset['validation']  # or 'test' depending on the split\n",
        "accuracy = evaluate_exact_match(dataset_split)\n",
        "print(f\"Exact Match Accuracy: {accuracy * 100:.2f}%\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Xu8C_g-wydNw"
      },
      "source": [
        "- For models on Hugging Face."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9zrAXmOFYZ94"
      },
      "outputs": [],
      "source": [
        "# Define your Hugging Face model and tokenizer\n",
        "model_name = \"gpt2\"  # Replace with the desired HF model\n",
        "tokenizer = AutoTokenizer.from_pretrained(model_name)\n",
        "model = AutoModelForCausalLM.from_pretrained(model_name)\n",
        "\n",
        "def get_model_prediction(prompt, model, tokenizer):\n",
        "    inputs = tokenizer(prompt, return_tensors=\"pt\")\n",
        "    with torch.no_grad():\n",
        "        output = model.generate(**inputs, max_length=inputs.input_ids.shape[1] + 10)\n",
        "    prediction = tokenizer.decode(output[0], skip_special_tokens=True)\n",
        "    return prediction.strip()\n",
        "\n",
        "def evaluate_exact_match(dataset, model, tokenizer):\n",
        "    correct_predictions = 0\n",
        "    total_questions = len(dataset)\n",
        "\n",
        "    for idx, example in enumerate(dataset):\n",
        "        # Construct the prompt for the model\n",
        "        prompt = example['question']\n",
        "        for i, choice in enumerate(example['choices']):\n",
        "            prompt += f\"\\n{i + 1}: {choice}\"\n",
        "        prompt += \"\\nAnswer with the number corresponding to the correct choice.\"\n",
        "\n",
        "        # Get model prediction\n",
        "        prediction = get_model_prediction(prompt, model, tokenizer)\n",
        "\n",
        "        # Extract the predicted answer (assuming it ends with the choice number)\n",
        "        predicted_answer = prediction.split()[-1]\n",
        "\n",
        "        # Check if the prediction matches the answer (convert to string)\n",
        "        correct_answer = str(example['answer'])\n",
        "        if predicted_answer == correct_answer:\n",
        "            correct_predictions += 1\n",
        "\n",
        "        print(f\"Question {idx + 1}/{total_questions} | Prediction: {predicted_answer} | Correct Answer: {correct_answer}\")\n",
        "\n",
        "    # Calculate and return exact match accuracy\n",
        "    accuracy = correct_predictions / total_questions\n",
        "    return accuracy\n",
        "\n",
        "# Evaluate the model on MMLU\n",
        "dataset_split = mmlu_dataset['validation']  # or 'test' depending on the split\n",
        "accuracy = evaluate_exact_match(dataset_split, model, tokenizer)\n",
        "print(f\"Exact Match Accuracy: {accuracy * 100:.2f}%\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "C2l7UGhpEM8b"
      },
      "source": [
        "### Text Completion"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qbRxIPH5EUJE"
      },
      "outputs": [],
      "source": [
        "import json\n",
        "import numpy as np\n",
        "from datasets import load_dataset\n",
        "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
        "import torch\n",
        "\n",
        "# Load MMLU dataset (example using \"high_school_geography\" topic from MMLU)\n",
        "mmlu_dataset = load_dataset(\"openai_humaneval\", \"mmlu/high_school_geography\")\n",
        "\n",
        "# Define your Hugging Face model and tokenizer\n",
        "model_name = \"gpt2\"  # Replace with the desired HF model\n",
        "tokenizer = AutoTokenizer.from_pretrained(model_name)\n",
        "model = AutoModelForCausalLM.from_pretrained(model_name)\n",
        "\n",
        "def get_perplexity(prompt, model, tokenizer):\n",
        "    inputs = tokenizer(prompt, return_tensors=\"pt\")\n",
        "    with torch.no_grad():\n",
        "        outputs = model(**inputs, labels=inputs.input_ids)\n",
        "    log_likelihood = outputs.loss.item() * inputs.input_ids.size(1)\n",
        "    perplexity = np.exp(log_likelihood / inputs.input_ids.size(1))\n",
        "    return perplexity\n",
        "\n",
        "def evaluate_perplexity(dataset, model, tokenizer):\n",
        "    total_perplexity = 0\n",
        "    total_questions = len(dataset)\n",
        "\n",
        "    for idx, example in enumerate(dataset):\n",
        "        question = example['question']\n",
        "        min_perplexity = float('inf')\n",
        "        best_choice = None\n",
        "\n",
        "        # Calculate perplexity for each choice\n",
        "        for choice in example['choices']:\n",
        "            prompt = f\"{question} {choice}\"\n",
        "            perplexity = get_perplexity(prompt, model, tokenizer)\n",
        "\n",
        "            if perplexity < min_perplexity:\n",
        "                min_perplexity = perplexity\n",
        "                best_choice = choice\n",
        "\n",
        "        # Check if the choice with lowest perplexity matches the answer\n",
        "        correct_answer = example['choices'][int(example['answer']) - 1]\n",
        "        if best_choice == correct_answer:\n",
        "            total_perplexity += 1\n",
        "\n",
        "        print(f\"Question {idx + 1}/{total_questions} | Best Choice: {best_choice} | Correct Answer: {correct_answer} | Min Perplexity: {min_perplexity}\")\n",
        "\n",
        "    # Calculate and return perplexity accuracy\n",
        "    accuracy = total_perplexity / total_questions\n",
        "    return accuracy\n",
        "\n",
        "# Evaluate the model on MMLU for perplexity\n",
        "dataset_split = mmlu_dataset['validation']  # or 'test' depending on the split\n",
        "accuracy = evaluate_perplexity(dataset_split, model, tokenizer)\n",
        "print(f\"Perplexity-Based Accuracy: {accuracy * 100:.2f}%\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3TY1lu9lEMtK"
      },
      "source": [
        "### Logits Comparison"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kUf7CPmeEU5J"
      },
      "outputs": [],
      "source": [
        "import json\n",
        "import numpy as np\n",
        "from datasets import load_dataset\n",
        "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
        "import torch\n",
        "import torch.nn.functional as F\n",
        "\n",
        "# Load MMLU dataset (example using \"high_school_geography\" topic from MMLU)\n",
        "mmlu_dataset = load_dataset(\"openai_humaneval\", \"mmlu/high_school_geography\")\n",
        "\n",
        "# Define your Hugging Face model and tokenizer\n",
        "model_name = \"gpt2\"  # Replace with the desired HF model\n",
        "tokenizer = AutoTokenizer.from_pretrained(model_name)\n",
        "model = AutoModelForCausalLM.from_pretrained(model_name)\n",
        "\n",
        "def get_logits_probability(prompt, model, tokenizer, options):\n",
        "    inputs = tokenizer(prompt, return_tensors=\"pt\")\n",
        "    with torch.no_grad():\n",
        "        outputs = model(**inputs)\n",
        "    logits = outputs.logits[0, -1, :]\n",
        "\n",
        "    probabilities = {}\n",
        "    for option in options:\n",
        "        option_token_id = tokenizer(option, add_special_tokens=False).input_ids[0]\n",
        "        probabilities[option] = F.softmax(logits, dim=-1)[option_token_id].item()\n",
        "    return probabilities\n",
        "\n",
        "def evaluate_logits(dataset, model, tokenizer):\n",
        "    correct_predictions = 0\n",
        "    total_questions = len(dataset)\n",
        "\n",
        "    for idx, example in enumerate(dataset):\n",
        "        question = example['question']\n",
        "        options = [\"A\", \"B\", \"C\", \"D\", \"a\", \"b\", \"c\", \"d\"]\n",
        "\n",
        "        # Calculate probabilities for each option\n",
        "        probabilities = get_logits_probability(question, model, tokenizer, options)\n",
        "\n",
        "        # Find the option with the highest probability\n",
        "        best_option = max(probabilities, key=probabilities.get)\n",
        "\n",
        "        # Check if the predicted option matches the answer\n",
        "        correct_answer = example['choices'][int(example['answer']) - 1]\n",
        "        correct_label = correct_answer[0]  # Assume the correct answer starts with A, B, C, or D\n",
        "\n",
        "        if best_option.lower() == correct_label.lower():\n",
        "            correct_predictions += 1\n",
        "\n",
        "        print(f\"Question {idx + 1}/{total_questions} | Predicted: {best_option} | Correct Answer: {correct_label} | Probability: {probabilities[best_option]:.4f}\")\n",
        "\n",
        "    # Calculate and return accuracy\n",
        "    accuracy = correct_predictions / total_questions\n",
        "    return accuracy\n",
        "\n",
        "# Evaluate the model on MMLU using logits comparison\n",
        "dataset_split = mmlu_dataset['validation']  # or 'test' depending on the split\n",
        "accuracy = evaluate_logits(dataset_split, model, tokenizer)\n",
        "print(f\"Logits-Based Accuracy: {accuracy * 100:.2f}%\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AcMLNn9nYqDY"
      },
      "source": [
        "## 2. Model Service Deployment (GPU)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 339
        },
        "collapsed": true,
        "id": "YEUsLFyNOQlg",
        "outputId": "decb7d96-8b86-4ae3-9459-7ce16d700b90"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "from vllm import LLM, SamplingParams"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 211
        },
        "collapsed": true,
        "id": "SLd8KYjgOOpf",
        "outputId": "3bd787e6-5589-4e1f-9c6d-668d08d92775"
      },
      "outputs": [],
      "source": [
        "prompts = [\n",
        "    \"Hello, my name is\",\n",
        "    \"The president of the United States is\",\n",
        "    \"The capital of France is\",\n",
        "    \"The future of AI is\",\n",
        "]\n",
        "\n",
        "# initialize\n",
        "sampling_params = SamplingParams(temperature=0.8, top_p=0.95)\n",
        "llm = LLM(model=checkpoint, dtype=\"float32\")\n",
        "\n",
        "# perform the inference\n",
        "outputs = llm.generate(prompts, sampling_params)\n",
        "\n",
        "# print outputs\n",
        "for output in outputs:\n",
        "    prompt = output.prompt\n",
        "    generated_text = output.outputs[0].text\n",
        "    print(f\"Prompt: {prompt!r}, Generated text: {generated_text!r}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "H_1RaOQGYtfd"
      },
      "source": [
        "## 3. Model On-device Deployment (CPU)\n",
        "\n",
        "In general, practical frameworks for CPU Deployment:\n",
        "ONNX Runtime with OpenVINO - Converts and optimizes models for efficient inference.\n",
        "DeepSparse by Neural Magic - Takes advantage of sparsity for ultra-fast inference on CPUs.\n",
        "TensorFlow Lite - Can be used for efficient CPU-based inference with optimizations.\n",
        "MLC & PowerInfer2 - Since you've previously used these frameworks, they could be adapted to a CPU-centric setting, leveraging the lightweight inference they support.\n",
        "\n",
        "- First, To use MLC in Python for model deployment and inference, which involves tools like TVM, MLC LLM, and Here's a step-by-step guide on how you can use MLC within Python. And let's install some important package."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "collapsed": true,
        "id": "W0ouQGQBYv_G",
        "outputId": "abbd7d60-7a64-4d35-cc8d-8d3a543ae019"
      },
      "outputs": [],
      "source": [
        "!python -m pip install --pre -U -f https://mlc.ai/wheels mlc-llm-nightly-cpu mlc-ai-nightly-cpu\n",
        "!pip install onnx onnxruntime"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rGqRKvQXPFqj"
      },
      "source": [
        "And import these package with other will-be-use library."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 31,
      "metadata": {
        "id": "_S0dntUS6KjO"
      },
      "outputs": [],
      "source": [
        "import tvm\n",
        "from tvm import relay\n",
        "from tvm.contrib import graph_executor\n",
        "import numpy as np\n",
        "import torch\n",
        "import torch.onnx\n",
        "from transformers import AutoTokenizer, AutoModelForCausalLM\n",
        "from transformers import GPT2Tokenizer\n",
        "import onnxruntime as ort\n",
        "import numpy as np"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4BCEbttyPM9_"
      },
      "source": [
        "We take GPT2 (127M) as an example, to see how it can work using CPU. At the very beginning, we will be exporting GPT-2 to ONNX."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "collapsed": true,
        "id": "rQHb9QlA7Djz",
        "outputId": "6fbe1a52-7572-48bb-8d02-20227c4cda4f"
      },
      "outputs": [],
      "source": [
        "# Define your Hugging Face model and tokenizer\n",
        "model_name = \"gpt2\"  # Replace with the desired HF model\n",
        "tokenizer = AutoTokenizer.from_pretrained(model_name)\n",
        "model = AutoModelForCausalLM.from_pretrained(model_name)\n",
        "\n",
        "# Assume you have a trained model saved in `model.pth`\n",
        "# model = torch.load(\"path/to/your_model.pth\")\n",
        "model.eval()  # Switch to evaluation mode for exporting\n",
        "\n",
        "# Prepare Dummy Input for Export:\n",
        "# You need a dummy input that matches the expected input type of GPT-2:\n",
        "text = \"Hello, how are you?\"\n",
        "inputs = tokenizer(text, return_tensors=\"pt\")\n",
        "\n",
        "input_ids = inputs[\"input_ids\"]  # This is the tensor containing tokenized input"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Nz0ug9bePngq"
      },
      "source": [
        "Export the GPT-2 model to ONNX format using torch.onnx.export:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {
        "collapsed": true,
        "id": "OGLGQ8_d71yf"
      },
      "outputs": [],
      "source": [
        "torch.onnx.export(\n",
        "    model,\n",
        "    input_ids,\n",
        "    \"gpt2.onnx\",\n",
        "    input_names=[\"input_ids\"],\n",
        "    output_names=[\"last_hidden_state\"],\n",
        "    dynamic_axes={\"input_ids\": {0: \"batch_size\", 1: \"sequence_length\"},\n",
        "                  \"last_hidden_state\": {0: \"batch_size\", 1: \"sequence_length\"}},  # Dynamic batch and sequence length\n",
        "    opset_version=14\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_TMv9nQvPrES"
      },
      "source": [
        "Explanation:\n",
        "\n",
        "`input_ids`: This is the input tensor containing tokenized input. GPT-2 expects the input tensor to be of type LongTensor (int64).\n",
        "**Dynamic Axes**: GPT-2 can work with variable batch sizes and sequence lengths. Defining \"batch_size\" and \"sequence_length\" as dynamic helps to run inference on inputs of different lengths.\n",
        "\n",
        "Then we will convert ONNX Model to TVM Relay.\n",
        "\n",
        "Once you have the GPT-2 model in ONNX format, you can convert it to TVM's Relay IR and then compile it. For GPT-2, the inputs are different compared to image models, and special attention should be given to the input tensor types and shapes."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "collapsed": true,
        "id": "ICI2syVT8HWJ",
        "outputId": "e1018b9b-689b-47d6-f819-25868a099cd3"
      },
      "outputs": [],
      "source": [
        "import onnx\n",
        "\n",
        "# Load the ONNX model\n",
        "onnx_model = onnx.load(\"gpt2.onnx\")\n",
        "\n",
        "# Print input names\n",
        "for input in onnx_model.graph.input:\n",
        "    print(f\"Input name: {input.name}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "B2wmQfsFP4eE"
      },
      "source": [
        "This will help you determine the input names, which you need when creating shape_dict. GPT-2 expects a tokenized input of variable length. We'll need to define the shape based on your typical input length:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {
        "id": "R-PlN5op8YpZ"
      },
      "outputs": [],
      "source": [
        "shape_dict = {\"input_ids\": (1, 10)}  # Batch size of 1, sequence length of 10 tokens (can be adjusted)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e7m2BSaeP9l8"
      },
      "source": [
        "Then, we use the following script to convert the ONNX model to TVM:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "IxzpRxWI85HS",
        "outputId": "f57a9693-2959-4a33-8b10-829b49de0761"
      },
      "outputs": [],
      "source": [
        "import tvm\n",
        "from tvm import relay\n",
        "\n",
        "# Convert the ONNX model to Relay format\n",
        "mod, params = relay.frontend.from_onnx(onnx_model, shape_dict)\n",
        "\n",
        "# Set target for CPU compilation\n",
        "target = \"llvm\"\n",
        "dev = tvm.cpu()\n",
        "\n",
        "# Compile the model with TVM\n",
        "with tvm.transform.PassContext(opt_level=3):\n",
        "    lib = relay.build(mod, target=target, params=params)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "i571lIsm_J29"
      },
      "source": [
        "The `target = \"llvm\"` setting in TVM refers to **LLVM (Low-Level Virtual Machine)** as the compilation target, which is used to generate optimized code for **CPU** inference. Here’s a detailed explanation:\n",
        "\n",
        "### Why Use `llvm` as the Target?\n",
        "1. **General CPU Backend**:\n",
        "   - In TVM, `\"llvm\"` is used to represent a **general-purpose CPU backend**. By setting the target to `\"llvm\"`, TVM will generate machine code optimized for your CPU. LLVM is a popular compiler infrastructure that allows TVM to generate highly optimized native code for a wide variety of CPU architectures, including x86, ARM, etc.\n",
        "   - This is why `\"llvm\"` is commonly used when you want to run inference on a **CPU**, regardless of the specific type of CPU in your machine.\n",
        "\n",
        "2. **Cross-Platform Compilation**:\n",
        "   - LLVM can be used to generate code for different types of CPUs, which makes it flexible and well-suited for **cross-compilation**. This means you can compile a model on one type of CPU and run it on another type.\n",
        "   - For example, you could use LLVM to compile on an x86 machine and then execute the code on an ARM-based CPU.\n",
        "\n",
        "3. **Optimized Performance**:\n",
        "   - LLVM enables various optimizations that are essential for deep learning workloads, such as **loop unrolling**, **vectorization**, and **instruction selection**. By using LLVM as the target, TVM can generate highly optimized machine-level code that takes advantage of CPU-specific features, such as vector instructions (e.g., AVX, NEON).\n",
        "   - This optimization process helps improve the performance of the inference on CPU.\n",
        "\n",
        "4. **Ease of Use**:\n",
        "   - The `\"llvm\"` target is the most straightforward target for many developers because it’s CPU-agnostic and doesn’t require specifying complex hardware details.\n",
        "   - It allows developers to use TVM without needing to worry about hardware-specific intricacies.\n",
        "\n",
        "### Examples of TVM Targets\n",
        "Besides `\"llvm\"`, TVM supports various targets, each suited to different hardware:\n",
        "\n",
        "- **\"cuda\"**: Used for compiling models to run on **NVIDIA GPUs**. This enables the use of CUDA to leverage GPU acceleration.\n",
        "- **\"opencl\"**: Used for **OpenCL-compatible GPUs**, FPGAs, or other hardware accelerators.\n",
        "- **\"rocm\"**: Target for **AMD GPUs** that support ROCm.\n",
        "- **\"metal\"**: Target for running models on **macOS and iOS** devices using Apple's Metal API for GPU support.\n",
        "- **\"llvm -mtriple=aarch64-linux-gnu\"**: Can be used for **cross-compiling** to an ARM architecture CPU, typically for ARM-based devices such as embedded systems or mobile devices.\n",
        "\n",
        "### Example Use Cases\n",
        "- If you are deploying a model on a **desktop server or a cloud server** that uses CPUs, using `\"llvm\"` is appropriate.\n",
        "- For deploying a model on an **NVIDIA GPU**, you would use `\"cuda\"` as the target to ensure the generated code takes full advantage of GPU capabilities.\n",
        "\n",
        "***Back to the Code***\n",
        "\n",
        ">Once compiled, we run inference using TVM's Graph Executor. You need to prepare an input tensor that matches the expected input format.\n",
        "\n",
        "Let's see the next logits first."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Llb2XFvC986A",
        "outputId": "7a8c91cc-f9e0-4d16-f90f-c6f59b398bde"
      },
      "outputs": [],
      "source": [
        "from tvm.contrib import graph_executor\n",
        "import numpy as np\n",
        "\n",
        "# Create runtime executor module\n",
        "module = graph_executor.GraphModule(lib[\"default\"](dev))\n",
        "\n",
        "# Create input tensor (make sure it matches the expected type and shape)\n",
        "dummy_input = np.random.randint(0, 50256, (1, 10)).astype(\"int64\")  # GPT-2 has a vocabulary size of 50256\n",
        "\n",
        "# Set the input and run the model\n",
        "module.set_input(\"input_ids\", dummy_input)\n",
        "module.run()\n",
        "\n",
        "# Get the output\n",
        "output = module.get_output(0).numpy()\n",
        "print(\"Model output:\", output)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qWFw8R8hQc8u"
      },
      "source": [
        "Then, we can see the next token"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "seFWp_coGeev",
        "outputId": "d7dea63d-8800-4781-a82d-b3121f916d96"
      },
      "outputs": [],
      "source": [
        "# Load the pretrained GPT-2 tokenizer\n",
        "tokenizer = AutoTokenizer.from_pretrained(\"gpt2\")\n",
        "\n",
        "# Create a natural language input\n",
        "input_text = \"Hello, what is your name\"\n",
        "\n",
        "# Tokenize the input text\n",
        "inputs = tokenizer(input_text, return_tensors=\"np\")\n",
        "input_ids = inputs[\"input_ids\"]  # This will be a NumPy array of token IDs\n",
        "\n",
        "# Load the ONNX model\n",
        "ort_session = ort.InferenceSession(\"gpt2.onnx\")\n",
        "\n",
        "# Use the tokenized input IDs for inference\n",
        "outputs = ort_session.run(None, {\"input_ids\": input_ids})\n",
        "\n",
        "# Get the logits output from the ONNX model\n",
        "logits = outputs[0]\n",
        "\n",
        "next_token_logits = logits[:, -1, :]  # Get the logits for the last token in the sequence\n",
        "next_token_id = np.argmax(next_token_logits, axis=-1)  # Choose the token with the highest probability\n",
        "next_token = tokenizer.decode(next_token_id)\n",
        "print(\"Generated word:\", next_token)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "m3m0w1lgQuz_"
      },
      "source": [
        "And, finally, we generate the whole sentence."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "7JimS6-tGqAh",
        "outputId": "6f45c765-664d-49f1-ef5f-50211813cea1"
      },
      "outputs": [],
      "source": [
        "generated_ids = list(input_ids[0]) # Start with the original input\n",
        "\n",
        "for _ in range(20):  # Generate up to 20 tokens\n",
        "    # Convert to numpy and run inference\n",
        "    input_array = np.array([generated_ids], dtype=np.int64)\n",
        "    logits = ort_session.run(None, {\"input_ids\": input_array})[0]\n",
        "\n",
        "    # Get logits for the last token\n",
        "    next_token_logits = logits[:, -1, :]\n",
        "    next_token_id = np.argmax(next_token_logits, axis=-1)\n",
        "\n",
        "    # Append next token to generated sequence\n",
        "    generated_ids.append(next_token_id[0])\n",
        "\n",
        "    # Stop if end-of-sequence token is generated (for GPT-2, token 50256)\n",
        "    if next_token_id[0] == 50256:\n",
        "        break\n",
        "\n",
        "# Convert generated token IDs back to text\n",
        "generated_text = tokenizer.decode(generated_ids)\n",
        "print(\"Generated text:\", generated_text)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tzHq5-00YwxH"
      },
      "source": [
        "## Reference"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}