{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Qwen2.5-7B QLoRA Training on Colab\n",
    "\n",
    "Google Colab Pro (A100) での学習用ノートブック\n",
    "\n",
    "**推奨**: Colab Pro ($10/月) 以上、A100 GPU"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. 環境セットアップ"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# GPU確認\n",
    "!nvidia-smi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Google Driveマウント（チェックポイント保存用）\n",
    "from google.colab import drive\n",
    "drive.mount('/content/drive')\n",
    "\n",
    "# 作業ディレクトリ作成\n",
    "!mkdir -p /content/drive/MyDrive/qwen-training/checkpoints\n",
    "!mkdir -p /content/drive/MyDrive/qwen-training/output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 依存関係インストール\n",
    "!pip install -q torch==2.2.0 torchvision==0.17.0\n",
    "!pip install -q transformers==4.46.0 datasets peft==0.13.0 trl==0.11.0\n",
    "!pip install -q bitsandbytes accelerate huggingface_hub safetensors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# HuggingFaceログイン\n",
    "from huggingface_hub import login\n",
    "login()  # トークンを入力"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. 設定"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 設定\n",
    "BASE_MODEL = \"Qwen/Qwen2.5-7B-Instruct\"\n",
    "OUTPUT_MODEL_ID = \"hajimemat/qwen2.5-7b-glaive-fc-lora-colab\"  # 変更可\n",
    "DATASET_NAME = \"glaiveai/glaive-function-calling-v2\"\n",
    "\n",
    "# Google Driveに保存\n",
    "CHECKPOINT_DIR = \"/content/drive/MyDrive/qwen-training/checkpoints\"\n",
    "FINAL_OUTPUT_DIR = \"/content/drive/MyDrive/qwen-training/output\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. データセット準備"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import load_dataset\n",
    "\n",
    "def convert_glaive_to_chatml(example):\n",
    "    parts = []\n",
    "    if example.get(\"system\"):\n",
    "        parts.append(f\"<|im_start|>system\\n{example['system']}<|im_end|>\")\n",
    "    \n",
    "    chat = example.get(\"chat\", \"\")\n",
    "    if chat:\n",
    "        current_role = None\n",
    "        current_content = []\n",
    "        for line in chat.split(\"\\n\"):\n",
    "            line = line.strip()\n",
    "            if line.startswith(\"USER:\"):\n",
    "                if current_role and current_content:\n",
    "                    content = \"\\n\".join(current_content).strip()\n",
    "                    if content:\n",
    "                        parts.append(f\"<|im_start|>{current_role}\\n{content}<|im_end|>\")\n",
    "                current_role = \"user\"\n",
    "                current_content = [line[5:].strip()]\n",
    "            elif line.startswith(\"ASSISTANT:\"):\n",
    "                if current_role and current_content:\n",
    "                    content = \"\\n\".join(current_content).strip()\n",
    "                    if content:\n",
    "                        parts.append(f\"<|im_start|>{current_role}\\n{content}<|im_end|>\")\n",
    "                current_role = \"assistant\"\n",
    "                current_content = [line[10:].strip()]\n",
    "            elif current_role:\n",
    "                current_content.append(line)\n",
    "        if current_role and current_content:\n",
    "            content = \"\\n\".join(current_content).strip()\n",
    "            if content:\n",
    "                parts.append(f\"<|im_start|>{current_role}\\n{content}<|im_end|>\")\n",
    "    return {\"text\": \"\\n\".join(parts)}\n",
    "\n",
    "print(f\"Loading dataset: {DATASET_NAME}\")\n",
    "dataset = load_dataset(DATASET_NAME, split=\"train\")\n",
    "print(f\"Original: {len(dataset)} examples\")\n",
    "\n",
    "dataset = dataset.map(convert_glaive_to_chatml, remove_columns=dataset.column_names, num_proc=4)\n",
    "dataset = dataset.filter(lambda x: len(x[\"text\"]) > 50)\n",
    "dataset = dataset.shuffle(seed=42)\n",
    "split = dataset.train_test_split(test_size=0.02, seed=42)\n",
    "\n",
    "print(f\"Train: {len(split['train'])}, Test: {len(split['test'])}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. モデル準備"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, TrainingArguments\n",
    "from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training\n",
    "\n",
    "# QLoRA量子化設定\n",
    "bnb_config = BitsAndBytesConfig(\n",
    "    load_in_4bit=True,\n",
    "    bnb_4bit_compute_dtype=torch.bfloat16,\n",
    "    bnb_4bit_quant_type=\"nf4\",\n",
    "    bnb_4bit_use_double_quant=True,\n",
    ")\n",
    "\n",
    "# LoRA設定\n",
    "lora_config = LoraConfig(\n",
    "    r=64,\n",
    "    lora_alpha=16,\n",
    "    lora_dropout=0.05,\n",
    "    target_modules=[\"q_proj\", \"k_proj\", \"v_proj\", \"o_proj\", \"gate_proj\", \"up_proj\", \"down_proj\"],\n",
    "    bias=\"none\",\n",
    "    task_type=\"CAUSAL_LM\",\n",
    ")\n",
    "\n",
    "# トークナイザー\n",
    "tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)\n",
    "tokenizer.padding_side = \"right\"\n",
    "if tokenizer.pad_token is None:\n",
    "    tokenizer.pad_token = tokenizer.eos_token\n",
    "\n",
    "# モデル\n",
    "print(f\"Loading model: {BASE_MODEL}\")\n",
    "model = AutoModelForCausalLM.from_pretrained(\n",
    "    BASE_MODEL,\n",
    "    quantization_config=bnb_config,\n",
    "    device_map=\"auto\",\n",
    "    attn_implementation=\"sdpa\",\n",
    "    trust_remote_code=True,\n",
    ")\n",
    "\n",
    "model = prepare_model_for_kbit_training(model)\n",
    "model = get_peft_model(model, lora_config)\n",
    "model.print_trainable_parameters()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. 学習実行"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from trl import SFTTrainer\n",
    "\n",
    "training_args = TrainingArguments(\n",
    "    output_dir=CHECKPOINT_DIR,\n",
    "    num_train_epochs=1,\n",
    "    per_device_train_batch_size=4,\n",
    "    per_device_eval_batch_size=4,\n",
    "    gradient_accumulation_steps=4,\n",
    "    learning_rate=2e-4,\n",
    "    weight_decay=0.01,\n",
    "    warmup_ratio=0.03,\n",
    "    lr_scheduler_type=\"cosine\",\n",
    "    optim=\"paged_adamw_8bit\",\n",
    "    bf16=True,\n",
    "    logging_steps=10,\n",
    "    save_steps=200,\n",
    "    save_total_limit=3,\n",
    "    eval_strategy=\"steps\",\n",
    "    eval_steps=200,\n",
    "    report_to=\"none\",\n",
    "    gradient_checkpointing=True,\n",
    "    save_safetensors=True,\n",
    ")\n",
    "\n",
    "trainer = SFTTrainer(\n",
    "    model=model,\n",
    "    train_dataset=split[\"train\"],\n",
    "    eval_dataset=split[\"test\"],\n",
    "    args=training_args,\n",
    "    peft_config=lora_config,\n",
    "    tokenizer=tokenizer,\n",
    "    max_seq_length=1024,\n",
    "    packing=False,\n",
    "    dataset_text_field=\"text\",\n",
    ")\n",
    "\n",
    "# チェックポイントから再開\n",
    "import os\n",
    "resume_from = None\n",
    "if os.path.exists(CHECKPOINT_DIR):\n",
    "    checkpoints = [d for d in os.listdir(CHECKPOINT_DIR) if d.startswith(\"checkpoint-\")]\n",
    "    if checkpoints:\n",
    "        latest = max(checkpoints, key=lambda x: int(x.split(\"-\")[1]))\n",
    "        resume_from = os.path.join(CHECKPOINT_DIR, latest)\n",
    "        print(f\"Resuming from: {resume_from}\")\n",
    "\n",
    "# 学習開始\n",
    "trainer.train(resume_from_checkpoint=resume_from)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6. 保存とアップロード"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ローカル保存\n",
    "print(f\"Saving to {FINAL_OUTPUT_DIR}\")\n",
    "trainer.save_model(FINAL_OUTPUT_DIR)\n",
    "tokenizer.save_pretrained(FINAL_OUTPUT_DIR)\n",
    "\n",
    "# HuggingFaceにアップロード\n",
    "print(f\"Uploading to {OUTPUT_MODEL_ID}\")\n",
    "try:\n",
    "    trainer.model.push_to_hub(OUTPUT_MODEL_ID, private=True)\n",
    "    tokenizer.push_to_hub(OUTPUT_MODEL_ID, private=True)\n",
    "    print(f\"Done! https://huggingface.co/{OUTPUT_MODEL_ID}\")\n",
    "except Exception as e:\n",
    "    print(f\"Upload failed: {e}\")\n",
    "    print(\"Model saved locally in Google Drive\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 7. クイックテスト（オプション）"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 簡単な推論テスト\n",
    "from peft import PeftModel\n",
    "\n",
    "test_prompt = \"\"\"<|im_start|>system\n",
    "You are a helpful assistant with access to functions.\n",
    "<|im_end|>\n",
    "<|im_start|>user\n",
    "What's the weather in Tokyo?\n",
    "<|im_end|>\n",
    "<|im_start|>assistant\n",
    "\"\"\"\n",
    "\n",
    "inputs = tokenizer(test_prompt, return_tensors=\"pt\").to(model.device)\n",
    "outputs = model.generate(**inputs, max_new_tokens=200, temperature=0.7)\n",
    "print(tokenizer.decode(outputs[0], skip_special_tokens=False))"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "gpuType": "A100",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}