LLMにやさしい言語SuiにやさしいLLMを作るJupyterNote

sui-lang-Qwen3-8b.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "a8d5bf55-2da8-4545-81bc-be1119e61f78",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "## 設定"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "df7c3391-e202-41a9-b206-2985f80b8e57",
   "metadata": {},
   "outputs": [],
   "source": [
    "max_seq_length = 4096\n",
    "# model_name = \"unsloth/Qwen3-8B\"\n",
    "model_name = \"unsloth/Qwen3-8B-unsloth-bnb-4bit\"\n",
    "# chat_template = \"qwen3-instruct\"\n",
    "use_chat_template = False # only for instruct model\n",
    "max_steps = 15\n",
    "use_wandb = False\n",
    "output_dir = \"sui_coder_qwen3\"\n",
    "save_lora = False\n",
    "save_merged = False\n",
    "save_gguf = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3c9db459-f107-4dd7-9697-d29ab64d9453",
   "metadata": {},
   "outputs": [],
   "source": [
    "report_to = \"none\"\n",
    "if use_wandb:\n",
    "    import wandb\n",
    "    report_to = \"wandb\"\n",
    "    wandb.login()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "600b2552-3542-4c36-b279-b96294a488ab",
   "metadata": {},
   "source": [
    "## モジュールインストール"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5ce7c931-0d5c-4e9f-861a-19efa38c1889",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture\n",
    "import os, re\n",
    "if \"COLAB_\" in \"\".join(os.environ.keys()):\n",
    "    # Do this only in Colab notebooks! Otherwise use pip install unsloth\n",
    "    import torch; v = re.match(r\"[0-9]{1,}\\.[0-9]{1,}\", str(torch.__version__)).group(0)\n",
    "    xformers = \"xformers==\" + (\"0.0.33.post1\" if v==\"2.9\" else \"0.0.32.post2\" if v==\"2.8\" else \"0.0.29.post3\")\n",
    "    !pip install --no-deps bitsandbytes accelerate {xformers} peft trl triton cut_cross_entropy unsloth_zoo\n",
    "    !pip install sentencepiece protobuf \"datasets==4.3.0\" \"huggingface_hub>=0.34.0\" hf_transfer\n",
    "    !pip install --no-deps unsloth\n",
    "    !pip install transformers==4.57.1\n",
    "    !pip install --no-deps trl==0.22.2"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f4864526-e082-44b6-bbe0-7e9ceab526b7",
   "metadata": {},
   "source": [
    "## モデルの準備"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6296d8ac-dbca-4c9c-bbd0-41ebb075e467",
   "metadata": {},
   "outputs": [],
   "source": [
    "from unsloth import FastLanguageModel\n",
    "model, tokenizer = FastLanguageModel.from_pretrained(\n",
    "    model_name = model_name, \n",
    "    max_seq_length = max_seq_length,\n",
    "    dtype = None,\n",
    "    load_in_4bit = True,\n",
    "    full_finetuning = False,\n",
    "    # token = \"hf_...\", # use one if using gated models like meta-llama/Llama-2-7b-hf\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "73711cac-2230-4ec5-95c7-4fc4c791769d",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = FastLanguageModel.get_peft_model(\n",
    "    model,\n",
    "    r = 32, # Choose any number > 0 ! Suggested 8, 16, 32, 64, 128\n",
    "    target_modules = [\"q_proj\", \"k_proj\", \"v_proj\", \"o_proj\",\n",
    "                      \"gate_proj\", \"up_proj\", \"down_proj\",\n",
    "                     ],\n",
    "    lora_alpha = 32,\n",
    "    lora_dropout = 0, # Supports any, but = 0 is optimized\n",
    "    bias = \"none\",    # Supports any, but = \"none\" is optimized\n",
    "    # [NEW] \"unsloth\" uses 30% less VRAM, fits 2x larger batch sizes!\n",
    "    use_gradient_checkpointing = \"unsloth\", # True or \"unsloth\" for very long context\n",
    "    random_state = 3407,\n",
    "    use_rslora = False,  # We support rank stabilized LoRA\n",
    "    loftq_config = None, # And LoftQ\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "77765bd3-d775-49e5-8160-79caa8e65101",
   "metadata": {},
   "outputs": [],
   "source": [
    "if use_chat_template:\n",
    "  from unsloth.chat_templates import get_chat_template\n",
    "  tokenizer = get_chat_template(\n",
    "    tokenizer,\n",
    "    chat_template = chat_template,\n",
    "  )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9b30812f-5fb4-4202-b035-a9bbed0a3758",
   "metadata": {},
   "source": [
    "## データの準備"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e40e7dda-5640-4853-bc15-e5f1602379b5",
   "metadata": {},
   "source": [
    "### 読み込み"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3bdef308-62a8-4e8b-822b-4808254d8bda",
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import load_dataset\n",
    "dataset = load_dataset(\"kishida/sui-lang-examples\", split = \"train\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8f0457e5-51f0-48a1-b2cd-6ce663de4251",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(len(dataset))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b8e8dfeb-fe34-4dd1-8972-2c26ed02a1f0",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset[3]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4530e85c-c9af-4110-bdac-0dc1fa29b84e",
   "metadata": {},
   "source": [
    "### 変換"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "358add3e-68f4-449e-a9e9-6a3387645c91",
   "metadata": {},
   "outputs": [],
   "source": [
    "lang_spec = \"\"\"\n",
    "## Sui言語仕様\n",
    "\n",
    "Suiは純粋なロジック言語。1行1命令、識別子は連番のみ。\n",
    "\n",
    "### 命令\n",
    "= VAR VAL        代入\n",
    "+ R A B          加算 (R = A + B)\n",
    "- R A B          減算\n",
    "* R A B          乗算\n",
    "/ R A B          除算\n",
    "% R A B          剰余\n",
    "< R A B          小なり (R = 1 if A < B else 0)\n",
    "> R A B          大なり\n",
    "~ R A B          等価\n",
    "! R A            NOT\n",
    "& R A B          AND\n",
    "| R A B          OR\n",
    "? COND LABEL     条件ジャンプ（LABELは整数のみ）\n",
    "@ LABEL          無条件ジャンプ（LABELは整数のみ）\n",
    ": LABEL          ラベル定義（LABELは整数のみ）\n",
    "# ID ARGC {      関数定義\n",
    "}                関数終了\n",
    "$ R FN ARGS...   関数呼び出し\n",
    "^ VAL            return\n",
    "[ VAR SIZE       配列作成\n",
    "] R ARR IDX      配列読み取り\n",
    "{ ARR IDX VAL    配列書き込み\n",
    ". VAL            出力\n",
    ", VAR            入力\n",
    "; COMMENT        コメント\n",
    "\n",
    "### 変数\n",
    "v0, v1, v2...    ローカル変数\n",
    "g0, g1, g2...    グローバル変数（状態、UIからアクセス可能）\n",
    "a0, a1, a2...    関数引数\n",
    "\n",
    "### ラベル（重要）\n",
    "- **LABELは必ず整数**（例: `0`, `1`, `2`）\n",
    "- `: start_label` のような **文字列ラベルは禁止**\n",
    "\n",
    "\n",
    "### 例\n",
    "\n",
    "```sui\n",
    "; 与えられた整数を3倍する\n",
    "# 0 1 {\n",
    "    * v0 a0 3\n",
    "    ^ v0\n",
    "}\n",
    "```\n",
    "\n",
    "```sui\n",
    "; Print 1 to 10\n",
    "= v0 1\n",
    ": 0\n",
    "> v1 v0 10\n",
    "? v1 1\n",
    ". v0\n",
    "+ v0 v0 1\n",
    "@ 0\n",
    ": 1\n",
    "```\n",
    "\n",
    "```sui\n",
    "; Sum all elements in a list\n",
    "; Calculate sum of array [10, 20, 30, 40, 50]\n",
    "\n",
    "; Create array (size 5)\n",
    "[ g0 5\n",
    "\n",
    "; Set elements\n",
    "{ g0 0 10\n",
    "{ g0 1 20\n",
    "{ g0 2 30\n",
    "{ g0 3 40\n",
    "{ g0 4 50\n",
    "\n",
    "; Calculate sum\n",
    "; g1 = sum, v0 = index, g2 = size\n",
    "= g1 0\n",
    "= g2 5\n",
    "= v0 0\n",
    "\n",
    "; Loop start\n",
    ": 0\n",
    "< v1 v0 g2\n",
    "! v2 v1\n",
    "? v2 1\n",
    "\n",
    "; Get arr[v0] and add to sum\n",
    "] v3 g0 v0\n",
    "+ g1 g1 v3\n",
    "\n",
    "; Increment index\n",
    "+ v0 v0 1\n",
    "@ 0\n",
    "\n",
    "; Loop end\n",
    ": 1\n",
    "; Output result\n",
    ". \"Sum:\"\n",
    ". g1\n",
    "```\n",
    "\n",
    "### 注意\n",
    "コメントは ; で始める。#でのコメントは禁止\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1bca54d4-85fe-4fc3-98e1-7dbe780c179b",
   "metadata": {},
   "outputs": [],
   "source": [
    "template = \"\"\"次の仕様を満たす{}を実装してください。\n",
    "{}\n",
    "\"\"\"\n",
    "system_prompt = \"\"\"あなたはSui言語を使ったアプリケーション開発者です。\n",
    "ユーザーに与えられた課題をSui言語の命令で実装します。\n",
    "コードブロックに入れたりせず実装のみを出力してください。\n",
    "\n",
    "Sui言語の仕様は次のとおり\n",
    "{}\n",
    "\"\"\".format(lang_spec)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "696095de-afba-4e85-bcaa-5d23366e6b89",
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_conv(example):\n",
    "    prompt = template.format(f\"{example['arg_num']}引数関数\" if example[\"type\"] == \"Function\" else \"スクリプト\", example[\"problem\"])\n",
    "    return {\"conversations\": [\n",
    "        {\"content\": system_prompt, \"role\": \"system\"},\n",
    "        {\"content\": prompt, \"role\": \"user\"},\n",
    "        {\"content\": example[\"sui_code\"], \"role\": \"assistant\"}]}\n",
    "dataset = dataset.map(make_conv, batched = False,)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3b6836d5-a354-49f2-8c77-00fc2e3faa26",
   "metadata": {},
   "outputs": [],
   "source": [
    "def formatting_prompts_func(examples):\n",
    "    convos = examples[\"conversations\"]\n",
    "    texts = [tokenizer.apply_chat_template(conv, tokenize = False, add_generation_prompt = False) for conv in convos]\n",
    "    return {\"text\": texts, }\n",
    "dataset = dataset.map(formatting_prompts_func, batched = True,)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b7093231-40ef-4235-9bbb-1441e03d6b2c",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset[30][\"text\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d38304a8-b44f-42f8-aa23-caccf3523a11",
   "metadata": {},
   "source": [
    "## 学習"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b78eeff9-1c3a-4cd9-bf8a-6837a14f338d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from trl import SFTConfig, SFTTrainer\n",
    "from transformers import DataCollatorForSeq2Seq\n",
    "trainer = SFTTrainer(\n",
    "    model = model,\n",
    "    tokenizer = tokenizer,\n",
    "    train_dataset = dataset,\n",
    "    max_seq_length = max_seq_length,\n",
    "    data_collator = DataCollatorForSeq2Seq(tokenizer = tokenizer),\n",
    "    packing = False, # Can make training 5x faster for short sequences.\n",
    "    args = SFTConfig(\n",
    "        dataset_text_field = \"text\",\n",
    "        per_device_train_batch_size = 2,\n",
    "        gradient_accumulation_steps = 4,\n",
    "        warmup_steps = 5,\n",
    "        #num_train_epochs = 1, # Set this for 1 full training run.\n",
    "        max_steps = max_steps,\n",
    "        learning_rate = 2e-4,\n",
    "        logging_steps = 1,\n",
    "        optim = \"adamw_8bit\",\n",
    "        weight_decay = 0.001,\n",
    "        lr_scheduler_type = \"linear\",\n",
    "        seed = 3407,\n",
    "        output_dir = output_dir,\n",
    "        report_to = report_to, # Use TrackIO/WandB etc\n",
    "    ),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ed5c3d82-1a7e-4b66-96b1-74fb8ee808e0",
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer.decode(trainer.train_dataset[3][\"input_ids\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9e78cf40-fb12-42dd-8059-2b3b3f0bb488",
   "metadata": {},
   "source": [
    "### 学習開始"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "764f3e5d-72be-4367-8c99-d30712ff12fc",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "# @title Show current memory stats\n",
    "gpu_stats = torch.cuda.get_device_properties(0)\n",
    "start_gpu_memory = round(torch.cuda.max_memory_reserved() / 1024 / 1024 / 1024, 3)\n",
    "max_memory = round(gpu_stats.total_memory / 1024 / 1024 / 1024, 3)\n",
    "print(f\"GPU = {gpu_stats.name}. Max memory = {max_memory} GB.\")\n",
    "print(f\"{start_gpu_memory} GB of memory reserved.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3632ac67-9699-48c9-b8af-b1337de64d08",
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer_stats = trainer.train()\n",
    "if use_wandb: wandb.finish()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0d87c1e1-6e42-4e84-8e5e-600ef6fafad5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# @title Show final memory and time stats\n",
    "used_memory = round(torch.cuda.max_memory_reserved() / 1024 / 1024 / 1024, 3)\n",
    "used_memory_for_lora = round(used_memory - start_gpu_memory, 3)\n",
    "used_percentage = round(used_memory / max_memory * 100, 3)\n",
    "lora_percentage = round(used_memory_for_lora / max_memory * 100, 3)\n",
    "print(f\"{trainer_stats.metrics['train_runtime']} seconds used for training.\")\n",
    "print(\n",
    "    f\"{round(trainer_stats.metrics['train_runtime']/60, 2)} minutes used for training.\"\n",
    ")\n",
    "print(f\"Peak reserved memory = {used_memory} GB.\")\n",
    "print(f\"Peak reserved memory for training = {used_memory_for_lora} GB.\")\n",
    "print(f\"Peak reserved memory % of max memory = {used_percentage} %.\")\n",
    "print(f\"Peak reserved memory for training % of max memory = {lora_percentage} %.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d659802f-6623-4899-9335-fd44db303421",
   "metadata": {},
   "source": [
    "## 推論"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fdfa15af-b13c-4789-a1f1-e84c8fb39ac3",
   "metadata": {},
   "outputs": [],
   "source": [
    "sample = dataset[3]\n",
    "messages = [\n",
    "    {\"role\": \"system\", \"content\": system_prompt},\n",
    "    {\"role\": \"user\", \"content\": \"/nothink\\n\" + template.format(\"2引数関数\", sample[\"problem\"])}]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5e12a6af-79a7-4c83-aede-abc1a8a93b94",
   "metadata": {},
   "outputs": [],
   "source": [
    "sample[\"sui_code\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "936d1ff9-fbed-4682-9ed7-60703daacec0",
   "metadata": {},
   "outputs": [],
   "source": [
    "FastLanguageModel.for_inference(model) # Enable native 2x faster inference\n",
    "inputs = tokenizer.apply_chat_template(\n",
    "    messages,\n",
    "    tokenize = True,\n",
    "    add_generation_prompt = True, # Must add for generation\n",
    "    return_tensors = \"pt\",\n",
    ").to(\"cuda\")\n",
    "from transformers import TextStreamer\n",
    "text_streamer = TextStreamer(tokenizer, skip_prompt = True)\n",
    "_ = model.generate(input_ids = inputs, streamer = text_streamer, max_new_tokens = 128,\n",
    "                   use_cache = True, temperature = 1.5, min_p = 0.1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2ccf3578-8fdf-4fb0-8893-5244ca92949b",
   "metadata": {},
   "source": [
    "## 保存"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ed39e906-d48f-401f-bc62-c5c1a83edadc",
   "metadata": {},
   "outputs": [],
   "source": [
    "if save_lora:\n",
    "    model.save_pretrained(output_dir + \"_lora\")  # Local saving\n",
    "    tokenizer.save_pretrained(output_dir + \"_lora\")\n",
    "if save_merged: model.save_pretrained_merged(output_dir + \"_merged\", tokenizer,)\n",
    "if save_gguf: model.save_pretrained_gguf(output_dir + \"_gguf\", tokenizer, quantization_method = \"q4_k_m\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "95a34ebb-f127-48e2-b33b-53c3c9fe3a1c",
   "metadata": {},
   "outputs": [],
   "source": [
    "%ls *.gguf"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "gpuType": "T4",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}