CS_BERT.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of the model checkpoint at bert-base-chinese were not used when initializing BertModel: ['cls.predictions.transform.LayerNorm.bias', 'cls.predictions.bias', 'cls.predictions.transform.dense.bias', 'cls.predictions.transform.LayerNorm.weight', 'cls.predictions.transform.dense.weight', 'cls.seq_relationship.bias', 'cls.seq_relationship.weight']\n",
      "- This IS expected if you are initializing BertModel from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
      "- This IS NOT expected if you are initializing BertModel from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n"
     ]
    }
   ],
   "source": [
    "from transformers import BertTokenizer, BertModel\n",
    "\n",
    "tokenizer = BertTokenizer.from_pretrained('bert-base-chinese')\n",
    "model = BertModel.from_pretrained('bert-base-chinese')\n",
    "\n",
    "tokenizer.save_pretrained('bert-base-chinese')\n",
    "model.save_pretrained('bert-base-chinese')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 读取数据\n",
    "import json\n",
    "\n",
    "with open('./data/z_medical_stroke_all.json', 'r', encoding='utf-8') as f:\n",
    "    data = json.load(f)\n",
    "print(data[0])\n",
    "\n",
    "# 将数据的instruction替换\n",
    "for i in range(len(data)):\n",
    "    data[i]['instruction'] = \"现在要做一个根据脑卒中病人的主诉、现病史、既往史和分级表进行严重程度分级的任务。案例会给出主诉、现病史、既往史。\\n\\t分级如下:\\n\\t正常\\n\\t轻度\\n\\t中度\\n\\t重度\\n\\t重度以上\\n请根据病人的主诉、现病史、既往史进行分级。\\n\"\n",
    "\n",
    "# 现在的数据的output字段是例如 2分 5分, 对其进行转换 0-1:正常 2-4:轻度 5-15:中度 16-20:重度 21-42:重度以上\n",
    "one = 0\n",
    "two = 0\n",
    "three = 0\n",
    "four = 0\n",
    "five = 0\n",
    "for i in range(len(data)):\n",
    "    if data[i]['output'] == '0分' or data[i]['output'] == '1分':\n",
    "        data[i]['output'] = \"正常\"\n",
    "        one += 1\n",
    "    elif data[i]['output'] == '2分' or data[i]['output'] == '3分' or data[i]['output'] == '4分':\n",
    "        data[i]['output'] = \"轻度\"\n",
    "        two += 1\n",
    "    elif data[i]['output'] == '5分' or data[i]['output'] == '6分' or data[i]['output'] == '7分' or data[i]['output'] == '8分' or data[i]['output'] == '9分' or data[i]['output'] == '10分' or data[i]['output'] == '11分' or data[i]['output'] == '12分' or data[i]['output'] == '13分' or data[i]['output'] == '14分' or data[i]['output'] == '15分':\n",
    "        data[i]['output'] = \"中度\"\n",
    "        three += 1\n",
    "    elif data[i]['output'] == '16分' or data[i]['output'] == '17分' or data[i]['output'] == '18分' or data[i]['output'] == '19分' or data[i]['output'] == '20分':\n",
    "        data[i]['output'] = \"重度\"\n",
    "        four += 1\n",
    "    elif data[i]['output'] == '21分' or data[i]['output'] == '22分' or data[i]['output'] == '23分' or data[i]['output'] == '24分' or data[i]['output'] == '25分' or data[i]['output'] == '26分' or data[i]['output'] == '27分' or data[i]['output'] == '28分' or data[i]['output'] == '29分' or data[i]['output'] == '30分' or data[i]['output'] == '31分' or data[i]['output'] == '32分' or data[i]['output'] == '33分' or data[i]['output'] == '34分' or data[i]['output'] == '35分' or data[i]['output'] == '36分' or data[i]['output'] == '37分' or data[i]['output'] == '38分' or data[i]['output'] == '39分' or data[i]['output'] == '40分' or data[i]['output'] == '41分' or data[i]['output'] == '42分':\n",
    "        data[i]['output'] = \"重度以上\"\n",
    "        five += 1\n",
    "print(f\"正常:{one} 轻度:{two} 中度:{three} 重度:{four} 重度以上:{five}\")\n",
    "\n",
    "\n",
    "\n",
    "# 将替换后的数据保存为jsonl格式 instuction + input 作为 input字段 output作为output字段\n",
    "input_len = []\n",
    "with open('./data/train.jsonl', 'w', encoding='utf-8') as f:\n",
    "    for i in range(len(data)):\n",
    "    # for i in range(64):\n",
    "        f.write(json.dumps({'input': data[i]['instruction'] + data[i]['input'], 'output': data[i]['output']}, ensure_ascii=False) + '\\n')\n",
    "        lens = len(data[i]['instruction'] + data[i]['input'])\n",
    "        input_len.append(lens)\n",
    "max_len = max(input_len)\n",
    "print(f\"max_len: {max_len}\")\n",
    "\n",
    "# 正常;轻度;中度各8个作为eval\n",
    "one = 0; two = 0; three = 0; four = 0; five = 0\n",
    "eval_data = []\n",
    "with open('./data/eval.jsonl', 'w', encoding='utf-8') as f:\n",
    "    for i in range(len(data)):\n",
    "        if data[i]['output'] == '正常' and one < 8:\n",
    "            f.write(json.dumps({'input': data[i]['instruction'] + data[i]['input'], 'output': data[i]['output']}, ensure_ascii=False) + '\\n')\n",
    "            one += 1\n",
    "        elif data[i]['output'] == '轻度' and two < 8:\n",
    "            f.write(json.dumps({'input': data[i]['instruction'] + data[i]['input'], 'output': data[i]['output']}, ensure_ascii=False) + '\\n')\n",
    "            two += 1\n",
    "        elif data[i]['output'] == '中度' and three < 8:\n",
    "            f.write(json.dumps({'input': data[i]['instruction'] + data[i]['input'], 'output': data[i]['output']}, ensure_ascii=False) + '\\n')\n",
    "            three += 1\n",
    "        else:\n",
    "            eval_data.append(data[i])\n",
    "print(f\"eval: 正常:{one} 轻度:{two} 中度:{three} 重度:{four} 重度以上:{five}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import BertTokenizer, BertConfig, BertForSequenceClassification\n",
    "import torch\n",
    "import os"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "config = BertConfig.from_pretrained('./bert-base-chinese')\n",
    "config.num_labels = 5\n",
    "config.problem_type = 'multi_label_classification'\n",
    "tokenizer = BertTokenizer.from_pretrained('./bert-base-chinese')\n",
    "\n",
    "model = BertForSequenceClassification.from_pretrained('./bert-base-chinese', config=config)\n",
    "\n",
    "if torch.cuda.is_available() and os.name == 'posix':\n",
    "    device = torch.compile(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(tokenizer.decode(tokenizer.encode('你好')))\n",
    "print(model, model.num_parameters())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.utils.data import Dataset, DataLoader\n",
    "import torch\n",
    "import json\n",
    "\n",
    "class MyDataset(Dataset):\n",
    "    def __init__(self, data_path, tokenizer, max_len=512, problem_type='multi_label_classification'):\n",
    "        self.data = []\n",
    "        self.tokenizer = tokenizer\n",
    "        self.max_len = max_len\n",
    "        with open(data_path, 'r', encoding='utf-8') as f:\n",
    "            for line in f:\n",
    "                line = json.loads(line)\n",
    "                self.data.append((line['input'], line['output']))\n",
    "        self.problem_type = problem_type\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.data)\n",
    "    \n",
    "    def __getitem__(self, idx):\n",
    "        input_text, label = self.data[idx]\n",
    "        input_ids = self.tokenizer.encode(input_text, max_length=self.max_len, padding='max_length', truncation=True, pad_to_max_length=True)\n",
    "        \n",
    "        if label == '正常':\n",
    "            label = 0\n",
    "        elif label == '轻度':\n",
    "            label = 1\n",
    "        elif label == '中度':\n",
    "            label = 2\n",
    "        elif label == '重度':\n",
    "            label = 3\n",
    "        elif label == '重度以上':\n",
    "            label = 4\n",
    "        else:\n",
    "            raise ValueError(\"Invalid label\")\n",
    "        \n",
    "        if self.problem_type == 'single_label_classification':\n",
    "            return {\n",
    "                'input_ids': torch.tensor(input_ids, dtype=torch.long),\n",
    "                'label': torch.tensor(label, dtype=torch.long)\n",
    "            }\n",
    "        elif self.problem_type == 'multi_label_classification':\n",
    "            return {\n",
    "                'input_ids': torch.tensor(input_ids, dtype=torch.long),\n",
    "                'label': torch.nn.functional.one_hot(torch.tensor(label), num_classes=5).float()\n",
    "            }\n",
    "        else:\n",
    "            raise ValueError(\"Invalid problem type\")\n",
    "        \n",
    "# 构建训练集和验证集\n",
    "train_dataset = MyDataset('./data/train.jsonl', tokenizer, config.max_position_embeddings, problem_type=config.problem_type)\n",
    "eval_dataset = MyDataset('./data/eval.jsonl', tokenizer, config.max_position_embeddings, problem_type=config.problem_type)\n",
    "\n",
    "# 构建DataLoader\n",
    "train_loader = DataLoader(train_dataset, batch_size=8, shuffle=True)\n",
    "eval_loader = DataLoader(eval_dataset, batch_size=8, shuffle=False)\n",
    "\n",
    "for i, data in enumerate(train_loader):\n",
    "    print(data['input_ids'].shape, data['label'].shape)\n",
    "    break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 测试模型在eval上的表现\n",
    "acc = 0\n",
    "for i, data in enumerate(eval_loader):\n",
    "    model.eval()\n",
    "    input_ids = data['input_ids']\n",
    "    label = data['label']\n",
    "    output = model(input_ids)\n",
    "    output = torch.argmax(output.logits, dim=1)\n",
    "    label = label.argmax(dim=-1)\n",
    "    print(f\"\\noutput:\\t{output}\\nlabel:\\t{label}\\n\")\n",
    "    acc += torch.sum(output == label)\n",
    "print(f\"acc: {acc / len(eval_dataset)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 学习率\n",
    "import matplotlib.pyplot as plt\n",
    "def lrate(\n",
    "    step: int,\n",
    "    warmup_steps: int = len(train_loader),\n",
    "    d_model: int = config.hidden_size,\n",
    ") -> float:\n",
    "    if step == 0:\n",
    "        return 0.0\n",
    "    return (d_model ** -0.5) * min(step ** -0.5, step * warmup_steps ** -1.5)\n",
    "plt.plot([lrate(step) for step in range(1, len(train_loader) * 10)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 优化器, 学习率调度器, 混合精度训练\n",
    "from torch.optim import AdamW\n",
    "from torch.optim.lr_scheduler import LambdaLR\n",
    "from torch.cuda.amp import GradScaler, autocast\n",
    "\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "optimizer = AdamW(model.parameters(), lr=1, betas=(0.9, 0.98), eps=1e-9)\n",
    "scheduler = LambdaLR(optimizer, lr_lambda=lrate)\n",
    "scaler = GradScaler(enabled=True) if torch.cuda.is_available() else None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 训练函数\n",
    "import torch.utils.tensorboard as tensorboard\n",
    "import os\n",
    "import torch.distributed as dist\n",
    "from torch.nn.parallel import DistributedDataParallel as DDP\n",
    "# 记录器\n",
    "writer = tensorboard.SummaryWriter('./logs/medical_stroke')\n",
    "\n",
    "torch.backends.cudnn.benchmark = True # 加速训练\n",
    "torch.manual_seed(233)\n",
    "\n",
    "def trainer(\n",
    "    model: torch.nn.Module,\n",
    "    train_loader: DataLoader,\n",
    "    eval_loader: DataLoader,\n",
    "    tokenizer: BertTokenizer,\n",
    "    device: torch.device,\n",
    "    optimizer: torch.optim.Optimizer,\n",
    "    scheduler: torch.optim.lr_scheduler._LRScheduler,\n",
    "    scaler: torch.cuda.amp.GradScaler,\n",
    "    writer: tensorboard.SummaryWriter,\n",
    "    epochs: int = 10,\n",
    "):\n",
    "    # 分布式训练命令 python -m torch.distributed.launch --nproc_per_node=2 train.py\n",
    "    if torch.cuda.device_count() > 1 and os.name == 'posix':\n",
    "        print(\"Using\", torch.cuda.device_count(), \"GPUs\")\n",
    "        rank = dist.get_rank()\n",
    "        torch.cuda.set_device(rank % torch.cuda.device_count())\n",
    "        dist.init_process_group(backend='nccl', init_method='env://')\n",
    "        model = DDP(model, device_ids=[rank % torch.cuda.device_count()])\n",
    "        model.to(device)\n",
    "    else:\n",
    "        print(\"Using\", torch.cuda.device_count(), \"GPUs\")\n",
    "        model = model.to(device)\n",
    "    global_step = 0\n",
    "    train_loss = []\n",
    "    acc = 0 \n",
    "    ppl = 0\n",
    "\n",
    "    for epoch in range(epochs):\n",
    "        model.train()\n",
    "        for step, batch in enumerate(train_loader):\n",
    "            input_ids = batch['input_ids'].to(device)\n",
    "            label = batch['label'].to(device)\n",
    "            \n",
    "            # 梯度清零\n",
    "            optimizer.zero_grad(set_to_none=True)\n",
    "            # 混合精度训练\n",
    "            if scaler is not None:\n",
    "                with autocast():\n",
    "                    # 前向传播\n",
    "                    output = model(input_ids, labels=label)\n",
    "                    loss = output.loss\n",
    "                # 反向传播\n",
    "                scaler.scale(loss).backward()\n",
    "                # 梯度裁剪\n",
    "                scaler.unscale_(optimizer)\n",
    "                torch.nn.utils.clip_grad_norm_(model.parameters(), 1)\n",
    "                # 更新权重\n",
    "                scaler.step(optimizer)\n",
    "                scaler.update()\n",
    "            else:\n",
    "                output = model(input_ids, labels=label)\n",
    "                print(output.logits)\n",
    "                loss = output.loss\n",
    "                loss.backward()\n",
    "                optimizer.step()\n",
    "            # 更新学习率\n",
    "            scheduler.step()\n",
    "            # 记录训练损失\n",
    "            train_loss.append(loss.item())\n",
    "            global_step += 1\n",
    "\n",
    "            # 打印训练信息\n",
    "            if (step+1) % (len(train_loader)//10) == 0:\n",
    "            # if (step+1) % 1 == 0:\n",
    "                train_loss = sum(train_loss) / len(train_loss)\n",
    "                print(f\"Epoch [{epoch+1}/{epochs}] Step [{step+1}/{len(train_loader)}] lr: {scheduler.get_last_lr()[0]} loss: {train_loss}\")\n",
    "\n",
    "                # 记录训练信息\n",
    "                writer.add_scalar('train_lr', scheduler.get_last_lr()[0], global_step)\n",
    "                writer.add_scalar('train_loss', train_loss, global_step)\n",
    "                train_loss = []\n",
    "\n",
    "        # 验证\n",
    "        if epoch % 1 == 0:\n",
    "            model.eval()\n",
    "            with torch.no_grad():\n",
    "                correct = 0\n",
    "                total = 0\n",
    "                for step, batch in enumerate(eval_loader):\n",
    "                    input_ids = batch['input_ids'].to(device)\n",
    "                    label = batch['label'].to(device)\n",
    "                    output = model(input_ids, labels=label)\n",
    "                    if config.problem_type == 'single_label_classification':\n",
    "                        pred = torch.argmax(output.logits, dim=-1)\n",
    "                        correct += (pred == label).sum().item()\n",
    "                        total += len(label)\n",
    "                    elif config.problem_type == 'multi_label_classification':\n",
    "                        pred = torch.argmax(output.logits, dim=-1)\n",
    "                        correct += (pred == label.argmax(dim=-1)).sum().item()\n",
    "                        total += len(label)\n",
    "                    else:\n",
    "                        raise ValueError(\"Invalid problem type\")\n",
    "                acc = correct / total\n",
    "                print(f\"Epoch [{epoch+1}/{epochs}] accuracy: {acc}\")\n",
    "                writer.add_scalar('eval_accuracy', acc, global_step)\n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = trainer(model, train_loader, eval_loader, tokenizer, device, optimizer, scheduler, scaler, writer, epochs=100)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 保存模型\n",
    "model.save_pretrained('./medical_stroke')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "self",
   "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.9.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}