Japanese_Stable_LM_2_Base_1_6B.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "gpuType": "T4",
      "authorship_tag": "ABX9TyNJr9vXGPrVvJTE5RiRLuXX",
      "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/Japanese_Stable_LM_2_Base_1_6B.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Japanese Stable LM 2 1.6B on Colab by AICU  https://j.aicu.ai/JSLM2B16B\n",
        "\n",
        "これは Stability AI が2024/5/9にリリースした日本語大規模言語モデル「[Japanese Stable LM 2 1.6B](https://ja.stability.ai/blog/japanese-stable-lm-2-16b)」に関するGoogle Colabでのサンプルコードです。\n",
        "\n",
        "T4 GPUなどGPUあり環境で実行してみてください（CPUでも動きます）。\n",
        "\n",
        "解説記事はこちら\n",
        "\n",
        "**Stability AI、日本語大規模言語モデル「Japanese Stable LM 2 1.6B」を発表。Google Colabのシークレット機能で使ってみた**\n",
        "https://note.com/aicu/n/nc4ebdc2a3eeb"
      ],
      "metadata": {
        "id": "kFxNJgjnZ2tB"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#!pip install transformers==4.40.1\n",
        "!pip install huggingface_hub\n",
        "!pip install torch\n",
        "!pip install tiktoken\n",
        "!pip install accelerate"
      ],
      "metadata": {
        "id": "Xn7h6jq1G4pN"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "from transformers import AutoTokenizer, AutoModelForCausalLM\n",
        "\n",
        "# 左側の「シークレット」から「HF_TOKEN」という環境変数を設定指定ください\n",
        "# https://twitter.com/GoogleColab/status/1719798406195867814\n",
        "# コードに直書きしたいひとは以下のように書いても動くと思います\n",
        "# token=\"(あなたのHugginFaceトークン)\"\n",
        "# !huggingface-cli login --token $token\n",
        "\n",
        "token = userdata.get('HF_TOKEN')\n"
      ],
      "metadata": {
        "id": "MiVQMISJK4jZ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qdiOurIfkZL1"
      },
      "outputs": [],
      "source": [
        "model_name = \"stabilityai/japanese-stablelm-2-base-1_6b\"\n",
        "tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=token)\n",
        "model = AutoModelForCausalLM.from_pretrained(\n",
        "    model_name,\n",
        "    token=token,\n",
        "    trust_remote_code=True,\n",
        "#   torch_dtype=torch.float16,\n",
        "    torch_dtype=torch.float32,  # ここをfloat32に変更しました\n",
        "    low_cpu_mem_usage=True,\n",
        "    device_map=\"auto\"\n",
        ")\n",
        "\n",
        "prompt = \"\"\"\n",
        "AI で科学研究を加速するには、\n",
        "\"\"\".strip()\n",
        "\n",
        "inputs = tokenizer(\n",
        "    prompt,\n",
        "    add_special_tokens=True,\n",
        "    return_tensors=\"pt\"\n",
        ").to(model.device)\n",
        "\n",
        "# this is for reproducibility.\n",
        "# feel free to change to get different result\n",
        "seed = 23\n",
        "torch.manual_seed(seed)\n",
        "\n",
        "tokens = model.generate(\n",
        "    **inputs,\n",
        "    max_new_tokens=128,\n",
        "    temperature=0.99,\n",
        "    top_p=0.95,\n",
        "    do_sample=True,\n",
        ")\n",
        "\n",
        "out = tokenizer.decode(tokens[0], skip_special_tokens=True)\n",
        "print(out)\n"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# シードは固定すると同じ結果が出力されます。以下の2行をコメントアウトすると毎回異なる結果になります。\n",
        "seed = 39\n",
        "torch.manual_seed(seed)\n",
        "# 自分のテキストで推論してみよう\n",
        "prompt = \"\"\"\n",
        "吾輩は猫である、名前は…\n",
        "\"\"\".strip()\n",
        "\n",
        "inputs = tokenizer(\n",
        "    prompt,\n",
        "    add_special_tokens=True,\n",
        "    return_tensors=\"pt\"\n",
        ").to(model.device)\n",
        "\n",
        "tokens = model.generate(\n",
        "    **inputs,\n",
        "    max_new_tokens=512,\n",
        "    temperature=0.99,\n",
        "    top_p=0.95,\n",
        "    do_sample=True,\n",
        ")\n",
        "\n",
        "out = tokenizer.decode(tokens[0], skip_special_tokens=True)\n",
        "print(out)\n"
      ],
      "metadata": {
        "id": "fEbMKnT9-dks"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "時同じくして、アドボケーター仲間のDELLさんが Japanese Stable LM 2 Instruct 1.6B の Google Colab版コードを公開してくれています。\n",
        "\n",
        "https://twitter.com/xqdior/status/1788459275809067424\n",
        "\n",
        "Gradioインタフェースで日本語チャットを楽しめます！"
      ],
      "metadata": {
        "id": "uhwT-UZfaa6r"
      }
    }
  ]
}