{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 命名实体识别"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.5\n",
      "              precision    recall  f1-score   support\n",
      "\n",
      "        MISC       0.00      0.00      0.00         1\n",
      "         PER       1.00      1.00      1.00         1\n",
      "\n",
      "   micro avg       0.50      0.50      0.50         2\n",
      "   macro avg       0.50      0.50      0.50         2\n",
      "weighted avg       0.50      0.50      0.50         2\n",
      "\n"
     ]
    }
   ],
   "source": [
    "from seqeval.metrics import accuracy_score\n",
    "from seqeval.metrics import classification_report\n",
    "from seqeval.metrics import f1_score\n",
    "y_true = [['O', 'O', 'O', 'B-MISC', 'I-MISC', 'I-MISC', 'O'], ['B-PER', 'I-PER', 'O']]\n",
    "y_pred = [['O', 'O', 'B-MISC', 'I-MISC', 'I-MISC', 'I-MISC', 'O'], ['B-PER', 'I-PER', 'O']]\n",
    "print(f1_score(y_true, y_pred))\n",
    "print(classification_report(y_true, y_pred))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "d:\\gitcode\\.venv\\Lib\\site-packages\\tqdm\\auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "#1导入相关包\n",
    "import evaluate\n",
    "from datasets import load_dataset\n",
    "from transformers import AutoTokenizer, AutoModelForTokenClassification, TrainingArguments, Trainer, DataCollatorForTokenClassification"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Generating train split: 100%|██████████| 20865/20865 [00:00<00:00, 205845.42 examples/s]\n",
      "Generating validation split: 100%|██████████| 2319/2319 [00:00<00:00, 278085.34 examples/s]\n",
      "Generating test split: 100%|██████████| 4637/4637 [00:00<00:00, 343925.51 examples/s]\n"
     ]
    }
   ],
   "source": [
    "#2加载数据集\n",
    "# 直接使用load_dataset进行加载\n",
    "ner_datasets = load_dataset(\"D:/gitcode/data/peoples-daily-ner\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'id': '0',\n",
       " 'tokens': ['海',\n",
       "  '钓',\n",
       "  '比',\n",
       "  '赛',\n",
       "  '地',\n",
       "  '点',\n",
       "  '在',\n",
       "  '厦',\n",
       "  '门',\n",
       "  '与',\n",
       "  '金',\n",
       "  '门',\n",
       "  '之',\n",
       "  '间',\n",
       "  '的',\n",
       "  '海',\n",
       "  '域',\n",
       "  '。'],\n",
       " 'ner_tags': [0, 0, 0, 0, 0, 0, 0, 5, 6, 0, 5, 6, 0, 0, 0, 0, 0, 0]}"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ner_datasets[\"train\"][0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'id': Value(dtype='string', id=None),\n",
       " 'tokens': Sequence(feature=Value(dtype='string', id=None), length=-1, id=None),\n",
       " 'ner_tags': Sequence(feature=ClassLabel(names=['O', 'B-PER', 'I-PER', 'B-ORG', 'I-ORG', 'B-LOC', 'I-LOC'], id=None), length=-1, id=None)}"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ner_datasets[\"train\"].features"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['O', 'B-PER', 'I-PER', 'B-ORG', 'I-ORG', 'B-LOC', 'I-LOC']"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "label_list = ner_datasets[\"train\"].features[\"ner_tags\"].feature.names\n",
    "label_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "#3数据预处理\n",
    "tokenizer = AutoTokenizer.from_pretrained(\"D:/gitcode/model/chinese-macbert-base\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'input_ids': [[101, 3862, 102], [101, 7157, 102], [101, 3683, 102], [101, 6612, 102], [101, 1765, 102], [101, 4157, 102], [101, 1762, 102], [101, 1336, 102], [101, 7305, 102], [101, 680, 102], [101, 7032, 102], [101, 7305, 102], [101, 722, 102], [101, 7313, 102], [101, 4638, 102], [101, 3862, 102], [101, 1818, 102], [101, 511, 102]], 'token_type_ids': [[0, 0, 0], [0, 0, 0], [0, 0, 0], [0, 0, 0], [0, 0, 0], [0, 0, 0], [0, 0, 0], [0, 0, 0], [0, 0, 0], [0, 0, 0], [0, 0, 0], [0, 0, 0], [0, 0, 0], [0, 0, 0], [0, 0, 0], [0, 0, 0], [0, 0, 0], [0, 0, 0]], 'attention_mask': [[1, 1, 1], [1, 1, 1], [1, 1, 1], [1, 1, 1], [1, 1, 1], [1, 1, 1], [1, 1, 1], [1, 1, 1], [1, 1, 1], [1, 1, 1], [1, 1, 1], [1, 1, 1], [1, 1, 1], [1, 1, 1], [1, 1, 1], [1, 1, 1], [1, 1, 1], [1, 1, 1]]}"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tokenizer(ner_datasets[\"train\"][0][\"tokens\"]) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'input_ids': [101, 3862, 7157, 3683, 6612, 1765, 4157, 1762, 1336, 7305, 680, 7032, 7305, 722, 7313, 4638, 3862, 1818, 511, 102], 'token_type_ids': [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], 'attention_mask': [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]}"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# 对于已经做好tokenize的数据，要指定is_split_into_words参数为True\n",
    "tokenizer(ner_datasets[\"train\"][0][\"tokens\"], is_split_into_words=True)   "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'input_ids': [101, 10673, 12865, 12921, 8181, 8681, 102], 'token_type_ids': [0, 0, 0, 0, 0, 0, 0], 'attention_mask': [1, 1, 1, 1, 1, 1, 1]}"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#有可能一个英文的词，拆成几个token\n",
    "#所以要对标签做一下额外的处理\n",
    "res = tokenizer(\"interesting word\")\n",
    "res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[None, 0, 0, 0, 0, 1, None]"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#若一个单词被拆成了多个token\n",
    "#如果这些token来自于同一个词,那么这些对应id是相同的\n",
    "res.word_ids()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 借助word_ids 实现标签映射\n",
    "def process_function(examples):\n",
    "    tokenized_exmaples = tokenizer(examples[\"tokens\"],\n",
    "                                   max_length=128, \n",
    "                                   truncation=True, \n",
    "                                   is_split_into_words=True)\n",
    "    #return tokenized_exmaples\n",
    "    labels = []\n",
    "    for i, label in enumerate(examples[\"ner_tags\"]):\n",
    "        word_ids = tokenized_exmaples.word_ids(batch_index=i)\n",
    "        label_ids = []\n",
    "        for word_id in word_ids:\n",
    "            if word_id is None:\n",
    "                #CLS SEP 不参与计算\n",
    "                #在cross loss 中-100不参与运算\n",
    "                label_ids.append(-100)\n",
    "            else:\n",
    "                label_ids.append(label[word_id])\n",
    "        labels.append(label_ids)\n",
    "    tokenized_exmaples[\"labels\"] = labels\n",
    "    return tokenized_exmaples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Map: 100%|██████████| 20865/20865 [00:02<00:00, 9362.95 examples/s] \n",
      "Map: 100%|██████████| 2319/2319 [00:00<00:00, 8915.19 examples/s] \n",
      "Map: 100%|██████████| 4637/4637 [00:00<00:00, 9948.85 examples/s] \n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "DatasetDict({\n",
       "    train: Dataset({\n",
       "        features: ['id', 'tokens', 'ner_tags', 'input_ids', 'token_type_ids', 'attention_mask', 'labels'],\n",
       "        num_rows: 20865\n",
       "    })\n",
       "    validation: Dataset({\n",
       "        features: ['id', 'tokens', 'ner_tags', 'input_ids', 'token_type_ids', 'attention_mask', 'labels'],\n",
       "        num_rows: 2319\n",
       "    })\n",
       "    test: Dataset({\n",
       "        features: ['id', 'tokens', 'ner_tags', 'input_ids', 'token_type_ids', 'attention_mask', 'labels'],\n",
       "        num_rows: 4637\n",
       "    })\n",
       "})"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tokenized_datasets = ner_datasets.map(process_function, batched=True)\n",
    "tokenized_datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'id': '0', 'tokens': ['海', '钓', '比', '赛', '地', '点', '在', '厦', '门', '与', '金', '门', '之', '间', '的', '海', '域', '。'], 'ner_tags': [0, 0, 0, 0, 0, 0, 0, 5, 6, 0, 5, 6, 0, 0, 0, 0, 0, 0], 'input_ids': [101, 3862, 7157, 3683, 6612, 1765, 4157, 1762, 1336, 7305, 680, 7032, 7305, 722, 7313, 4638, 3862, 1818, 511, 102], 'token_type_ids': [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], 'attention_mask': [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], 'labels': [-100, 0, 0, 0, 0, 0, 0, 0, 5, 6, 0, 5, 6, 0, 0, 0, 0, 0, 0, -100]}\n"
     ]
    }
   ],
   "source": [
    "print(tokenized_datasets[\"train\"][0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of BertForTokenClassification were not initialized from the model checkpoint at /teacher_train/model/chinese-macbert-base and are newly initialized: ['classifier.bias', 'classifier.weight']\n",
      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
     ]
    }
   ],
   "source": [
    "#4创建模型\n",
    "# 对于所有的非二分类任务，切记要指定num_labels，否则就会device错误\n",
    "model = AutoModelForTokenClassification.from_pretrained(\"/teacher_train/model/chinese-macbert-base\",\n",
    "                                                        num_labels=len(label_list))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "7"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.config.num_labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#5创建评估函数\n",
    "# 这里方便大家加载，替换成了本地的加载方式，无需额外下载\n",
    "seqeval = evaluate.load(\"seqeval\")\n",
    "seqeval"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    " # predictions = [['O', 'O', 'B-MISC', 'I-MISC', 'I-MISC', 'I-MISC', 'O'], ['B-PER', 'I-PER', 'O']]\n",
    " # references = [['O', 'O', 'O', 'B-MISC', 'I-MISC', 'I-MISC', 'O'], ['B-PER', 'I-PER', 'O']]\n",
    "\n",
    "def eval_metric(pred):\n",
    "    predictions, labels = pred\n",
    "    predictions = np.argmax(predictions, axis=-1)\n",
    "\n",
    "    # 将id转换为原始的字符串类型的标签\n",
    "    true_predictions = [\n",
    "        [label_list[p] for p, l in zip(prediction, label) if l != -100]\n",
    "        for prediction, label in zip(predictions, labels) \n",
    "    ]\n",
    "\n",
    "    true_labels = [\n",
    "        [label_list[l] for p, l in zip(prediction, label) if l != -100]\n",
    "        for prediction, label in zip(predictions, labels) \n",
    "    ]\n",
    "\n",
    "    result = seqeval.compute(predictions=true_predictions, references=true_labels, scheme=\"IOB2\")\n",
    "\n",
    "    return {\n",
    "        \"f1\": result[\"overall_f1\"]\n",
    "    }\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "args = TrainingArguments(\n",
    "    output_dir=\"models_for_ner\",\n",
    "    logging_dir=\"./logs\",\n",
    "    per_device_train_batch_size=64,\n",
    "    per_device_eval_batch_size=128,\n",
    "    eval_strategy=\"epoch\",\n",
    "    save_strategy=\"epoch\",\n",
    "    metric_for_best_model=\"f1\",\n",
    "    load_best_model_at_end=True,\n",
    "    logging_steps=50,\n",
    "    num_train_epochs=1\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_3129756/746877450.py:9: FutureWarning: `tokenizer` is deprecated and will be removed in version 5.0.0 for `Trainer.__init__`. Use `processing_class` instead.\n",
      "  trainer = Trainer(\n"
     ]
    }
   ],
   "source": [
    "#7创建训练器\n",
    "\"\"\"\n",
    "DataCollatorForTokenClassification类适用于序列标注任务（token classification），\n",
    "比如命名实体识别（NER），每个token都会对应一个预测label，\n",
    "当一个序列的token比label多时，对额外的label位置添加-100处理，\n",
    "使得在计算交叉熵损失时，-100位置的label损失为0不用考虑。\n",
    "使得所有序列保持同一长度，同时对token多于label的场景使用-100进行填充。\n",
    "\"\"\"\n",
    "trainer = Trainer(\n",
    "    model=model,\n",
    "    args=args,\n",
    "    tokenizer=tokenizer,\n",
    "    train_dataset=tokenized_datasets[\"train\"],\n",
    "    eval_dataset=tokenized_datasets[\"validation\"],\n",
    "    compute_metrics=eval_metric,\n",
    "    data_collator=DataCollatorForTokenClassification(tokenizer=tokenizer)\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#8模型训练\n",
    "trainer.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#9评估\n",
    "trainer.evaluate(eval_dataset=tokenized_datasets[\"test\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#10模型预测\n",
    "from transformers import pipeline\n",
    "# 使用pipeline进行推理，要指定id2label\n",
    "model.config.id2label = {idx: label for idx, label in enumerate(label_list)}\n",
    "model.config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 如果模型是基于GPU训练的，那么推理时要指定device\n",
    "# 对于NER任务，可以指定aggregation_strategy为simple，得到具体的实体的结果，而不是token的结果\n",
    "ner_pipe = pipeline(\"token-classification\", model=model, tokenizer=tokenizer, device=0, aggregation_strategy=\"simple\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "gitcode",
   "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.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}