cookbook/System_prompt/System_prompt.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "gpuType": "L4"
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "# System Instruction Tutorial - Quick Start: Common Scenarios for Using the Yi-Coder Model\n",
        "\n",
        "**Yi-Coder**, the most powerful language model under 10 billion parameters, boasts an impressive 128k long text capability.  Yi-Coder is a game-changer for developers, pushing the boundaries of code understanding and generation. With its unparalleled accuracy and efficiency, Yi-Coder empowers you to write cleaner, more robust code, faster than ever before.\n",
        "\n",
        "This tutorial will guide you through a fast-paced exploration of Yi-Coder's capabilities, focusing on practical applications like:\n",
        "\n",
        "* **Code Completion:** Experience seamless code suggestions as you type, saving valuable time and effort.\n",
        "* **Code Insertion:** Effortlessly insert complex code snippets, enhancing your coding workflow.\n",
        "* **Repo Q&A:** Ask Yi-Coder questions about your codebase, getting precise answers and insightful feedback.\n",
        "\n",
        "Get ready to unlock the full potential of Yi-Coder and revolutionize your coding experience. Let's dive in!\n"
      ],
      "metadata": {
        "id": "Bch3RhmuXsy-"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## 1. Code Completion\n",
        "\n",
        "Here, we will use the example of \"writing a quick sort algorithm\" to illustrate system prompts.\n",
        "\n",
        "**Recommended System Prompt:**\n",
        "\n",
        "```\n",
        "You are Yi-Coder, you are exceptionally skilled in programming, coding, and any computer-related issues.\n",
        "```"
      ],
      "metadata": {
        "id": "eU4LYptRc_UD"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from transformers import AutoTokenizer, AutoModelForCausalLM\n",
        "\n",
        "device = \"cuda\" # the device to load the model onto\n",
        "model_path = '01-ai/Yi-Coder-9B-Chat'\n",
        "# Now you do not need to add \"trust_remote_code=True\"\n",
        "tokenizer = AutoTokenizer.from_pretrained(model_path)\n",
        "model = AutoModelForCausalLM.from_pretrained(model_path, device_map=\"auto\").eval()\n",
        "\n",
        "# tokenize the input into tokens\n",
        "# Instead of using model.chat(), we directly use model.generate()\n",
        "# But you need to use tokenizer.apply_chat_template() to format your inputs as shown below\n",
        "prompt = \"Write a quick sort algorithm.\"\n",
        "messages = [\n",
        "    {\"role\": \"system\", \"content\": \"You are Yi-Coder, you are exceptionally skilled in programming, coding, and any computer-related issues.\"},\n",
        "    {\"role\": \"user\", \"content\": prompt}\n",
        "]\n",
        "text = tokenizer.apply_chat_template(\n",
        "    messages,\n",
        "    tokenize=False,\n",
        "    add_generation_prompt=True\n",
        ")\n",
        "model_inputs = tokenizer([text], return_tensors=\"pt\").to(device)\n",
        "\n",
        "# Directly use generate() and tokenizer.decode() to get the output.\n",
        "# Use `max_new_tokens` to control the maximum output length.\n",
        "generated_ids = model.generate(\n",
        "    model_inputs.input_ids,\n",
        "    max_new_tokens=1024,\n",
        "    eos_token_id=tokenizer.eos_token_id\n",
        ")\n",
        "generated_ids = [\n",
        "    output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)\n",
        "]\n",
        "\n",
        "response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]\n",
        "print(response)"
      ],
      "metadata": {
        "id": "EHrKmCUZKEzk"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## 2. Code Insertion\n",
        "\n",
        "This scenario demonstrates how Yi-Coder can identify errors and insert the correct code to fix them.\n",
        "\n",
        "**Recommended System Prompt:**\n",
        "\n",
        "```\n",
        "You are Yi-Coder, you are exceptionally skilled in programming, coding, and any computer-related issues.\n",
        "```\n"
      ],
      "metadata": {
        "id": "oNU-ttDhdg3n"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Instead of using model.chat(), we directly use model.generate()\n",
        "# But you need to use tokenizer.apply_chat_template() to format your inputs as shown below\n",
        "prompt = \"\"\"\n",
        "```python\n",
        "def quick_sort(arr):\n",
        "    if len(arr) <= 1:\n",
        "        return arr\n",
        "    else:\n",
        "        pivot = arr[len(arr) // 2]\n",
        "        left = [x for x in arr if x < pivot]\n",
        "\n",
        "        right = [x for x in arr if x > pivot]\n",
        "        return quick_sort(left) + middle + quick_sort(right)\n",
        "\n",
        "print(quick_sort([3,6,8,10,1,2,1]))\n",
        "# Prints \"[1, 1, 2, 3, 6, 8, 10]\"\n",
        "```\n",
        "Is there a problem with this code?\n",
        "\"\"\"\n",
        "messages = [\n",
        "    {\"role\": \"system\", \"content\": \"You are Yi-Coder, you are exceptionally skilled in programming, coding, and any computer-related issues.\"},\n",
        "    {\"role\": \"user\", \"content\": prompt}\n",
        "]\n",
        "text = tokenizer.apply_chat_template(\n",
        "    messages,\n",
        "    tokenize=False,\n",
        "    add_generation_prompt=True\n",
        ")\n",
        "model_inputs = tokenizer([text], return_tensors=\"pt\").to(device)\n",
        "\n",
        "# Directly use generate() and tokenizer.decode() to get the output.\n",
        "# Use `max_new_tokens` to control the maximum output length.\n",
        "generated_ids = model.generate(\n",
        "    model_inputs.input_ids,\n",
        "    max_new_tokens=1024,\n",
        "    eos_token_id=tokenizer.eos_token_id\n",
        ")\n",
        "generated_ids = [\n",
        "    output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)\n",
        "]\n",
        "\n",
        "response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]\n",
        "print(response)"
      ],
      "metadata": {
        "id": "KUWhFBW9KeiM"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## 3. Repo Q&A\n",
        "\n",
        "This scenario focuses on leveraging Yi-Coder to answer questions about an entire codebase.\n",
        "\n",
        "**Recommended System Prompt:**\n",
        "\n",
        "```\n",
        "You are Yi-Coder, you are exceptionally skilled in programming, coding, and any computer-related issues.\n",
        "```"
      ],
      "metadata": {
        "id": "12qI1E_Gd6aJ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Instead of using model.chat(), we directly use model.generate()\n",
        "# But you need to use tokenizer.apply_chat_template() to format your inputs as shown below\n",
        "prompt = \"\"\"\n",
        "  <<File1 Path>>\n",
        "  {Code Content}\n",
        "  <<File2 Path>>\n",
        "  {Code Content}\n",
        "  <<File3 Path>>\n",
        "  {Code Content}\n",
        "  <<File4 Path>>\n",
        "  {Code Content}\n",
        "  <<File5 Path>>\n",
        "  {Code Content}\n",
        "  <<Query>>\"\"\"\n",
        "messages = [\n",
        "    {\"role\": \"system\", \"content\": \"You are a helpful assistant.\"},\n",
        "    {\"role\": \"user\", \"content\": prompt}\n",
        "]\n",
        "text = tokenizer.apply_chat_template(\n",
        "    messages,\n",
        "    tokenize=False,\n",
        "    add_generation_prompt=True\n",
        ")\n",
        "model_inputs = tokenizer([text], return_tensors=\"pt\").to(device)\n",
        "\n",
        "# Directly use generate() and tokenizer.decode() to get the output.\n",
        "# Use `max_new_tokens` to control the maximum output length.\n",
        "generated_ids = model.generate(\n",
        "    model_inputs.input_ids,\n",
        "    max_new_tokens=1024,\n",
        "    eos_token_id=tokenizer.eos_token_id\n",
        ")\n",
        "generated_ids = [\n",
        "    output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)\n",
        "]\n",
        "\n",
        "response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]\n",
        "print(response)"
      ],
      "metadata": {
        "id": "qvnW_5OfKmZY"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "KnJ5_csqXGts"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}