{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "license-cell",
   "metadata": {},
   "source": [
    "**Author:** Siwen Yu (yusiwen@gmail.com)\n",
    "**License:** MIT"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "02667909",
   "metadata": {},
   "source": [
    "# LoRA 微调 BERT-base：原理与代码逐块详解\n",
    "\n",
    "本 Notebook 逐块解析两个核心 Python 文件：\n",
    "- **`train_lora.py`** — 使用 LoRA 对 BERT-base 做二分类（情感分析）的**训练脚本**\n",
    "- **`inference.py`** — 加载训练好的 LoRA 适配器做**推理**的脚本\n",
    "\n",
    "---\n",
    "## 整体工作流\n",
    "```\n",
    "训练: BERT-base (冻结) → 注入 LoRA 适配器 → Rotten Tomatoes 情感数据 → 保存适配器 (~2MB)\n",
    "推理: BERT-base (冻结) → 加载 LoRA 适配器 → 预测新句子情感\n",
    "```\n",
    "\n",
    "> **注意**：BERT-base 仅 110M 参数，不需要 4-bit 量化。本示例在普通 GPU（~2GB VRAM）即可运行。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f9986a27",
   "metadata": {},
   "source": [
    "---\n",
    "# 第一部分：train_lora.py — 训练脚本"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3b3e61fc",
   "metadata": {},
   "source": [
    "## 1.1 导入依赖\n",
    "\n",
    "| 库 | 作用 |\n",
    "|---|---|\n",
    "| `torch` | PyTorch 核心框架 |\n",
    "| `datasets` | HuggingFace 数据集加载 |\n",
    "| `transformers` | BERT 模型、分词器、Trainer |\n",
    "| `peft` | **LoRA 核心实现**：`LoraConfig`、`get_peft_model`、`TaskType` |\n",
    "| `evaluate` | HuggingFace 评估库 |\n",
    "| `numpy` | 数值计算 |\n",
    "\n",
    "**与 Qwen 的区别**：这里没有 `BitsAndBytesConfig`（无量化），没有 `prepare_model_for_kbit_training`，因为 BERT-base 足够小。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e2f7deee",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from datasets import load_dataset\n",
    "from transformers import AutoTokenizer, AutoModelForSequenceClassification, Trainer, TrainingArguments\n",
    "from peft import LoraConfig, get_peft_model, TaskType\n",
    "import evaluate\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8e760567",
   "metadata": {},
   "source": [
    "## 1.2 GPU 检测\n",
    "\n",
    "检测 CUDA 可用性。BERT-base 号称 CPU 也能跑，但为了演示效率仍推荐 GPU。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "68c650ea",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"CUDA Available: {torch.cuda.is_available()}\")\n",
    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "print(f\"Using device: {torch.cuda.get_device_name(0) if device == 'cuda' else 'CPU'}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14ac2ef1",
   "metadata": {},
   "source": [
    "## 1.3 加载数据集与 Tokenizer\n",
    "\n",
    "### Rotten Tomatoes 数据集\n",
    "- 来源：烂番茄电影评论\n",
    "- 任务：二分类（Positive / Negative）\n",
    "- 包含 `train`、`validation`、`test` 三个 split\n",
    "\n",
    "### Tokenizer 配置\n",
    "- `padding=\"max_length\"` — 将所有句子填充到统一长度（非 Qwen 的 `padding=False`）\n",
    "- `truncation=True` — 超过 128 token 的句子截断\n",
    "- `max_length=128` — BERT 的典型输入长度\n",
    "\n",
    "**为什么这里要 padding**：Trainer 对分类任务需要批次内 tensor 形状一致；而 Qwen 的 DataCollatorForLanguageModeling 会动态 padding。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dcc85cc9",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = load_dataset(\"cornell-movie-review-data/rotten_tomatoes\")\n",
    "tokenizer = AutoTokenizer.from_pretrained(\"bert-base-uncased\")\n",
    "\n",
    "def tokenize_function(examples):\n",
    "    return tokenizer(examples[\"text\"], padding=\"max_length\", truncation=True, max_length=128)\n",
    "\n",
    "tokenized_datasets = dataset.map(tokenize_function, batched=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cbb20a87",
   "metadata": {},
   "source": [
    "## 1.4 加载基础模型\n",
    "\n",
    "`AutoModelForSequenceClassification.from_pretrained(\"bert-base-uncased\", num_labels=2)`\n",
    "- BERT-base：12 层 Transformer，768 hidden size，110M 参数\n",
    "- `num_labels=2` — 替换 BERT 的 pooler 输出层为 2 分类头（Positive/Negative）\n",
    "- **无量化**：110M 参数全精度 FP32 也仅 ~440MB\n",
    "\n",
    "> 与 Qwen 的关键区别：Qwen 需要 4-bit 量化才能放进 12GB 显存；BERT 完全不需要。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2523cdf1-2336-4a41-a840-78606da04bf0",
   "metadata": {},
   "source": [
    "## 1.4 加载基础模型\n",
    "\n",
    "`AutoModelForSequenceClassification.from_pretrained(\"bert-base-uncased\", num_labels=2)`\n",
    "- BERT-base：12 层 Transformer，768 hidden size，110M 参数\n",
    "- `num_labels=2` — 替换 BERT 的 pooler 输出层为 2 分类头（Positive/Negative）\n",
    "- **无量化**：110M 参数全精度 FP32 也仅 ~440MB\n",
    "\n",
    "> 与 Qwen 的关键区别：Qwen 需要 4-bit 量化才能放进 12GB 显存；BERT 完全不需要。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "548a71ae",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = AutoModelForSequenceClassification.from_pretrained(\"bert-base-uncased\", num_labels=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1453286f-4f33-4735-ad76-fe6762903066",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "### 关于num_labels=2\n",
    "\n",
    "传入 num_labels=2 时，会移除 BERT 原本用于预训练的输出层，并在 BERT 的 pooler_output（即代表整句话语义的 [CLS] 向量，维度为 768）后面，添加一个全新的线性全连接层（Linear Layer）。这个新层的结构是 Linear(in_features=768, out_features=2)。\n",
    "\n",
    "如果去掉 num_labels=2 这个设置，代码依然可以运行，但模型会发生根本性的改变：\n",
    "\n",
    "- 模型会自动读取 bert-base-uncased 官方配置文件（config.json）中的默认设置。\n",
    "- BERT 的默认设置正好也是 num_labels: 2（对应预训练时的“下一句预测”任务）。⚠️ 致命隐患：虽然它默认也是 2，但它会尝试加载原版预训练模型的权重。这会导致你的分类头权重不是随机初始化的，而是带有预训练残留权重的，这会严重干扰你自己的二分类微调。\n",
    "\n",
    "这个分类头也是微调的重要训练学习的部分，显式指定 num_labels 的核心核心作用是显式触发分类头的权重初始化（Weight Initialization），训练开始时用随机数初始化，逐步学习分类。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a609df36",
   "metadata": {},
   "source": [
    "## 1.5 LoRA 配置与注入\n",
    "\n",
    "### 与 Qwen 的 LoRA 差异\n",
    "\n",
    "| 参数 | BERT（本文件） | Qwen（之前的训练） | 原因 |\n",
    "|---|---|---|---|\n",
    "| `task_type` | `SEQ_CLS` | `CAUSAL_LM` | BERT 是编码器分类；Qwen 是解码器生成 |\n",
    "| `r=8` | 相同 | 相同 | 标准配置 |\n",
    "| `lora_alpha=32` | 32 | 16 | 不同缩放（α/r = 4 vs 2） |\n",
    "| `target_modules` | `[\"query\", \"value\"]` | `[\"q_proj\", \"v_proj\", \"k_proj\", \"o_proj\"]` | BERT 的命名是 `query`/`value`（无 `_proj` 后缀），且不加 K/O |\n",
    "| `lora_dropout=0.1` | 0.1 | 0.05 | BERT 任务小，略高 dropout |\n",
    "\n",
    "**原理回顾**：\n",
    "$$W' = W + BA, \\quad B \\in \\mathbb{R}^{d \\times r}, A \\in \\mathbb{R}^{r \\times k}$$\n",
    "\n",
    "BERT 的 attention 子层中，`query` 和 `value` 投影矩阵被插入 LoRA。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "de3c3a21",
   "metadata": {},
   "outputs": [],
   "source": [
    "lora_config = LoraConfig(\n",
    "    task_type=TaskType.SEQ_CLS, \n",
    "    r=8,\n",
    "    lora_alpha=32,\n",
    "    lora_dropout=0.1,\n",
    "    target_modules=[\"query\", \"value\"]\n",
    ")\n",
    "\n",
    "model = get_peft_model(model, lora_config)\n",
    "model.print_trainable_parameters()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e9b4c822",
   "metadata": {},
   "source": [
    "## 1.6 评估指标\n",
    "\n",
    "使用 HuggingFace `evaluate` 库加载 accuracy 指标。\n",
    "\n",
    "`compute_metrics` 函数会在每个 epoch 结束后自动调用，输入是 `(logits, labels)`，输出打分字典。\n",
    "\n",
    "**这与 Qwen 不同**：Qwen 的训练脚本没有评估步骤（纯训练），而 BERT 的分类任务天然适合在验证集上监控准确率。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "07bd95c1",
   "metadata": {},
   "outputs": [],
   "source": [
    "metric = evaluate.load(\"accuracy\")\n",
    "def compute_metrics(eval_pred):\n",
    "    logits, labels = eval_pred\n",
    "    predictions = np.argmax(logits, axis=-1)\n",
    "    return metric.compute(predictions=predictions, references=labels)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b825f0cc",
   "metadata": {},
   "source": [
    "## 1.7 训练参数\n",
    "\n",
    "| 参数 | 值 | 与 Qwen 的差异 | 原因 |\n",
    "|---|---|---|---|\n",
    "| `per_device_train_batch_size=32` | 32 vs 2 | BERT 仅 110M 参数，batch 可以很大 |\n",
    "| `learning_rate=2e-4` | 相同 | LoRA 通用学习率 |\n",
    "| `fp16=True` | 相同 | 混合精度通用 |\n",
    "| `num_train_epochs=1` | 相同 | 演示用 |\n",
    "| `eval_strategy=\"epoch\"` | 有 vs 无 | 分类任务需要评估 |\n",
    "| `load_best_model_at_end=True` | 有 vs 无 | 自动加载最佳 checkpoint |\n",
    "| `weight_decay=0.01` | 有 vs 无 | 防过拟合（L2 正则化） |"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c80f455b",
   "metadata": {},
   "outputs": [],
   "source": [
    "training_args = TrainingArguments(\n",
    "    output_dir=\"./lora_results\",\n",
    "    learning_rate=2e-4,\n",
    "    per_device_train_batch_size=32,\n",
    "    per_device_eval_batch_size=32,\n",
    "    num_train_epochs=1,\n",
    "    weight_decay=0.01,\n",
    "    eval_strategy=\"epoch\",\n",
    "    save_strategy=\"epoch\",\n",
    "    load_best_model_at_end=True,\n",
    "    logging_steps=10,\n",
    "    fp16=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "95e7dae2",
   "metadata": {},
   "source": [
    "## 1.8 训练执行与保存\n",
    "\n",
    "### Trainer（与 Qwen 的关键区别）\n",
    "- **有 `eval_dataset`**：Qwen 的 Trainer 只传了 `train_dataset`，这里额外传入验证集\n",
    "- **有 `compute_metrics`**：每个 epoch 结束后计算 accuracy\n",
    "- **无 `DataCollatorForLanguageModeling`**：分类任务不需要 LM 的 collator（默认 `DataCollatorWithPadding`）\n",
    "\n",
    "### 保存\n",
    "- `model.save_pretrained(\"./my_saved_lora\")` — 只保存 LoRA 权重（~2MB）\n",
    "- `tokenizer.save_pretrained(\"./my_saved_lora\")` — 同时保存 tokenizer，供推理时直接加载"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "971d123e",
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer = Trainer(\n",
    "    model=model,\n",
    "    args=training_args,\n",
    "    train_dataset=tokenized_datasets[\"train\"],\n",
    "    eval_dataset=tokenized_datasets[\"validation\"],\n",
    "    compute_metrics=compute_metrics,\n",
    ")\n",
    "\n",
    "print(\"\\n--- Starting LoRA Training and Validation ---\")\n",
    "trainer.train()\n",
    "\n",
    "model.save_pretrained(\"./my_saved_lora\")\n",
    "tokenizer.save_pretrained(\"./my_saved_lora\")\n",
    "print(\"LoRA Model successfully trained and saved!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "04610435",
   "metadata": {},
   "source": [
    "---\n",
    "# 第二部分：inference.py — 推理脚本\n",
    "\n",
    "训练完成后，用推理脚本加载适配器并在自定义样本上测试。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "99af8175",
   "metadata": {},
   "source": [
    "## 2.1 导入与路径配置\n",
    "\n",
    "与 Qwen 的 `chat_qlora.py` 模式一致：\n",
    "- `PeftModel` — 加载 LoRA 适配器\n",
    "- `PeftConfig` — 可选，用于读取适配器配置（本脚本未使用但导入了）\n",
    "- `base_model_path = \"bert-base-uncased\"` — 原始基础模型\n",
    "- `lora_weights_path = \"./my_saved_lora\"` — 训练保存的适配器 + tokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6a6c7385",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from transformers import AutoTokenizer, AutoModelForSequenceClassification\n",
    "from peft import PeftModel, PeftConfig\n",
    "\n",
    "base_model_path = \"bert-base-uncased\"\n",
    "lora_weights_path = \"./my_saved_lora\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1e4bd467",
   "metadata": {},
   "source": [
    "## 2.2 加载 Tokenizer 与基础模型\n",
    "\n",
    "### Tokenizer\n",
    "- **从 `lora_weights_path` 加载**：因为训练时我们用 `tokenizer.save_pretrained()` 把 tokenizer 和适配器存在同一目录\n",
    "- 这样推理时无需再次下载 tokenizer\n",
    "\n",
    "### 基础模型\n",
    "- 与训练时完全一致：`bert-base-uncased`，`num_labels=2`\n",
    "- **无量化**：BERT-base 不需要量化即可推理"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "425075f9",
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer = AutoTokenizer.from_pretrained(lora_weights_path)\n",
    "base_model = AutoModelForSequenceClassification.from_pretrained(base_model_path, num_labels=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bd6a413a",
   "metadata": {},
   "source": [
    "## 2.3 挂载 LoRA 适配器\n",
    "\n",
    "`PeftModel.from_pretrained(base_model, lora_weights_path)` 做的事情：\n",
    "1. 读取 `adapter_config.json` 获取 LoRA 配置（`r=8`、`target_modules` 等）\n",
    "2. 加载 `adapter_model.safetensors` 中的 BA 矩阵\n",
    "3. 注入到 BERT 的 `query` 和 `value` 层\n",
    "\n",
    "`model.to(\"cuda\")` — 将模型移到 GPU（与 Qwen 的 `device_map=\"auto\"` 方式不同，这里手动指定）\n",
    "\n",
    "`model.eval()` — 关闭 dropout，进入推理模式。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b2d8af1f",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = PeftModel.from_pretrained(base_model, lora_weights_path)\n",
    "model = model.to(\"cuda\")\n",
    "model.eval()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "19b3f45c",
   "metadata": {},
   "source": [
    "## 2.4 运行推理\n",
    "\n",
    "### 测试样本\n",
    "两条典型的正面/负面电影评论：\n",
    "1. \"This movie was an absolute masterpiece...\" → 预期 Positive (1)\n",
    "2. \"A total waste of time...\" → 预期 Negative (0)\n",
    "\n",
    "### 推理流程\n",
    "```\n",
    "句子 → tokenizer → input_ids (1×128) → BERT → logits (1×2) → argmax → label\n",
    "```\n",
    "\n",
    "### 与 Qwen 推理的区别\n",
    "| | BERT（本文件） | Qwen（chat_qlora.py） |\n",
    "|---|---|---|\n",
    "| 任务 | 分类（序列级输出） | 生成（token 级自回归） |\n",
    "| 输出 | `logits` → `argmax` | `generate()` → token 序列 |\n",
    "| 前向方式 | `model(**inputs)` | `model.generate(**inputs)` |\n",
    "| 显存需求 | ~1.5GB | ~6-8GB |"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "101110b1",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_sentences = [\n",
    "    \"This movie was an absolute masterpiece with breathtaking visuals!\",\n",
    "    \"A total waste of time. The acting was horrible and the plot made no sense.\"\n",
    "]\n",
    "\n",
    "print(\"\\n--- Running Validation Inference ---\")\n",
    "for text in test_sentences:\n",
    "    inputs = tokenizer(text, return_tensors=\"pt\", padding=True, truncation=True, max_length=128).to(\"cuda\")\n",
    "    \n",
    "    with torch.no_grad():\n",
    "        outputs = model(**inputs)\n",
    "        prediction = torch.argmax(outputs.logits, dim=-1).item()\n",
    "        \n",
    "    sentiment = \"Positive\" if prediction == 1 else \"Negative\"\n",
    "    print(f\"Review: '{text}' -> Predicted Sentiment: {sentiment}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ccef289f",
   "metadata": {},
   "source": [
    "---\n",
    "# 总结\n",
    "\n",
    "## 关键技术栈\n",
    "\n",
    "| 技术 | 作用 |\n",
    "|---|---|\n",
    "| **LoRA** | 低秩分解适配器，只训练 ~0.1% 参数 |\n",
    "| **HuggingFace Trainer** | 封装训练 + 评估循环 |\n",
    "| **PEFT** | LoRA 注入/加载的标准化 API |\n",
    "| **FP16 Mixed Precision** | 半精度训练加速 |\n",
    "\n",
    "## BERT (本示例) vs Qwen 的关键差异\n",
    "\n",
    "| 维度 | BERT (110M) | Qwen2.5-7B |\n",
    "|---|---|---|\n",
    "| 模型架构 | Encoder-only | Decoder-only |\n",
    "| 任务类型 | SEQ_CLS (分类) | CAUSAL_LM (生成) |\n",
    "| 量化 | 不需要 | 需要 4-bit QLoRA |\n",
    "| LoRA target | `query`, `value` | `q_proj`, `v_proj`, `k_proj`, `o_proj` |\n",
    "| 推理方式 | `model(**inputs)` → logits | `model.generate()` → tokens |\n",
    "| 参数总量 | 110M | 7B |\n",
    "| 单卡训练显存 | ~2GB | ~8-10GB (4-bit) |\n",
    "| LoRA 适配器大小 | ~2MB | ~30MB |\n",
    "\n",
    "## 训练→推理完整流程\n",
    "\n",
    "1. **加载** BERT-base，冻结全部参数\n",
    "2. **注入** LoRA 适配器到 Attention 的 `query`/`value` 层\n",
    "3. **训练**：在 Rotten Tomatoes 数据上优化分类头，每个 epoch 后评估 accuracy\n",
    "4. **保存**：~2MB 适配器文件 + tokenizer\n",
    "5. **推理**：BERT-base + LoRA 适配器 → 预测新评论的情感"
   ]
  }
 ],
 "metadata": {
  "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.14.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}