20231220_SwallowLLM.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "machine_shape": "hm",
      "gpuType": "V100",
      "authorship_tag": "ABX9TyOZ44F0djSP3T54pf9h6Ody",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/aicuai/GenAI-Steam/blob/main/20231220_SwallowLLM.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Demo for Swallow LLM by Tokyo Tech\n",
        "\n",
        "公式サイト\n",
        "https://tokyotech-llm.github.io/\n",
        "\n",
        "モデル\n",
        "https://huggingface.co/tokyotech-llm\n",
        "\n",
        "★AICU media 「[東工大と産総研、英語の言語理解や対話で高い能力を持つ大規模言語モデル「Swallow」を公開 #SwallowLLM](https://note.com/aicu/n/n3eb8c1f2df02)」\n",
        "\n",
        "★AICU media 「[東工大LLM「Swallow」を使ってGoogle Colabで遊んでみよう #SwallowLLM](https://note.com/aicu/n/nd0337d4952f3)」の解説コードです \n",
        "\n",
        "参考：比較的初心者向けのGoogle Colabでの「Japanese Stable LM Gamma 7B」を動かす記事\n",
        "\n",
        "■[Stability AI Japanが公開した30億パラメーターの日本語向けLLMを動かしてみた - 生成AIストリーム - 窓の杜](https://forest.watch.impress.co.jp/docs/serial/aistream/1544320.html)\n",
        "\n",
        "Coded by Akihiko SHIRAI (kaitas[@o_ob](https://twitter.com/o_ob)) PoC開発や技術発信のお仕事歓迎です"
      ],
      "metadata": {
        "id": "UQlni64Jqps2"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Step.1\n",
        "\n",
        "[Google Colab Pro](https://colab.research.google.com/signup/pricing?hl=ja) を使ってGPUが利用できるインスタンスを作ります。具体的には「A100 GPU」もしくは「V100 GPU」以上を選びましょう。「A100」を選んでしばらく待てば、運が良ければ割り当てられます（利用できない時は「V100」になります）。\n",
        "\n"
      ],
      "metadata": {
        "id": "pstgLy4Oq9-I"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jDa0yN1gqgla"
      },
      "outputs": [],
      "source": [
        "# パッケージのインストール、いろいろあるけど Colab環境ならこれだけで動くはず\n",
        "!pip install sentencepiece accelerate"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# @title Step.2 Tokenizer & Model Loading\n",
        "import torch\n",
        "from transformers import AutoTokenizer, AutoModelForCausalLM\n",
        "\n",
        "# @markdown [https://huggingface.co/tokyotech-llm](https://huggingface.co/tokyotech-llm) から利用したいモデルを選択してください。最初は 7b-instruct から始めるのがおすすめです。13bは ColabPro では動いています。70bはColabProでもダウンロードが難しいです。\n",
        "tokenizer_model_name = \"tokyotech-llm/Swallow-13b-instruct-hf\" # @param ['tokyotech-llm/Swallow-7b-hf','tokyotech-llm/Swallow-7b-instruct-hf','tokyotech-llm/Swallow-13b-hf','tokyotech-llm/Swallow-13b-instruct-hf','tokyotech-llm/Swallow-70b-hf','tokyotech-llm/Swallow-70b-instruct-hf']\n",
        "model_name = \"tokyotech-llm/Swallow-13b-instruct-hf\" # @param ['tokyotech-llm/Swallow-7b-hf','tokyotech-llm/Swallow-7b-instruct-hf','tokyotech-llm/Swallow-13b-hf','tokyotech-llm/Swallow-13b-instruct-hf','tokyotech-llm/Swallow-70b-hf','tokyotech-llm/Swallow-70b-instruct-hf']\n",
        "\n",
        "# Load\n",
        "tokenizer = AutoTokenizer.from_pretrained(tokenizer_model_name)\n",
        "model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16, low_cpu_mem_usage=True ,device_map=\"auto\")\n",
        "\n",
        "# @title\n",
        "PROMPT_DICT = {\n",
        "    \"prompt_input\": (\n",
        "        \"以下に、あるタスクを説明する指示があり、それに付随する入力が更なる文脈を提供しています。\"\n",
        "        \"リクエストを適切に完了するための回答を記述してください。\\n\\n\"\n",
        "        \"### 指示:\\n{instruction}\\n\\n### 入力:\\n{input}\\n\\n### 応答:\"\n",
        "\n",
        "    ),\n",
        "    \"prompt_no_input\": (\n",
        "        \"以下に、あるタスクを説明する指示があります。\"\n",
        "        \"リクエストを適切に完了するための回答を記述してください。\\n\\n\"\n",
        "        \"### 指示:\\n{instruction}\\n\\n### 応答:\"\n",
        "    ),\n",
        "}\n",
        "\n",
        "def create_prompt(instruction, input=None):\n",
        "    \"\"\"\n",
        "    Generates a prompt based on the given instruction and an optional input.\n",
        "    If input is provided, it uses the 'prompt_input' template from PROMPT_DICT.\n",
        "    If no input is provided, it uses the 'prompt_no_input' template.\n",
        "\n",
        "    Args:\n",
        "        instruction (str): The instruction describing the task.\n",
        "        input (str, optional): Additional input providing context for the task. Default is None.\n",
        "\n",
        "    Returns:\n",
        "        str: The generated prompt.\n",
        "    \"\"\"\n",
        "    if input:\n",
        "        # Use the 'prompt_input' template when additional input is provided\n",
        "        return PROMPT_DICT[\"prompt_input\"].format(instruction=instruction, input=input)\n",
        "    else:\n",
        "        # Use the 'prompt_no_input' template when no additional input is provided\n",
        "        return PROMPT_DICT[\"prompt_no_input\"].format(instruction=instruction)\n"
      ],
      "metadata": {
        "cellView": "form",
        "id": "Q_xj3OlBsiKy"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# @title Step.3 Settings & Prompts\n",
        "instruction_example = \"以下のトピックに関する詳細な情報を提供してください。\" # @param {type: \"string\"}\n",
        "input_example = \"\\u6771\\u4EAC\\u5DE5\\u696D\\u5927\\u5B66\\u306E\\u6A2A\\u6D5C\\u306B\\u3042\\u308B\\u30AD\\u30E3\\u30F3\\u30D1\\u30B9\\u306E\\u98DF\\u5802\\u306B\\u3064\\u3044\\u3066\\u8AAC\\u660E\\u3057\\u3066\\u304F\\u3060\\u3055\\u3044\" # @param {type: \"string\"}\n",
        "Do_sample=True #@param {type:\"boolean\"}\n",
        "\n",
        "if Do_sample:\n",
        "  temperature = 0.99 #@param {type:\"slider\", min:0, max:2, step:0.1}\n",
        "  top_p = 0.95 #@param {type:\"slider\", min:0, max:1, step:0.01}\n",
        "\n",
        "max_new_tokens=128 #@param {type:\"slider\", min:128, max:1024, step:64}\n",
        "\n",
        "# Example usage\n",
        "# instruction_example = \"以下のトピックに関する詳細な情報を提供してください。\"\n",
        "# input_example = \"東京工業大学の主なキャンパスについて教えてください\"\n",
        "prompt = create_prompt(instruction_example, input_example)\n",
        "\n",
        "input_ids = tokenizer.encode(\n",
        "    prompt,\n",
        "    add_special_tokens=False,\n",
        "    return_tensors=\"pt\"\n",
        ")\n",
        "\n",
        "tokens = model.generate(\n",
        "    input_ids.to(device=model.device),\n",
        "    max_new_tokens=max_new_tokens,\n",
        "    temperature=temperature,\n",
        "    top_p=top_p,\n",
        "    do_sample=Do_sample,\n",
        ")\n",
        "\n",
        "out = tokenizer.decode(tokens[0], skip_special_tokens=True)\n",
        "print(out)\n"
      ],
      "metadata": {
        "cellView": "form",
        "id": "BvRt-mTPsQow"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "QTeTSsSzI3UM"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}