notebook.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "509d27c7",
   "metadata": {},
   "source": [
    "#### 本篇博客希望展示如何基于transformers提供的功能进行模型的开发，减少代码量，提高开发速度。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "3faa5227",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import warnings\n",
    "import torch.nn as nn\n",
    "\n",
    "from torch import Tensor\n",
    "from typing import List, Dict\n",
    "from dataclasses import dataclass, field\n",
    "from torch.nn.utils.rnn import pad_sequence\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from transformers.file_utils import logger, logging\n",
    "from transformers.trainer_utils import EvalPrediction\n",
    "from transformers.modeling_outputs import TokenClassifierOutput\n",
    "from sklearn.metrics import f1_score, precision_score, recall_score\n",
    "from transformers import TrainingArguments, Trainer, BertTokenizer, BertModel, BertPreTrainedModel\n",
    "\n",
    "warnings.filterwarnings(\"ignore\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "78d01ea5",
   "metadata": {},
   "source": [
    "### 一、定义参数"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "d2e82946",
   "metadata": {},
   "outputs": [],
   "source": [
    "@dataclass\n",
    "class ModelArguments:\n",
    "    use_lstm: bool = field(default=True, metadata={\"help\": \"是否使用LSTM\"})\n",
    "    lstm_hidden_size: int = field(default=500, metadata={\"help\": \"LSTM隐藏层输出的维度\"})\n",
    "    lstm_layers: int = field(default=1, metadata={\"help\": \"堆叠LSTM的层数\"})\n",
    "    lstm_dropout: float = field(default=0.5, metadata={\"help\": \"LSTM的dropout\"})\n",
    "    hidden_dropout: float = field(default=0.5, metadata={\"help\": \"预训练模型输出向量表示的dropout\"})\n",
    "    ner_num_labels: int = field(default=12, metadata={\"help\": \"需要预测的标签数量\"})\n",
    "\n",
    "\n",
    "@dataclass\n",
    "class OurTrainingArguments:\n",
    "    checkpoint_dir: str = field(default=\"./models/checkpoints\", metadata={\"help\": \"训练过程中的checkpoints的保存路径\"})\n",
    "    best_dir: str = field(default=\"./models/best\", metadata={\"help\": \"最优模型的保存路径\"})\n",
    "    do_eval: bool = field(default=True, metadata={\"help\": \"是否在训练时进行评估\"})\n",
    "    epoch: int = field(default=5, metadata={\"help\": \"训练的epoch\"})\n",
    "    train_batch_size: int = field(default=8, metadata={\"help\": \"训练时的batch size\"})\n",
    "    eval_batch_size: int = field(default=8, metadata={\"help\": \"评估时的batch size\"})\n",
    "\n",
    "\n",
    "@dataclass\n",
    "class DataArguments:\n",
    "    train_file: str = field(default=\"./data/train.txt\", metadata={\"help\": \"训练数据的路径\"})\n",
    "    dev_file: str = field(default=\"./data/dev.txt\", metadata={\"help\": \"测试数据的路径\"})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "32460633",
   "metadata": {},
   "source": [
    "### 二、读取数据"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ea4db3c1",
   "metadata": {},
   "source": [
    "这里定义了一个用于保存数据的数据结构，这样的方法能够提高代码的可阅读性。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "b3d61bfc",
   "metadata": {},
   "outputs": [],
   "source": [
    "@dataclass\n",
    "class Example:\n",
    "    text: List[str] # ner的文本\n",
    "    label: List[str] = None # ner的标签\n",
    "\n",
    "    def __post_init__(self):\n",
    "        if self.label:\n",
    "            assert len(self.text) == len(self.label)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1e160541",
   "metadata": {},
   "source": [
    "定义将文件中的ner数据保存为Example列表的函数"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "9d31e5d8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Example(text=['回', '眸', '’', '９', '７', '房', '地', '产', '景', '气', '水', '平', '缓', '缓', '回', '升'], label=['O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O'])\n"
     ]
    }
   ],
   "source": [
    "def read_data(path):\n",
    "    examples = []\n",
    "    with open(path, \"r\", encoding=\"utf-8\") as file:\n",
    "        text = []\n",
    "        label = []\n",
    "        for line in file:\n",
    "            line = line.strip()\n",
    "            # 一条文本结束\n",
    "            if len(line) == 0:\n",
    "                examples.append(Example(text, label))\n",
    "                text = []\n",
    "                label = []\n",
    "                continue\n",
    "            text.append(line.split()[0])\n",
    "            label.append(line.split()[1])\n",
    "    return examples\n",
    "\n",
    "train_data = read_data(\"./data/train.txt\")\n",
    "eval_data = read_data(\"./data/dev.txt\")\n",
    "print(train_data[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6832ad3e",
   "metadata": {},
   "source": [
    "加载标签数据并分配对于的id"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "afcad60f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'<pad>': 0, 'O': 1, 'B-ORG': 2, 'I-ORG': 3, 'B-LOC': 4, 'I-LOC': 5, 'B-TIME': 6, 'I-TIME': 7, 'B-PER': 8, 'I-PER': 9, '<start>': 10, '<eos>': 11}\n",
      "{0: '<pad>', 1: 'O', 2: 'B-ORG', 3: 'I-ORG', 4: 'B-LOC', 5: 'I-LOC', 6: 'B-TIME', 7: 'I-TIME', 8: 'B-PER', 9: 'I-PER', 10: '<start>', 11: '<eos>'}\n"
     ]
    }
   ],
   "source": [
    "def load_tag(path):\n",
    "    with open(path, \"r\", encoding=\"utf-8\") as file:\n",
    "        lines = file.readlines()\n",
    "        tag2id = {tag.strip(): idx for idx, tag in enumerate(lines)}\n",
    "        id2tag = {idx: tag for tag, idx in tag2id.items()}\n",
    "    return tag2id, id2tag\n",
    "\n",
    "tag2id, id2tag = load_tag(\"./data/tag.txt\")\n",
    "print(tag2id)\n",
    "print(id2tag)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ccc3eee1",
   "metadata": {},
   "source": [
    "读取tokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "c0a4d3ed",
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer = BertTokenizer.from_pretrained(\"bert-base-chinese\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "edb71215",
   "metadata": {},
   "source": [
    "### 三、构建Dataset和collate_fn"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e64aba82",
   "metadata": {},
   "source": [
    "构建Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "983edd18",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'input_ids': tensor([ 101, 1726, 4704,  100, 8037, 8035, 2791, 1765,  772, 3250, 3698, 3717,\n",
      "        2398, 5353, 5353, 1726, 1285,  102]), 'labels': tensor([10,  1,  1,  1,  1,  1,  1,  1,  1,  1,  1,  1,  1,  1,  1,  1,  1, 11])}\n"
     ]
    }
   ],
   "source": [
    "class NERDataset(Dataset):\n",
    "    def __init__(self, examples: List[Example], max_length=128):\n",
    "        self.max_length = 512 if max_length > 512 else max_length\n",
    "        \"\"\"\n",
    "        1. 将文本的长度控制在max_length - 2，减2的原因是为[CLS]和[SEP]空出位置； \n",
    "        2. 将文本转换为id序列；\n",
    "        3. 将id序列转换为Tensor；\n",
    "        \"\"\"\n",
    "        self.texts = [torch.LongTensor(tokenizer.encode(example.text[: self.max_length - 2])) for example in examples]\n",
    "        self.labels = []\n",
    "        for example in examples:\n",
    "            label = example.label\n",
    "            \"\"\"\n",
    "            1. 将字符的label转换为对于的id；\n",
    "            2. 控制label的最长长度；\n",
    "            3. 添加开始位置和结束位置对应的标签，这里<start>对应输入中的[CLS],<eos>对于[SEP]；\n",
    "            4. 转换为Tensor；\n",
    "            \"\"\"\n",
    "            label = [tag2id[\"<start>\"]] + [tag2id[l] for l in label][: self.max_length - 2] + [tag2id[\"<eos>\"]]\n",
    "            self.labels.append(torch.LongTensor(label))\n",
    "        assert len(self.texts) == len(self.labels)\n",
    "        for text, label in zip(self.texts, self.labels):\n",
    "            assert len(text) == len(label)\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.texts)\n",
    "\n",
    "    def __getitem__(self, item):\n",
    "        return {\n",
    "            \"input_ids\": self.texts[item],\n",
    "            \"labels\": self.labels[item]\n",
    "        }\n",
    "\n",
    "train_dataset = NERDataset(train_data)\n",
    "eval_dataset = NERDataset(eval_data)\n",
    "print(train_dataset[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f067e39b",
   "metadata": {},
   "source": [
    "定义collate_fn，collate_fn的作用在Dataloader生成batch数据时会被调用。\n",
    "这里的作用是对每个batch进行padding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "3f7a62d0",
   "metadata": {},
   "outputs": [],
   "source": [
    "def collate_fn(features) -> Dict[str, Tensor]:\n",
    "    batch_input_ids = [feature[\"input_ids\"] for feature in features]\n",
    "    batch_labels = [feature[\"labels\"] for feature in features]\n",
    "    batch_attentiton_mask = [torch.ones_like(feature[\"input_ids\"]) for feature in features]\n",
    "    # padding\n",
    "    batch_input_ids = pad_sequence(batch_input_ids, batch_first=True, padding_value=tokenizer.pad_token_id)\n",
    "    batch_labels = pad_sequence(batch_labels, batch_first=True, padding_value=tag2id[\"<pad>\"])\n",
    "    batch_attentiton_mask = pad_sequence(batch_attentiton_mask, batch_first=True, padding_value=0)\n",
    "    assert batch_input_ids.shape == batch_labels.shape\n",
    "    return {\"input_ids\": batch_input_ids, \"labels\": batch_labels, \"attention_mask\": batch_attentiton_mask}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "319b7fe1",
   "metadata": {},
   "source": [
    "测试一下collate_fn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "d223b32c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "dict_keys(['input_ids', 'labels', 'attention_mask'])\n",
      "<class 'torch.Tensor'>\n",
      "torch.Size([2, 19])\n",
      "<class 'torch.Tensor'>\n",
      "torch.Size([2, 19])\n",
      "<class 'torch.Tensor'>\n",
      "torch.Size([2, 19])\n"
     ]
    }
   ],
   "source": [
    "dataloader = DataLoader(train_dataset, shuffle=True, batch_size=2, collate_fn=collate_fn)\n",
    "batch = next(iter(dataloader))\n",
    "print(batch.keys())\n",
    "print(type(batch[\"input_ids\"]))\n",
    "print(batch[\"input_ids\"].shape)\n",
    "print(type(batch[\"labels\"]))\n",
    "print(batch[\"labels\"].shape)\n",
    "print(type(batch[\"attention_mask\"]))\n",
    "print(batch[\"attention_mask\"].shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f4ff64a7",
   "metadata": {},
   "source": [
    "### 四、定义一个评估函数"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "2ea6b115",
   "metadata": {},
   "outputs": [],
   "source": [
    "def ner_metrics(eval_output: EvalPrediction) -> Dict[str, float]:\n",
    "    \"\"\"\n",
    "    该函数是回调函数，Trainer会在进行评估时调用该函数。\n",
    "    (如果使用Pycharm等IDE进行调试，可以使用断点的方法来调试该函数，该函数在进行评估时被调用)\n",
    "    \"\"\"\n",
    "    preds = eval_output.predictions\n",
    "    preds = np.argmax(preds, axis=-1).flatten()\n",
    "    labels = eval_output.label_ids.flatten()\n",
    "    # labels为0表示为<pad>，因此计算时需要去掉该部分\n",
    "    mask = labels != 0\n",
    "    preds = preds[mask]\n",
    "    labels = labels[mask]\n",
    "    metrics = dict()\n",
    "    metrics[\"f1\"] = f1_score(labels, preds, average=\"macro\")\n",
    "    metrics[\"precision\"] = precision_score(labels, preds, average=\"macro\")\n",
    "    metrics[\"recall\"] = recall_score(labels, preds, average=\"macro\")\n",
    "    # 必须以字典的形式返回，后面会用到字典的key\n",
    "    return metrics"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4cae549d",
   "metadata": {},
   "source": [
    "### 五、构建模型\n",
    "+ 自定义的模型需要继承BertPreTrainedModel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "9a9b53e7",
   "metadata": {},
   "outputs": [],
   "source": [
    "class BertForNER(BertPreTrainedModel):\n",
    "    def __init__(self, config, *model_args, **model_kargs):\n",
    "        super().__init__(config) # 初始化父类(必要的步骤)\n",
    "        if \"model_args\" in model_kargs:\n",
    "            model_args = model_kargs[\"model_args\"]\n",
    "            \"\"\"\n",
    "            必须将额外的参数更新至self.config中，这样在调用save_model保存模型时才会将这些参数保存；\n",
    "            这种在使用from_pretrained方法加载模型时才不会出错；\n",
    "            \"\"\"\n",
    "            self.config.__dict__.update(model_args.__dict__)\n",
    "        self.num_labels = self.config.ner_num_labels\n",
    "        self.bert = BertModel(config, add_pooling_layer=False)\n",
    "        self.dropout = nn.Dropout(self.config.hidden_dropout)\n",
    "        self.lstm = nn.LSTM(self.config.hidden_size, # 输入的维度\n",
    "                            self.config.lstm_hidden_size, # 输出维度\n",
    "                            num_layers=self.config.lstm_layers, # 堆叠lstm的层数\n",
    "                            dropout=self.config.lstm_dropout,\n",
    "                            bidirectional=True, # 是否双向\n",
    "                            batch_first=True)\n",
    "        if self.config.use_lstm:\n",
    "            self.classifier = nn.Linear(self.config.lstm_hidden_size * 2, self.num_labels)\n",
    "        else:\n",
    "            self.classifier = nn.Linear(self.config.hidden_size, self.num_labels)\n",
    "        self.init_weights()\n",
    "\n",
    "    def forward(\n",
    "            self,\n",
    "            input_ids=None,\n",
    "            attention_mask=None,\n",
    "            token_type_ids=None,\n",
    "            position_ids=None,\n",
    "            head_mask=None,\n",
    "            inputs_embeds=None,\n",
    "            labels=None,\n",
    "            pos=None,\n",
    "            output_attentions=None,\n",
    "            output_hidden_states=None,\n",
    "            return_dict=None,\n",
    "    ):\n",
    "        return_dict = return_dict if return_dict is not None else self.config.use_return_dict\n",
    "        outputs = self.bert(\n",
    "            input_ids,\n",
    "            attention_mask=attention_mask,\n",
    "            token_type_ids=token_type_ids,\n",
    "            position_ids=position_ids,\n",
    "            head_mask=head_mask,\n",
    "            inputs_embeds=inputs_embeds,\n",
    "            output_attentions=output_attentions,\n",
    "            output_hidden_states=output_hidden_states,\n",
    "            return_dict=return_dict,\n",
    "        )\n",
    "        sequence_output = self.dropout(outputs[0])\n",
    "        if self.config.use_lstm:\n",
    "            sequence_output, _ = self.lstm(sequence_output)\n",
    "        logits = self.classifier(sequence_output)\n",
    "\n",
    "        loss = None\n",
    "        if labels is not None:\n",
    "            loss_fct = nn.CrossEntropyLoss()\n",
    "            # 如果attention_mask不为空，则只计算attention_mask中为1部分的Loss\n",
    "            if attention_mask is not None:\n",
    "                active_loss = attention_mask.view(-1) == 1\n",
    "                active_logits = logits.view(-1, self.num_labels)\n",
    "                active_labels = torch.where(\n",
    "                    active_loss, labels.view(-1), torch.tensor(loss_fct.ignore_index).type_as(labels)\n",
    "                )\n",
    "                loss = loss_fct(active_logits, active_labels)\n",
    "            else:\n",
    "                loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))\n",
    "\n",
    "        if not return_dict:\n",
    "            output = (logits,) + outputs[2:]\n",
    "            return ((loss,) + output) if loss is not None else output\n",
    "\n",
    "        return TokenClassifierOutput(\n",
    "            loss=loss,\n",
    "            logits=logits, # 该部分在评估时，会作为EvalPrediction对象的predictions进行返回\n",
    "            hidden_states=outputs.hidden_states,\n",
    "            attentions=outputs.attentions,\n",
    "        )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "259732dd",
   "metadata": {},
   "source": [
    "测试一下模型是否符合预期"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "eadf84fb",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of the model checkpoint at E:/pretrained/bert-base-chinese were not used when initializing BertForNER: ['cls.predictions.transform.dense.bias', 'cls.seq_relationship.weight', 'cls.predictions.decoder.weight', 'cls.predictions.bias', 'bert.pooler.dense.bias', 'cls.predictions.transform.LayerNorm.weight', 'cls.predictions.transform.LayerNorm.bias', 'cls.predictions.transform.dense.weight', 'bert.pooler.dense.weight', 'cls.seq_relationship.bias']\n",
      "- This IS expected if you are initializing BertForNER 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 BertForNER from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n",
      "Some weights of BertForNER were not initialized from the model checkpoint at E:/pretrained/bert-base-chinese and are newly initialized: ['lstm.bias_hh_l0_reverse', 'classifier.weight', 'lstm.weight_hh_l0_reverse', 'lstm.bias_ih_l0_reverse', 'lstm.bias_hh_l0', 'classifier.bias', 'lstm.weight_ih_l0', 'lstm.weight_hh_l0', 'lstm.weight_ih_l0_reverse', 'lstm.bias_ih_l0']\n",
      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<class 'transformers.modeling_outputs.TokenClassifierOutput'>\n",
      "tensor(2.5061, grad_fn=<NllLossBackward>)\n",
      "torch.Size([2, 19, 12])\n"
     ]
    }
   ],
   "source": [
    "model_args = ModelArguments(use_lstm=True)\n",
    "model = BertForNER.from_pretrained(\"bert-base-chinese\", model_args=model_args)\n",
    "output = model(**batch)\n",
    "print(type(output))\n",
    "print(output.loss)\n",
    "print(output.logits.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c10baab8",
   "metadata": {},
   "source": [
    "### 六、模型训练"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "id": "436a3573",
   "metadata": {},
   "outputs": [],
   "source": [
    "def run(model_args: ModelArguments, data_args: DataArguments, args: OurTrainingArguments):\n",
    "    # 设定训练参数\n",
    "    training_args = TrainingArguments(output_dir=args.checkpoint_dir,  # 训练中的checkpoint保存的位置\n",
    "                                      num_train_epochs=args.epoch,\n",
    "                                      do_eval=args.do_eval,  # 是否进行评估\n",
    "                                      evaluation_strategy=\"epoch\",  # 每个epoch结束后进行评估\n",
    "                                      per_device_train_batch_size=args.train_batch_size,\n",
    "                                      per_device_eval_batch_size=args.eval_batch_size,\n",
    "                                      load_best_model_at_end=True,  # 训练完成后加载最优模型\n",
    "                                      metric_for_best_model=\"f1\"  # 评估最优模型的指标，该指标是ner_metrics返回评估指标中的key\n",
    "                                      )\n",
    "    # 构建dataset\n",
    "    train_dataset = NERDataset(read_data(data_args.train_file))\n",
    "    eval_dataset = NERDataset(read_data(data_args.dev_file))\n",
    "    # 加载预训练模型\n",
    "    model = BertForNER.from_pretrained(\"bert-base-chinese\", model_args=model_args)\n",
    "    # 初始化Trainer\n",
    "    trainer = Trainer(model=model,\n",
    "                      args=training_args,\n",
    "                      train_dataset=train_dataset,\n",
    "                      eval_dataset=eval_dataset,\n",
    "                      tokenizer=tokenizer,\n",
    "                      data_collator=collate_fn,\n",
    "                      compute_metrics=ner_metrics)\n",
    "    # 模型训练\n",
    "    trainer.train()\n",
    "    # 训练完成后，加载最优模型并进行评估\n",
    "    logger.info(trainer.evaluate(eval_dataset))\n",
    "    # 保存训练好的模型\n",
    "    trainer.save_model(args.best_dir)\n",
    "\n",
    "# 定义各类参数并训练模型\n",
    "model_args = ModelArguments(use_lstm=True)\n",
    "data_args = DataArguments()\n",
    "training_args = OurTrainingArguments(train_batch_size=16, eval_batch_size=32)\n",
    "run(model_args, data_args, training_args)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Pytorch1.8",
   "language": "python",
   "name": "pytorch1.8"
  },
  "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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}